[mlir][Linalg] Retire Linalg generic interchange pattern and pass

[guraypp]

Summary

This revision removes the Linalg generic interchange pattern and pass.

It also changes transform-patterns test to make use of transform dialect.

Context: https://discourse.llvm.org/t/psa-retire-linalg-filter-based-patterns/63785

[nicolasvasilache]

Thanks!

[guraypp]

run clang-format

[guraypp]

merged branch

[guraypp]

merge branch

mlir/include/mlir/Dialect/Linalg/Passes.h

StringRef opName = "", const linalg::LinalgTransformationFilter &filter =
                                linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyDecomposePass.
 // TODO: if/when we need finer control add an `opName` parameter.
 std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyDecomposePass(
     const linalg::LinalgTransformationFilter &filter =
         linalg::LinalgTransformationFilter());
 
-/// Create a LinalgStrategyInterchangePass.
-std::unique_ptr<OperationPass<func::FuncOp>>
-createLinalgStrategyInterchangePass(
-    ArrayRef<int64_t> iteratorInterchange = {},
-    const linalg::LinalgTransformationFilter &filter =
-        linalg::LinalgTransformationFilter());
-
 /// Create a LinalgStrategyPeelPass.
 std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyPeelPass(
     StringRef opName = "",
     const linalg::LinalgPeelOptions &opt = linalg::LinalgPeelOptions(),
     const linalg::LinalgTransformationFilter &filter =
         linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyVectorizePass.

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgStrategyDecomposePass
   let constructor = "mlir::createLinalgStrategyDecomposePass()";
   let dependentDialects = ["linalg::LinalgDialect"];
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,
   ];
 }
 
-def LinalgStrategyInterchangePass
-    : Pass<"linalg-strategy-interchange-pass", "func::FuncOp"> {
-  let summary = "Configurable pass to apply pattern-based iterator interchange.";
-  let constructor = "mlir::createLinalgStrategyInterchangePass()";
-  let dependentDialects = ["linalg::LinalgDialect"];
-  let options = [
-    Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
-      "Which func op is the anchor to latch on.">,
-  ];
-}
-
 def LinalgStrategyPeelPass
     : Pass<"linalg-strategy-peel-pass", "func::FuncOp"> {
   let summary = "Configurable pass to apply pattern-based linalg peeling.";
   let constructor = "mlir::createLinalgStrategyPeelPass()";
   let dependentDialects = [
     "linalg::LinalgDialect",
     "scf::SCFDialect"
   ];

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

void addToPassPipeline(OpPassManager &pm,
                          LinalgTransformationFilter m) const override {
     pm.addPass(createLinalgStrategyGeneralizePass(opName, m));
   }
 
 private:
   std::string opName;
 };
 
