[Linalg-Vectorizer] Add pattern for mixed precision vector.contract. (#1067)

This pattern tries to fold arith.ext* ops present in high level mixed
precision linalg contraction ops into mixed precision vector.contract.

Author: shahidact

lib/TPP/Transforms/LinalgVectorize.cpp

@@ -78,6 +78,7 @@ struct LinalgVectorize
     tensor::populateFoldTensorSubsetIntoVectorTransferPatterns(patterns);
     patterns.add<linalg::CopyVectorizationPattern>(ctx);
     vector::populateVectorStepLoweringPatterns(patterns);
+    vector::populateFoldArithExtensionPatterns(patterns);
 
     if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
       return signalPassFailure();

lib/TPP/Transforms/Vectorization.cpp

@@ -11,6 +11,7 @@
 #include "TPP/Dialect/Xsmm/XsmmUtils.h"
 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -37,8 +38,6 @@ struct LinalgGenericToVector : OpRewritePattern<linalg::GenericOp> {
 
   LogicalResult matchAndRewrite(linalg::GenericOp linalgOp,
                                 PatternRewriter &rewriter) const override {
-    if (!linalgOp.hasPureBufferSemantics())
-      return failure();
     if (xsmm::utils::getDataType(rewriter, linalgOp.getOperand(0).getType()) ==
             xsmm::DataTypeAttr::get(rewriter.getContext(),
                                     xsmm::DataType::BF16) &&
@@ -107,6 +106,7 @@ struct VectorizationPass
 
   void populateCombinePatterns(RewritePatternSet &patterns) {
     patterns.add<LinalgToVector<linalg::BatchReduceMatmulOp>,
+                 LinalgToVector<linalg::ContractOp>,
                  LinalgToVector<linalg::TransposeOp>,
                  LinalgToVector<linalg::FillOp>>(patterns.getContext());
     patterns.add<LinalgGenericToVector>(patterns.getContext());
@@ -117,6 +117,7 @@ struct VectorizationPass
     populateCombinePatterns(patterns);
     vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
     vector::populateVectorReductionToContractPatterns(patterns);
+    vector::populateFoldArithExtensionPatterns(patterns);
     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
   }
 };

test/Passes/linalg-vectorize.mlir

@@ -54,7 +54,7 @@ func.func @vectorize_contract_mixed_precision_float(
 // CHECK-NOT: vector.transpose
 // CHECK: vector.transfer_read{{.*}}: tensor<128x256x2xbf16>, vector<128x256x2xbf16>
 // CHECK: vector.transfer_read{{.*}}: tensor<256x256xf32>, vector<256x256xf32>
-// CHECK-COUNT-2: arith.extf
+// CHECK-NOT: arith.extf
 // CHECK: vector.contract
 // CHECK: vector.transfer_write
 
@@ -92,7 +92,7 @@ module {
 // CHECK-NOT: vector.transpose
 // CHECK: vector.transfer_read{{.*}}: tensor<2x2x8x32x4xi8>, vector<2x2x8x32x4xi8>
 // CHECK: vector.transfer_read{{.*}}: tensor<1x2x32x32xi32>, vector<1x2x32x32xi32>
-// CHECK-COUNT-2: arith.extsi
+// CHECK-NOT: arith.extsi
 // CHECK: vector.contract
 // CHECK: vector.transfer_write
 

test/Passes/pass-vectorization.mlir

@@ -113,8 +113,70 @@ module {
 
 // CHECK:   func.func @entry(%[[ARG0:.*]]: tensor<2x4x8x1x2xbf16>) -> tensor<2x2x8x4xbf16> {
 // CHECK:       vector.transfer_write 
-// CHECK-NOT:       %[[vec1:.*]] = vector.transfer_read 
-// CHECK-NOT:       %[[vec2:.*]] = vector.transfer_read 
-// CHECK-NOT:       %[[vec3:.*]] = vector.transfer_read 
-// CHECK-NOT:       %[[vec4:.*]] = vector.contract 
-// CHECK-NOT:       vector.transfer_write %[[vec4]]
+// CHECK:       vector.transfer_read
+// CHECK:       vector.transfer_read
+// CHECK:       vector.contract
+// CHECK:       vector.transfer_write
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2, d4, d6, d3)>
+#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d2, d6, d5, d3)>
+#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d4, d5)>
+module {
+  func.func @vectorize_contract_mixed_precision_int(
+      %arg0: tensor<1x2x32x8x4xi8>, %arg1: tensor<2x2x8x32x4xi8>,
+      %arg2: tensor<1x2x32x32xi32>) -> tensor<1x2x32x32xi32> {
+    %0 = linalg.generic {
+      indexing_maps = [#map, #map1, #map2],
+      iterator_types = ["parallel", "parallel", "reduction", "reduction", "parallel", "parallel", "reduction"]}
+      ins(%arg0, %arg1 : tensor<1x2x32x8x4xi8>, tensor<2x2x8x32x4xi8>)
+      outs(%arg2 : tensor<1x2x32x32xi32>) {
+    ^bb0(%in: i8, %in_0: i8, %out: i32):
+      %0 = arith.extsi %in : i8 to i32
+      %1 = arith.extsi %in_0 : i8 to i32
+      %2 = arith.muli %0, %1 : i32
+      %3 = arith.addi %out, %2 : i32
+      linalg.yield %3 : i32
+    } -> tensor<1x2x32x32xi32>
+    return %0 : tensor<1x2x32x32xi32>
+  }
+}
+
+// CHECK-LABEL: @vectorize_contract_mixed_precision_int
+// CHECK: vector.transfer_read{{.*}}: tensor<1x2x32x8x4xi8>, vector<1x2x32x8x4xi8>
+// CHECK-NOT: vector.broadcast
+// CHECK-NOT: vector.transpose
+// CHECK: vector.transfer_read{{.*}}: tensor<2x2x8x32x4xi8>, vector<2x2x8x32x4xi8>
+// CHECK: vector.transfer_read{{.*}}: tensor<1x2x32x32xi32>, vector<1x2x32x32xi32>
+// CHECK-NOT: arith.extsi
+// CHECK: vector.contract
+// CHECK: vector.transfer_write
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+func.func @vectorize_contract_mixed_precision_float(
+    %arg0: tensor<256x128x2xbf16>, %arg1: tensor<128x256x2xbf16>,
+    %arg2: tensor<256x256xf32>) -> tensor<256x256xf32> {
+  %0 = linalg.contract
+    indexing_maps = [#map, #map1, #map2]
+    ins(%arg0, %arg1 : tensor<256x128x2xbf16>, tensor<128x256x2xbf16>)
+    outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
+  return %0 : tensor<256x256xf32>
+}
+
+// Ensure that mixed precision contraction vectorizes cleanly
+// without extra operations and/or dimensions.
+
+// CHECK-LABEL: @vectorize_contract_mixed_precision_float
+// CHECK: vector.transfer_read{{.*}}: tensor<256x128x2xbf16>, vector<256x128x2xbf16>
+// CHECK-NOT: vector.broadcast
+// CHECK-NOT: vector.transpose
+// CHECK: vector.transfer_read{{.*}}: tensor<128x256x2xbf16>, vector<128x256x2xbf16>
+// CHECK: vector.transfer_read{{.*}}: tensor<256x256xf32>, vector<256x256xf32>
+// CHECK-NOT: arith.extf
+// CHECK: vector.contract
+// CHECK: vector.transfer_write