[mlir][linalg] Add loop interchange to CodegenStrategy.

[gysit]

Summary

Add a loop interchange pass and integrate it with CodegenStrategy.

This patch depends on https://reviews.llvm.org/D110728 and https://reviews.llvm.org/D110746.

[nicolasvasilache]

can we s/loop/iterator everywhere?
We are really interchanging the order of iterators in a linalg.generic, there are no loops yet per-se.

[gysit]

Address comments.

[gysit]

Rebase.

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

linalg::LinalgTransformationFilter filter =
         linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyGeneralizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgStrategyGeneralizePass(StringRef opName = "",
                                    linalg::LinalgTransformationFilter filter =
                                        linalg::LinalgTransformationFilter());
 
+/// Create a LinalgStrategyInterchangePass.
+std::unique_ptr<OperationPass<FuncOp>>
+createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange = {},
+                                    linalg::LinalgTransformationFilter filter =
+                                        linalg::LinalgTransformationFilter());
+
 /// Create a LinalgStrategyVectorizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgStrategyVectorizePass(StringRef opName = "",
                                   linalg::LinalgVectorizationOptions opt =
                                       linalg::LinalgVectorizationOptions(),
                                   linalg::LinalgTransformationFilter filter =
                                       linalg::LinalgTransformationFilter());
 

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

def LinalgStrategyGeneralizePass
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,
     Option<"anchorOpName", "anchor-op", "std::string", /*default=*/"",
       "Which linalg op within the func is the anchor to latch on.">,
   ];
 }
 
