[mlir][Linalg] Extend linalg vectorization to MatmulOp

Nicolas Vasilache · Nicolas Vasilache · commit 64c4dcb5eef2 · 2020-01-17T17:09:47.000-05:00
Summary: This is a simple extension to allow vectorization to work not only on GenericLinalgOp but more generally across named ops too. For now, this still only vectorizes matmul-like ops but is a step towards more generic vectorization of Linalg ops. Reviewers: ftynse Subscribers: mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, aartbik, liufengdb, llvm-commits Tags: #llvm Differential Revision: https://reviews.llvm.org/D72942
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td
@@ -86,10 +86,11 @@ class LinalgOpToAffineLoops<string OpType> : NativeCodeCall<
 //===----------------------------------------------------------------------===//
 // Linalg to vector patterns precondition and DRR.
 //===----------------------------------------------------------------------===//
-def PreconditionVectorizeGenericLinalgOp : CPred<
-  "succeeded(vectorizeGenericLinalgOpPrecondition(op))">;
-def VectorizeGenericLinalgOp : NativeCodeCall<
-  "vectorizeGenericLinalgOp($_builder, op)">;
+def PreconditionVectorizeLinalgOp : CPred<
+  "succeeded(vectorizeLinalgOpPrecondition(op))">;
+def VectorizeLinalgOp : NativeCodeCall<
+  "vectorizeLinalgOp($_builder, op)">;
+
 
 //===----------------------------------------------------------------------===//
 // Linalg generic permutation patterns precondition and DRR.
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h
@@ -79,9 +79,9 @@ template <typename ConcreteOp>
 LogicalResult linalgOpToAffineLoops(PatternRewriter &rewriter, Operation *op);
 
 /// Rewrite a linalg.generic into a suitable vector.contraction op.
-LogicalResult vectorizeGenericLinalgOpPrecondition(Operation *op);
-SmallVector<Value, 0> vectorizeGenericLinalgOp(PatternRewriter &rewriter,
-                                               Operation *op);
+LogicalResult vectorizeLinalgOpPrecondition(Operation *op);
+SmallVector<Value, 0> vectorizeLinalgOp(PatternRewriter &rewriter,
+                                        Operation *op);
 
 /// Emits a `generic` or `indexed_generic` operation with the `indexing_maps`
 /// and `iterator_types` permutated according to `permutation`.
diff --git a/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp
@@ -158,10 +158,20 @@ static bool isMatmul(linalg::GenericOp genericOp) {
          genericOp.indexing_maps() == maps && hasMultiplyAddBody(genericOp);
 }
 