-/// Represent one application of createLinalgStrategyInterchangePass.
-struct Interchange : public Transformation {
-  explicit Interchange(ArrayRef<int64_t> iteratorInterchange,
-                       LinalgTransformationFilter::FilterFunction f = nullptr)
-      : Transformation(std::move(f)),
-        iteratorInterchange(iteratorInterchange.begin(),
-                            iteratorInterchange.end()) {}
-
-  void addToPassPipeline(OpPassManager &pm,
-                         LinalgTransformationFilter m) const override {
-    pm.addPass(createLinalgStrategyInterchangePass(iteratorInterchange, m));
-  }
-
-private:
-  SmallVector<int64_t> iteratorInterchange;
-};
-
 /// Represent one application of createLinalgStrategyDecomposePass.
 struct Decompose : public Transformation {
   explicit Decompose(LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(std::move(f)) {}
 
   void addToPassPipeline(OpPassManager &pm,
                          LinalgTransformationFilter m) const override {
     pm.addPass(createLinalgStrategyDecomposePass(m));
generalize(StringRef opName,
     return *this;
   }
   /// Conditionally append a pattern to generalize named operations.
   CodegenStrategy &
   generalizeIf(bool b, StringRef opName,
                LinalgTransformationFilter::FilterFunction f = nullptr) {
     return b ? generalize(opName, std::move(f)) : *this;
   }
-  /// Append a pattern to interchange iterators.
-  CodegenStrategy &
-  interchange(ArrayRef<int64_t> iteratorInterchange,
-              const LinalgTransformationFilter::FilterFunction &f = nullptr) {
-    transformationSequence.emplace_back(
-        std::make_unique<Interchange>(iteratorInterchange, f));
-    return *this;
-  }
-  /// Conditionally append a pattern to interchange iterators.
-  CodegenStrategy &
-  interchangeIf(bool b, ArrayRef<int64_t> iteratorInterchange,
-                LinalgTransformationFilter::FilterFunction f = nullptr) {
-    return b ? interchange(iteratorInterchange, std::move(f)) : *this;
-  }
   /// Append patterns to decompose convolutions.
   CodegenStrategy &
   decompose(const LinalgTransformationFilter::FilterFunction &f = nullptr) {
     transformationSequence.emplace_back(std::make_unique<Decompose>(f));
     return *this;
   }
   /// Conditionally append patterns to decompose convolutions.
   CodegenStrategy &

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

 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
   /// Tile sizes and interchange used to tile the root operation.
   LinalgTilingAndFusionOptions options;
 };
 
 ///
-/// Linalg generic interchange pattern.
-///
-/// Apply the `interchange` transformation on a RewriterBase.
-/// `filter` controls LinalgTransformMarker matching and update when specified.
-/// See `interchange` for more details.
-struct GenericOpInterchangePattern : public OpRewritePattern<GenericOp> {
-  using OpRewritePattern<GenericOp>::OpRewritePattern;
-
-  /// GenericOp-specific constructor with an optional `filter`.
-  GenericOpInterchangePattern(
-      MLIRContext *context, ArrayRef<unsigned> interchangeVector,
-      LinalgTransformationFilter f = LinalgTransformationFilter(),
-      PatternBenefit benefit = 1);
-
-  /// `matchAndRewrite` implementation that returns the significant transformed
-  /// pieces of IR.
-  FailureOr<GenericOp>
-  returningMatchAndRewrite(GenericOp op, PatternRewriter &rewriter) const;
-
-  LogicalResult matchAndRewrite(GenericOp op,
-                                PatternRewriter &rewriter) const override {
-    return returningMatchAndRewrite(op, rewriter);
-  }
-
-private:
-  /// LinalgTransformMarker handles special attribute manipulations.
-  LinalgTransformationFilter filter;
-  /// The interchange vector to reorder the iterators and indexing_maps dims.
-  SmallVector<unsigned, 8> interchangeVector;
-};
-
-///
 /// Linalg generalization pattern.
 ///
 /// Apply the `generalization` transformation as a pattern.
 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `generalization` for more details.
 struct LinalgGeneralizationPattern
     : public OpInterfaceRewritePattern<LinalgOp> {
   /// Construct a generic pattern applied to all LinalgOp that verify `filter`.

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

void runOnOperation() override {
     if (failed(applyPatternsAndFoldGreedily(funcOp,
                                             std::move(decompositionPattern))))
       signalPassFailure();
   }
 
   LinalgTransformationFilter filter;
 };
 
-/// Configurable pass to apply pattern-based linalg generalization.
-struct LinalgStrategyInterchangePass
-    : public LinalgStrategyInterchangePassBase<LinalgStrategyInterchangePass> {
-
-  LinalgStrategyInterchangePass() = default;
-
-  LinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
-                                LinalgTransformationFilter filter)
-      : iteratorInterchange(iteratorInterchange.begin(),
-                            iteratorInterchange.end()),
-        filter(std::move(filter)) {}
-
-  void runOnOperation() override {
-    auto funcOp = getOperation();
-    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();
-  }
-
-  SmallVector<int64_t> iteratorInterchange;
-  LinalgTransformationFilter filter;
-};
-
 /// Configurable pass to apply pattern-based linalg peeling.
 struct LinalgStrategyPeelPass
     : public LinalgStrategyPeelPassBase<LinalgStrategyPeelPass> {
 
   LinalgStrategyPeelPass() = default;
 
   LinalgStrategyPeelPass(StringRef opName, LinalgPeelOptions opt,
                          LinalgTransformationFilter filt)

 /// Create a LinalgStrategyDecomposePass.
 // TODO: if/when we need finer control add an `opName` parameter.
 std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::createLinalgStrategyDecomposePass(
     const LinalgTransformationFilter &filter) {
   return std::make_unique<LinalgStrategyDecomposePass>(filter);
 }
 
-/// Create a LinalgStrategyInterchangePass.
-std::unique_ptr<OperationPass<func::FuncOp>>
-mlir::createLinalgStrategyInterchangePass(
-    ArrayRef<int64_t> iteratorInterchange,
-    const LinalgTransformationFilter &filter) {
-  return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
-                                                         filter);
-}
-
 /// Create a LinalgStrategyPeelPass.
 std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::createLinalgStrategyPeelPass(StringRef opName,
                                    const LinalgPeelOptions &opt,
                                    const LinalgTransformationFilter &filter) {
   return std::make_unique<LinalgStrategyPeelPass>(opName, opt, filter);
 }
 

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

mlir::linalg::LinalgTileAndFuseTensorOpsPattern::returningMatchAndRewrite(
   rootOp->replaceAllUsesWith(tileLoopNest->getRootOpReplacementResults());
 
   // Apply the filter if specified.
   for (LinalgOp linalgOp : tileLoopNest->getAllTiledAndFusedOps())
     filter.replaceLinalgTransformationFilter(rewriter, linalgOp);
   return tileLoopNest;
 }
 
-/// Linalg generic interchange pattern.
-mlir::linalg::GenericOpInterchangePattern::GenericOpInterchangePattern(
-    MLIRContext *context, ArrayRef<unsigned> interchangeVector,
-    LinalgTransformationFilter f, PatternBenefit benefit)
-    : OpRewritePattern(context, benefit), filter(std::move(f)),
-      interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
-
-FailureOr<GenericOp>
-mlir::linalg::GenericOpInterchangePattern::returningMatchAndRewrite(
-    GenericOp genericOp, PatternRewriter &rewriter) const {
-  if (failed(filter.checkAndNotify(rewriter, genericOp)))
-    return failure();
-
-  FailureOr<GenericOp> transformedOp =
-      interchangeGenericOp(rewriter, genericOp, interchangeVector);
-  if (failed(transformedOp))
-    return failure();
-
-  // New filter if specified.
-  filter.replaceLinalgTransformationFilter(rewriter, genericOp);
-  return transformedOp;
-}
-
 /// Linalg generalization pattern.
 mlir::linalg::LinalgGeneralizationPattern::LinalgGeneralizationPattern(
     MLIRContext *context, LinalgTransformationFilter f, PatternBenefit benefit)
     : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
       filter(std::move(f)) {}
 
 mlir::linalg::LinalgGeneralizationPattern::LinalgGeneralizationPattern(
     StringRef opName, MLIRContext *context, LinalgTransformationFilter f,

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

-// RUN: mlir-opt %s -test-linalg-transform-patterns=test-patterns -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-patterns -split-input-file -test-transform-dialect-interpreter | FileCheck %s
 
 // CHECK-DAG: #[[$STRIDED_1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // Map corresponding to a 2D memory access where the stride along the last dim is known to be 1.
 // CHECK-DAG: #[[$STRIDED_2D_u_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECK-DAG: #[[$kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
 // CHECK-DAG: #[[$nm:.*]] = affine_map<(d0, d1, d2) -> (d1, d0)>
 // CHECK-DAG: #[[$km:.*]] = affine_map<(d0, d1, d2) -> (d2, d0)>
 
func.func @permute_generic(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
    outs(%C : memref<?x?xf32, offset: ?, strides: [?, 1]>) {
     ^bb(%a: f32, %b: f32, %c: f32):
       %d = arith.mulf %a, %b: f32
       %e = arith.addf %c, %d: f32
       linalg.yield %e: f32
   }
   return
 }
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+    transform.structured.interchange %0 { iterator_interchange = [1, 2, 0]}
+  }
+}
 // CHECK-LABEL:  func @permute_generic
 // CHECK:        linalg.generic {
 // CHECK-SAME:   indexing_maps = [#[[$kn]], #[[$nm]], #[[$km]]],
 // CHECK-SAME:   iterator_types = ["parallel", "reduction", "parallel"],
 // CHECK-SAME:   library_call = "linalg_matmul"}
 // CHECK:          memref<?x?xf32, #[[$STRIDED_2D_u_1]]>,
 // CHECK-SAME:     memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
 // CHECK-SAME:     memref<?x?xf32, #[[$STRIDED_2D_u_1]]>

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

static void applyPatterns(func::FuncOp funcOp) {
   //===--------------------------------------------------------------------===//
   // Linalg to vector contraction patterns.
   //===--------------------------------------------------------------------===//
   patterns.add<LinalgVectorizationPattern>(
       ctx, LinalgTransformationFilter(StringAttr::get(ctx, "VECTORIZE"))
                .addOpFilter<MatmulOp, FillOp, GenericOp>());
   patterns.add<CopyVectorizationPattern>(ctx);
 
-  //===--------------------------------------------------------------------===//
-  // Linalg generic interchange pattern.
-  //===--------------------------------------------------------------------===//
-  patterns.add<GenericOpInterchangePattern>(
-      ctx,
-      /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
-      LinalgTransformationFilter(ArrayRef<StringAttr>{},
-                                 StringAttr::get(ctx, "PERMUTED")));
-
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });
 }