[mlir][Linalg] Extend linalg vectorization to MatmulOp

[nicolasvasilache]

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.

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[ftynse]

The change is mostly renaming/clean-up. I only have minor comments, feel free to land after addressing.

mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp

192–193

Why SmallVector<Value, 0> ? Maybe just use std::vector if you don't need stack elements.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp

192–193

This connects to RewriterGen.cpp in subtle ways and the only thing I have found to work with brace initializations and without changing RewriterGen too much is this.
Ther eis an internal bug I filed if you want to try and make the world a better place.
For me, I have shelved it for now :)

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td

 
 class LinalgOpToAffineLoops<string OpType> : NativeCodeCall<
   "if (failed(linalgOpToAffineLoops<" # OpType # ">($_builder, op))) " #
   "  return matchFailure();">;
 
 //===----------------------------------------------------------------------===//
 // 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.
 //===----------------------------------------------------------------------===//
 class PreconditionPermuteGenericLinalgOp<list<int> permutation> : CPred<
   "succeeded(permuteGenericLinalgOpPrecondition(op, {" #
   StrJoinInt<permutation>.result # "}))">;
 class PermuteGenericLinalgOp<list<int> permutation, string value> :

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h

 template <typename ConcreteOp>
 LogicalResult linalgOpToLoops(PatternRewriter &rewriter, Operation *op);
 
 /// Emits a loop nest of `affine.for` with the proper body for `op`.
 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,
 
 /// Emits a `generic` or `indexed_generic` operation with the `indexing_maps`
 /// and `iterator_types` permutated according to `permutation`.
 LogicalResult
 permuteGenericLinalgOpPrecondition(Operation *op,
                                    ArrayRef<unsigned> permutation);
 SmallVector<Value, 0> permuteGenericLinalgOp(PatternRewriter &rewriter,
                                              Operation *op,

mlir/lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp

static bool isMatmul(linalg::GenericOp genericOp) {
   auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}));
   auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}));
   auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}));
   auto maps = ArrayAttr::get({mapA, mapB, mapC}, ctx);
   return genericOp.getNumInputs() == 2 && genericOp.getNumOutputs() == 1 &&
          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.
+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();
 
   // TODO(ntv): non-identity layout.
   auto isStaticMemRefWithIdentityLayout = [](Value v) {
     auto m = v.getType().dyn_cast<MemRefType>();
     if (!m || !m.hasStaticShape() || !m.getAffineMaps().empty())
       return false;
     return true;
   };
   if (!llvm::all_of(genericOp.getInputsAndOutputBuffers(),
                     isStaticMemRefWithIdentityLayout))
     return failure();
   return success();
 }
 
-SmallVector<Value, 0>
-mlir::linalg::vectorizeGenericLinalgOp(PatternRewriter &rewriter,
+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 {};
 }
 
 //============================================================================//
 // Precondition and transformation for permutation of Linalg generic ops.
 //============================================================================//
 LogicalResult mlir::linalg::permuteGenericLinalgOpPrecondition(

mlir/test/Dialect/Linalg/transform-patterns.mlir

#matmul_trait = {
   args_in = 2,
   args_out = 1,
   indexing_maps = [
     affine_map<(m, n, k) -> (m, k)>,
     affine_map<(m, n, k) -> (k, n)>,
     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>) {
   linalg.generic #matmul_trait %A, %B, %C {
     ^bb(%a: f32, %b: f32, %c: f32) :
       %d = mulf %a, %b: f32
       %e = addf %c, %d: f32
       linalg.yield %e : f32
   } : 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>>
 //       CHECK: vector.type_cast %{{.*}} : memref<16x32xf32> to memref<vector<16x32xf32>>
 //       CHECK: load %{{.*}}[] : memref<vector<16x32xf32>>
 //       CHECK: vector.type_cast %{{.*}} : memref<8x32xf32> to memref<vector<8x32xf32>>
 //       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
           return %e: f32
         }
 #matmul_accesses = [
           affine_map<(m, n, k) -> (m, k)>,
           affine_map<(m, n, k) -> (k, n)>,
           affine_map<(m, n, k) -> (m, n)>
 ]
 #generic_matmul_trait = {
           args_in = 2,
           args_out = 1,
           fun = @fma,
           indexing_maps = #matmul_accesses,
           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]>) {
   linalg.generic #generic_matmul_trait %A, %B, %C : memref<?x?xf32, offset: ?, strides: [?, 1]>, memref<?x?xf32, offset: ?, strides: [?, 1]>, memref<?x?xf32, offset: ?, strides: [?, 1]>
 
   return
 }
 // CHECK-LABEL : func @fma

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

 //===----------------------------------------------------------------------===//
 def : Pattern<(DotOp:$op $_, $_, $_),
               [(LinalgOpToLoops<"DotOp">)],
               [(Constraint<HasLinalgTransformMarker<"REG">>)]>;
 
 //===----------------------------------------------------------------------===//
 // Linalg to vector contraction patterns.
 //===----------------------------------------------------------------------===//
+def : Pattern<(MatmulOp:$op $_, $_, $_),
+              [(VectorizeLinalgOp)],
+              [(Constraint<And<[
+                HasLinalgTransformMarker<"VECTORIZE">,
+                PreconditionVectorizeLinalgOp
+               ]>>)]>;
 def : Pattern<(GenericOp:$op $_, $_, $_, $_, $_, $_, $_, $_),
-             [(VectorizeGenericLinalgOp)],
+              [(VectorizeLinalgOp)],
               [(Constraint<And<[
-               HasLinalgTransformMarker<"_marked_matmul_">,
-               PreconditionVectorizeGenericLinalgOp
+                HasLinalgTransformMarker<"VECTORIZE">,
+                PreconditionVectorizeLinalgOp
                ]>>)]>;
 
 //===----------------------------------------------------------------------===//
 // Linalg generic permutation patterns.
 //===----------------------------------------------------------------------===//
 def : Pat<(GenericOp:$op $_, $_, $_, $_, $_, $_, $_, $_),
               (PermuteGenericLinalgOp<[1, 2, 0], "PERMUTE"> $op),
               [(Constraint<And<[

 //===----------------------------------------------------------------------===//
 def : Pat<(MatmulOp:$op $_, $_, $_),
           (PromoteSubviewsLinalgOp),
           [(Constraint<And<[
               PreconditionPromoteSubviewsLinalgOp,
               HasOperandsOfType<"SubViewOp">,
               HasLinalgTransformMarker<"_promote_views_">]>>
            )]>;
+
 #endif // TEST_LINALG_TRANSFORMS_PATTERNS