[mlir][linalg] Update transformation marker if no other pattern matches.

[gysit]

Summary

The Linalg Codegenstrategy uses transformation markers to keep track of which transformations have been applied to a specific operation. When working with multiple operations, for example after fusion, it may happen that specific transformations apply only to a subset of the operations. In this scenario, we nevertheless want to increment the transformation markers of all fused operations. This revision introduced a LinalgSkipPattern that has a low benefit and updates all operations that are not matched by another pattern of the current transformation stage. Only if the transformation markers of all operations are up to date, follow up transformations will run.

Depends On D114855

[gysit]

Improve comment.

[nicolasvasilache]

This looks like a footgun to me, have we reached the point where we should outline to an scf.execute op?

[gysit]

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

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

LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
 };
 
 ///
+/// Linalg skip pattern.
+///
+/// Apply no transformation and if `filter` matches increment the
+/// LinalgTransformMarker. The pattern has a low benefit and applies only if no
+/// other pattern of the current transformation stage matches.
+struct LinalgSkipPattern : public OpInterfaceRewritePattern<LinalgOp> {
+  LinalgSkipPattern(MLIRContext *context, LinalgTransformationFilter filter,
+                    PatternBenefit benefit = 0);
+
+  LogicalResult matchAndRewrite(LinalgOp linalgOp,
+                                PatternRewriter &rewriter) const override;
+
+private:
+  /// LinalgTransformMarker handles special attribute manipulations.
+  LinalgTransformationFilter filter;
+};
+
+///
 /// Linalg promotion patterns.
 ///
 /// Apply the `promoteSubViews` transformation as a pattern.
 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `promoteSubViews` for more details.
 struct LinalgBasePromotionPattern : public RewritePattern {
   /// Entry point to match any LinalgOp OpInterface.
   /// MatchAnyOpTag-based constructor with a mandatory `filter`.

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

LinalgStrategyDecomposePass(LinalgTransformationFilter filter)
       : filter(filter) {}
 
   void runOnFunction() override {
     auto funcOp = getFunction();
     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
       return;
     RewritePatternSet decompositionPattern(funcOp.getContext());
     populateDecomposeConvolutionPatterns(decompositionPattern, filter);
+    // Update the transformation marker if no decompose pattern matches.
+    decompositionPattern.add<LinalgSkipPattern>(funcOp.getContext(), filter);
     if (failed(applyPatternsAndFoldGreedily(funcOp,
                                             std::move(decompositionPattern))))
       signalPassFailure();
   }
 
   LinalgTransformationFilter filter;
 };
 
void runOnFunction() override {
     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
       return;
 
     SmallVector<unsigned> interchangeVector(iteratorInterchange.begin(),
                                             iteratorInterchange.end());
     RewritePatternSet interchangePattern(funcOp.getContext());
     interchangePattern.add<GenericOpInterchangePattern>(
         funcOp.getContext(), interchangeVector, filter);
-    if (failed(applyPatternsAndFoldGreedily(funcOp,
-                                            std::move(interchangePattern))))
-      signalPassFailure();
+    // Update the transformation marker for all non generic operations.
+    interchangePattern.add<LinalgSkipPattern>(funcOp.getContext(), filter);
+    (void)applyPatternsAndFoldGreedily(funcOp, std::move(interchangePattern));
   }
 
   SmallVector<int64_t> iteratorInterchange;
   LinalgTransformationFilter filter;
 };
 
 /// Configurable pass to apply pattern-based linalg promotion.
 struct LinalgStrategyPromotePass

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

if (failed(generalizeNamedOpPrecondition(op)))
     return failure();
 
   GenericOp genericOp = generalizeNamedOp(rewriter, op);
   rewriter.replaceOp(op, genericOp.getResults());
   filter.replaceLinalgTransformationFilter(rewriter, genericOp);
   return success();
 }
 
+// LinalgSkipPattern
+mlir::linalg::LinalgSkipPattern::LinalgSkipPattern(
+    MLIRContext *context, LinalgTransformationFilter filter,
+    PatternBenefit benefit)
+    : mlir::OpInterfaceRewritePattern<LinalgOp>(context, benefit),
+      filter(filter) {}
+
+LogicalResult mlir::linalg::LinalgSkipPattern::matchAndRewrite(
+    LinalgOp linalgOp, PatternRewriter &rewriter) const {
+  if (failed(filter.checkAndNotify(rewriter, linalgOp)))
+    return failure();
+  filter.replaceLinalgTransformationFilter(rewriter, linalgOp);
+  return success();
+}
+
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     MLIRContext *context, LinalgTransformationFilter filter,
     LinalgPromotionOptions options, PatternBenefit benefit)
     : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter),
       options(options) {}
 
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     StringRef opName, MLIRContext *context, LinalgPromotionOptions options,

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-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
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matvec iterator-interchange=1,0 vectorize" -split-input-file | FileCheck %s --check-prefix=CHECK-INTERCHANGE-AND-VECTORIZE
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=conv decompose generalize" -split-input-file | FileCheck %s --check-prefix=CHECK-DECOMPOSE-AND-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}

 
 // -----
 
 #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 both operations are interchanged despite the missing generic generalization pattern.
+  //      CHECK-GENERALIZE:      linalg.generic
+  // CHECK-GENERALIZE-SAME:      iterator_types = ["reduction", "parallel"]
   //      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>
+  %1 = linalg.matvec ins(%arg0, %0: tensor<72x72xf32>, tensor<72xf32>) outs(%arg2: tensor<72xf32>) -> tensor<72xf32>
+  return %1 : tensor<72xf32>
+}
 
-  // Check matvec is generalized and its iterators are interchanged.
-  //      CHECK-GENERALIZE:      linalg.generic
-  // CHECK-GENERALIZE-SAME:      iterator_types = ["reduction", "parallel"]
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-INTERCHANGE-AND-VECTORIZE: func @matvec(
+func @matvec(%arg0: tensor<72x72xf32>, %arg1: tensor<72xf32>, %arg2: tensor<72xf32>) -> tensor<72xf32> {
+
+  // Check both operations are vectorized despite the missing matvec interchange pattern.
+  //      CHECK-INTERCHANGE-AND-VECTORIZE-NOT:      linalg.generic
+  //      CHECK-INTERCHANGE-AND-VECTORIZE-NOT:      linalg.matvec
+  %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>
   %1 = linalg.matvec ins(%arg0, %0: tensor<72x72xf32>, tensor<72xf32>) outs(%arg2: tensor<72xf32>) -> tensor<72xf32>
   return %1 : tensor<72xf32>
 }
+
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-DECOMPOSE-AND-GENERALIZE: func @conv(
+func @conv(%input: tensor<4x1x6x3xf32>, %filter: tensor<1x2x3x8xf32>, %init: tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32> {
+  %cst = arith.constant 0.0 : f32
+
+  // Check both operations are generalized despite the missing fill decompose pattern.
+  //      CHECK-DECOMPOSE-AND-GENERALIZE-NOT:      linalg.fill
+  //      CHECK-DECOMPOSE-AND-GENERALIZE-NOT:      linalg.conv_2d_nhwc_hwcf
+  %0 = linalg.fill(%cst, %init) : f32, tensor<4x1x2x8xf32> -> tensor<4x1x2x8xf32>
+  %1 = linalg.conv_2d_nhwc_hwcf
+    {dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<[3, 2]> : tensor<2xi64>}
+    ins(%input, %filter : tensor<4x1x6x3xf32>, tensor<1x2x3x8xf32>)
+    outs(%0 : tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32>
+  return %1 : tensor<4x1x2x8xf32>
+}