-LogicalResult
-mlir::linalg::vectorizeGenericLinalgOpPrecondition(Operation *op) {
-  // TODO(ntv): This is in fact much more general than just vectorization for
-  // matmul ops.
+// TODO(ntv): This is in fact much more general than just vectorization for
+// matmul ops.
+LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
+  auto linalgOp = cast<linalg::LinalgOp>(op);
+  // All types must be static shape to go to vector.
+  for (Value operand : linalgOp.getInputsAndOutputBuffers())
+    if (!operand.getType().cast<ShapedType>().hasStaticShape())
+      return failure();
+  for (Type outputTensorType : linalgOp.getOutputTensorTypes())
+    if (!outputTensorType.cast<ShapedType>().hasStaticShape())
+      return failure();
+  if (isa<linalg::MatmulOp>(op))
+    return success();
+
   auto genericOp = dyn_cast<linalg::GenericOp>(op);
   if (!genericOp || !isMatmul(genericOp))
     return failure();
@@ -179,30 +189,29 @@ mlir::linalg::vectorizeGenericLinalgOpPrecondition(Operation *op) {
   return success();
 }
 
-SmallVector<Value, 0>
-mlir::linalg::vectorizeGenericLinalgOp(PatternRewriter &rewriter,
-                                       Operation *op) {
+SmallVector<Value, 0> mlir::linalg::vectorizeLinalgOp(PatternRewriter &rewriter,
+                                                      Operation *op) {
   LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE
                        "]: Rewrite linalg op as vector.contract: "
                     << *op << ":\n");
 
-  assert(succeeded(vectorizeGenericLinalgOpPrecondition(op)) &&
+  assert(succeeded(vectorizeLinalgOpPrecondition(op)) &&
          "DRR failure case must be a precondition");
 
-  auto genericOp = cast<linalg::GenericOp>(op);
-  assert(genericOp.hasBufferSemantics() &&
+  auto linalgOp = cast<linalg::LinalgOp>(op);
+  assert(linalgOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   edsc::ScopedContext scope(rewriter, op->getLoc());
   using edsc::intrinsics::std_load;
   using edsc::intrinsics::std_store;
   using vector_contract = edsc::intrinsics::ValueBuilder<vector::ContractionOp>;
   using vector_type_cast = edsc::intrinsics::ValueBuilder<vector::TypeCastOp>;
-  auto vA = std_load(vector_type_cast(genericOp.getInput(0)));
-  auto vB = std_load(vector_type_cast(genericOp.getInput(1)));
-  auto vectorMemRefC = vector_type_cast(genericOp.getOutputBuffer(0));
+  auto vA = std_load(vector_type_cast(linalgOp.getInput(0)));
+  auto vB = std_load(vector_type_cast(linalgOp.getInput(1)));
+  auto vectorMemRefC = vector_type_cast(linalgOp.getOutputBuffer(0));
   auto vC = std_load(vectorMemRefC);
-  auto vRes = vector_contract(vA, vB, vC, genericOp.indexing_maps(),
-                              genericOp.iterator_types());
+  auto vRes = vector_contract(vA, vB, vC, linalgOp.indexing_maps(),
+                              linalgOp.iterator_types());
   std_store(vRes, vectorMemRefC);
   return {};
 }
diff --git a/mlir/test/Dialect/Linalg/transform-patterns.mlir b/mlir/test/Dialect/Linalg/transform-patterns.mlir
@@ -173,7 +173,7 @@ func @fusion_test(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
     affine_map<(m, n, k) -> (m, n)>
   ],
   iterator_types = ["parallel", "parallel", "reduction"],
-  __internal_linalg_transform__ = "_marked_matmul_"
+  __internal_linalg_transform__ = "VECTORIZE"
 }
 func @vectorization_test(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
                          %C: memref<8x32xf32>) {
@@ -185,7 +185,6 @@ func @vectorization_test(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
   } : memref<8x16xf32>, memref<16x32xf32>, memref<8x32xf32>
   return
 }
-
 // CHECK-LABEL: func @vectorization_test
 //       CHECK: vector.type_cast %{{.*}} : memref<8x16xf32> to memref<vector<8x16xf32>>
 //       CHECK: load %{{.*}}[] : memref<vector<8x16xf32>>
@@ -195,6 +194,17 @@ func @vectorization_test(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
 //       CHECK: load %{{.*}}[] : memref<vector<8x32xf32>>
 //       CHECK: vector.contract {indexing_maps = [#[[mk]], #[[kn]], #[[mn]]], iterator_types = ["parallel", "parallel", "reduction"]} %{{.*}}, %{{.*}}, %{{.*}} : vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
 //       CHECK: store %{{.*}}, %{{.*}}[] : memref<vector<8x32xf32>>
+
+func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
+                         %C: memref<8x32xf32>) {
+  linalg.matmul(%A, %B, %C) { __internal_linalg_transform__ = "VECTORIZE"} :
+    memref<8x16xf32>, memref<16x32xf32>, memref<8x32xf32>
+  return
+}
+// CHECK-LABEL: func @vectorization_test_2
+//       CHECK: vector.contract {{.*}} :
+//                vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
+
 func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
           %d = mulf %a, %b: f32
           %e = addf %c, %d: f32
@@ -213,7 +223,6 @@ func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
           library_call = "linalg_matmul",
           iterator_types = ["parallel", "parallel", "reduction"]
         }
-
 func @permute_generic(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
            %B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
            %C: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
diff --git a/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td b/mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td
@@ -99,11 +99,17 @@ def : Pattern<(DotOp:$op $_, $_, $_),
 //===----------------------------------------------------------------------===//
 // Linalg to vector contraction patterns.
 //===----------------------------------------------------------------------===//
+def : Pattern<(MatmulOp:$op $_, $_, $_),
+              [(VectorizeLinalgOp)],
+              [(Constraint<And<[
+                HasLinalgTransformMarker<"VECTORIZE">,
+                PreconditionVectorizeLinalgOp
+               ]>>)]>;
 def : Pattern<(GenericOp:$op $_, $_, $_, $_, $_, $_, $_, $_),
-             [(VectorizeGenericLinalgOp)],
-             [(Constraint<And<[
-               HasLinalgTransformMarker<"_marked_matmul_">,
-               PreconditionVectorizeGenericLinalgOp
+              [(VectorizeLinalgOp)],
+              [(Constraint<And<[
+                HasLinalgTransformMarker<"VECTORIZE">,
+                PreconditionVectorizeLinalgOp
                ]>>)]>;
 
 //===----------------------------------------------------------------------===//
@@ -135,4 +141,5 @@ def : Pat<(MatmulOp:$op $_, $_, $_),
               HasOperandsOfType<"SubViewOp">,
               HasLinalgTransformMarker<"_promote_views_">]>>
            )]>;
+
 #endif // TEST_LINALG_TRANSFORMS_PATTERNS