+def LinalgStrategyInterchangePass
+    : FunctionPass<"linalg-strategy-interchange-pass"> {
+  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 LinalgStrategyVectorizePass
     : FunctionPass<"linalg-strategy-vectorize-pass"> {
   let summary = "Configurable pass to apply pattern-based linalg vectorization.";
   let constructor = "mlir::createLinalgStrategyVectorizePass()";
   let dependentDialects = ["linalg::LinalgDialect"];
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,

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(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 createLinalgStrategyVectorizePass.
 struct Vectorize : public Transformation {
   explicit Vectorize(linalg::LinalgVectorizationOptions options,
                      LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(f), opName(), options(options) {}
 
   Vectorize(StringRef name, linalg::LinalgVectorizationOptions options,
             LinalgTransformationFilter::FilterFunction f = nullptr)
struct CodegenStrategy {
   }
   /// Conditionally append a pattern to generalize named operations.
   CodegenStrategy &
   generalizeIf(bool b, StringRef opName,
                LinalgTransformationFilter::FilterFunction f = nullptr) {
     return b ? generalize(opName, f) : *this;
     return *this;
   }
+  /// Append a pattern to interchange iterators.
+  CodegenStrategy &
+  interchange(ArrayRef<int64_t> iteratorInterchange,
+              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, f) : *this;
+    return *this;
+  }
   /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
   CodegenStrategy &
   vectorize(StringRef opName,
             LinalgTransformationFilter::FilterFunction f = nullptr) {
     assert(!opName.empty() && "expected an op name");
     transformationSequence.emplace_back(std::make_unique<Vectorize>(
         opName, linalg::LinalgVectorizationOptions(), f));
     return *this;

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

void runOnFunction() override {
     if (failed(applyPatternsAndFoldGreedily(funcOp,
                                             std::move(generalizationPattern))))
       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(filter) {}
+
+  void runOnFunction() override {
+    auto funcOp = getFunction();
+    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 promotion.
 struct LinalgStrategyPromotePass
     : public LinalgStrategyPromotePassBase<LinalgStrategyPromotePass> {
 
   LinalgStrategyPromotePass() = default;
 
   LinalgStrategyPromotePass(StringRef opName, LinalgPromotionOptions opt,
                             LinalgTransformationFilter filt)

 
 /// Create a LinalgStrategyGeneralizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyGeneralizePass(StringRef opName,
                                          LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
 }
 
+/// Create a LinalgStrategyInterchangePass.
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
+                                          LinalgTransformationFilter filter) {
+  return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
+                                                         filter);
+}
+
 /// Create a LinalgStrategyVectorizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyVectorizePass(StringRef opName,
                                         LinalgVectorizationOptions opt,
                                         LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter);
 }
 

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

 // Test that both anchor-op name and MatmulOp-based codegen strategy produce the same result.
 // 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" | FileCheck %s
 // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=2,4,8 tile-interchange=1,2,0 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" | FileCheck %s
 // 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" | FileCheck %s --check-prefix=OUTER
 // 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" | FileCheck %s
 // 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" | FileCheck %s --check-prefix=OUTER
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 generalize" | FileCheck %s --check-prefix=GENER
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 generalize iterator-interchange=0,2,1" | FileCheck %s --check-prefix=GENER
 
 
 // CHECK-LABEL: func @matmul(
 // OUTER-LABEL: func @matmul(
 // GENER-LABEL: func @matmul(
 func @matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) {
   linalg.matmul
    ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>)
    outs(%C: memref<1584x1584xf32>)
 
   //      CHECK: vector.matrix_multiply
   // CHECK-SAME: {lhs_columns = 8 : i32, lhs_rows = 2 : i32, rhs_columns = 4 : i32}
   // CHECK-SAME: (vector<16xf32>, vector<32xf32>) -> vector<8xf32>
 
-  // OUTER: vector.outerproduct {{.*}} : vector<2xf32>, vector<4xf32>
-  // GENER: linalg.generic
+  //      OUTER: vector.outerproduct {{.*}} : vector<2xf32>, vector<4xf32>
+
+  //      GENER: linalg.generic
+  // GENER-SAME: iterator_types = ["parallel", "reduction", "parallel"]
   return
 }
 

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

Option<bool> registerPromote{
       llvm::cl::init(false)};
   Option<bool> registerPromoteFullTile{
       *this, "register-promote-full-tile-pad",
       llvm::cl::desc("Pad the small aligned memory buffer to the tile sizes."),
       llvm::cl::init(false)};
   Option<bool> generalize{*this, "generalize",
                           llvm::cl::desc("Generalize named operations."),
                           llvm::cl::init(false)};
+  ListOption<int64_t> iteratorInterchange{
+      *this, "iterator-interchange", llvm::cl::MiscFlags::CommaSeparated,
+      llvm::cl::desc("Specifies the iterator interchange.")};
   Option<bool> vectorize{
       *this, "vectorize",
       llvm::cl::desc("Rewrite the linalg op as a vector operation."),
       llvm::cl::init(false)};
   Option<std::string> splitVectorTransfersTo{
       *this, "split-transfers",
       llvm::cl::desc(
           "Split vector transfers between slow (masked) and fast "
strategy.tileIf(!tileSizes.empty(), anchorOpName, tilingOptions)
                      .setAlignment(16)
                      .setUseFullTileBuffersByDefault(promoteFullTile))
       .tileIf(!registerTileSizes.empty(), anchorOpName, registerTilingOptions)
       .promoteIf(registerPromote, anchorOpName,
                  LinalgPromotionOptions()
                      .setAlignment(16)
                      .setUseFullTileBuffersByDefault(registerPromoteFullTile))
       .generalizeIf(generalize, anchorOpName)
+      .interchangeIf(!iteratorInterchange.empty(), iteratorInterchange)
       .vectorizeIf(vectorize, generalize ? genericOpName : anchorOpName)
       .setEnableVectorTransferPartialRewrite(true)
       .setEnableVectorContractLowering(true)
       .setEnableVectorToSCFConversion(true)
       .setVectorTransformsOptions(
           vector::VectorTransformsOptions()
               .setVectorTransformsOptions(vectorContractLowering)
               .setVectorTransferSplit(vectorTransferSplit))