[mlir][linalg] Support generalization of generic.

[gysit]

Summary

Adapt the generalization pattern to increment the transformation marker if generalization is not needed. Add a test to exercise the behavior in combination with subsequent interchange.

[gysit]

Updated comment.

[gysit]

As discussed offline, let`s abandon this revision until we have a better solution not relying on markers.

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

 using namespace mlir;
 using namespace mlir::linalg;
 
 LogicalResult mlir::linalg::generalizeNamedOpPrecondition(Operation *op) {
   LinalgOp namedOp = dyn_cast<LinalgOp>(op);
   // Check if the operation is a LinalgOp but not a GenericOp.
   if (!namedOp || isa<GenericOp>(op))
     return failure();
-  // Check if the operation has a region builder.
-  if (!namedOp.getRegionBuilder())
-    return failure();
   return success();
 }
 
 GenericOp mlir::linalg::generalizeNamedOp(PatternRewriter &rewriter,
                                           LinalgOp namedOp) {
   SmallVector<Value> inputOperands = namedOp.getInputOperands();
   SmallVector<Value> outputOperands = namedOp.getOutputOperands();
   SmallVector<AffineMap> indexingMaps = namedOp.getIndexingMaps();

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

void runOnFunction() override {
     RewritePatternSet generalizationPattern(funcOp.getContext());
     if (!anchorOpName.empty()) {
       generalizationPattern.add<LinalgGeneralizationPattern>(
           anchorOpName, funcOp.getContext(), filter);
     } else {
       generalizationPattern.add<LinalgGeneralizationPattern>(
           funcOp.getContext(), filter);
     }
-    if (failed(applyPatternsAndFoldGreedily(funcOp,
-                                            std::move(generalizationPattern))))
-      signalPassFailure();
+    (void)applyPatternsAndFoldGreedily(funcOp,
+                                       std::move(generalizationPattern));
   }
 
   LinalgTransformationFilter filter;
 };
 
 /// Configurable pass to apply lowering of coarser-grained named linalg ops into
 /// finer-grained named versions.
 struct LinalgStrategyDecomposePass

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

mlir::linalg::LinalgGeneralizationPattern::LinalgGeneralizationPattern(
     StringRef opName, MLIRContext *context, LinalgTransformationFilter filter,
     PatternBenefit benefit)
     : RewritePattern(opName, benefit, context, {}), filter(filter) {}
 
 LogicalResult mlir::linalg::LinalgGeneralizationPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   if (failed(filter.checkAndNotify(rewriter, op)))
     return failure();
+  if (isa<GenericOp>(op)) {
+    filter.replaceLinalgTransformationFilter(rewriter, op);
+    return success();
+  }
   if (failed(generalizeNamedOpPrecondition(op)))
     return failure();
 
   GenericOp genericOp = generalizeNamedOp(rewriter, op);
   rewriter.replaceOp(op, genericOp.getResults());
   filter.replaceLinalgTransformationFilter(rewriter, genericOp);
   return success();
 }

mlir/test/Dialect/Linalg/codegen-strategy.mlir

 // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=2,4,8 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" -split-input-file | FileCheck %s --check-prefix=CHECK-INTRINSIC
 // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" -split-input-file | FileCheck %s --check-prefix=CHECK-OUTER
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 tile-interchange=1,2,0 generalize iterator-interchange=0,2,1" -split-input-file | FileCheck %s --check-prefix=CHECK-INTERCHANGE
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 tile-interchange=1,2,0 generalize iterator-interchange=0,2,1" -split-input-file | FileCheck %s --check-prefix=CHECK-TILE-AND-INTERCHANGE
 // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 pad pack-paddings=1,1,0 hoist-paddings=3,3,0" -split-input-file | FileCheck %s --check-prefix=CHECK-PAD
 // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 fuse pad vectorize" -split-input-file | FileCheck %s --check-prefix=CHECK-FUSE
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matvec generalize iterator-interchange=1,0" -split-input-file | FileCheck %s --check-prefix=CHECK-GENERALIZE
 
 // CHECK-INTRINSIC: func @matmul(
 //     CHECK-OUTER: func @matmul(
 func @matmul(%arg0: memref<72x72xf32>, %arg1: memref<72x72xf32>, %arg2: memref<72x72xf32>) {
 
   // Check the matrix intrinsic lowering is triggered.
   //      CHECK-INTRINSIC: vector.matrix_multiply
   // CHECK-INTRINSIC-SAME: {lhs_columns = 8 : i32, lhs_rows = 2 : i32, rhs_columns = 4 : i32}
   // CHECK-INTRINSIC-SAME: (vector<16xf32>, vector<32xf32>) -> vector<8xf32>
 
   // Check the outer product lowering is triggered.
   //          CHECK-OUTER: vector.outerproduct {{.*}} : vector<2xf32>, vector<4xf32>
   linalg.matmul ins(%arg0, %arg1: memref<72x72xf32>, memref<72x72xf32>) outs(%arg2: memref<72x72xf32>)
   return
 }
 
 // -----
 
-// CHECK-INTERCHANGE: func @matmul(
+// CHECK-TILE-AND-INTERCHANGE: func @matmul(
 func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> {
-  //  CHECK-INTERCHANGE-DAG: %[[C16:.*]] = arith.constant 16
-  //  CHECK-INTERCHANGE-DAG: %[[C32:.*]] = arith.constant 32
-  //  CHECK-INTERCHANGE-DAG: %[[C64:.*]] = arith.constant 64
+  //  CHECK-TILE-AND-INTERCHANGE-DAG: %[[C16:.*]] = arith.constant 16
+  //  CHECK-TILE-AND-INTERCHANGE-DAG: %[[C32:.*]] = arith.constant 32
+  //  CHECK-TILE-AND-INTERCHANGE-DAG: %[[C64:.*]] = arith.constant 64
 
   // Check the tile loops are interchanged.
-  //      CHECK-INTERCHANGE: scf.for {{.*}} step %[[C32]]
-  //      CHECK-INTERCHANGE:   scf.for {{.*}} step %[[C64]]
-  //      CHECK-INTERCHANGE:    scf.for {{.*}} step %[[C16]]
+  //      CHECK-TILE-AND-INTERCHANGE: scf.for {{.*}} step %[[C32]]
+  //      CHECK-TILE-AND-INTERCHANGE:   scf.for {{.*}} step %[[C64]]
+  //      CHECK-TILE-AND-INTERCHANGE:    scf.for {{.*}} step %[[C16]]
 
   // Check the operation has been generalized and interchanged.
-  //      CHECK-INTERCHANGE:      linalg.generic
-  // CHECK-INTERCHANGE-SAME:      iterator_types = ["parallel", "reduction", "parallel"]
+  //      CHECK-TILE-AND-INTERCHANGE:      linalg.generic
+  // CHECK-TILE-AND-INTERCHANGE-SAME:      iterator_types = ["parallel", "reduction", "parallel"]
   %0 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%arg2: tensor<72x72xf32>) -> tensor<72x72xf32>
   return %0 : tensor<72x72xf32>
 }
 
 // -----
 
 //     CHECK-PAD-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0) -> (16, -d0 + 72)>
 
func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> {
   %cst = arith.constant 0.0 : f32
   %0 = linalg.fill(%cst, %arg0) : f32, tensor<72x72xf32> -> tensor<72x72xf32>
 
   // Check the matmul is padded and vectorized despite the empty anchor op string.
   //        CHECK-FUSE:  vector.outerproduct
   %1 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%0: tensor<72x72xf32>) -> tensor<72x72xf32>
   return %1 : tensor<72x72xf32>
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-GENERALIZE: func @matvec(
+func @matvec(%arg0: tensor<72x72xf32>, %arg1: tensor<72xf32>, %arg2: tensor<72xf32>) -> tensor<72xf32> {
+  // Check the generic op iterators are interchanged although generalization was not needed.
+  //      CHECK-GENERALIZE:      linalg.generic
+  // CHECK-GENERALIZE-SAME:      iterator_types = ["reduction", "parallel"]
+  %0 = linalg.generic {indexing_maps = [#map0, #map1],
+                       iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<72x72xf32>)
+      outs(%arg1 : tensor<72xf32>) {
+    ^bb0(%arg3: f32, %arg5: f32):       // no predecessors
+      %2 = arith.addf %arg3, %arg5 : f32
+      linalg.yield %2 : f32
+  } -> tensor<72xf32>
+
+  // Check matvec is generalized and its iterators are interchanged.
+  //      CHECK-GENERALIZE:      linalg.generic
+  // CHECK-GENERALIZE-SAME:      iterator_types = ["reduction", "parallel"]
+  %1 = linalg.matvec ins(%arg0, %0: tensor<72x72xf32>, tensor<72xf32>) outs(%arg2: tensor<72xf32>) -> tensor<72xf32>
+  return %1 : tensor<72xf32>
+}