[mlir][linalg] Add decompose to CodegenStrategy.

[gysit]

Summary

Add the decompose patterns that lower higher dimensional convolutions to lower dimensional ones to CodegenStrategy and use CodegenStrategy to test the decompose patterns. Additionally, remove the assertion that checks the anchor op name is set in the CodegenStrategyTest pass. Removing the assertion allows us to simplify the pipelines used in the interchange and decompose tests.

Depends On D114797

[gysit]

Move decompose in front of generalization.

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

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(f) {}
+
+  void addToPassPipeline(OpPassManager &pm,
+                         LinalgTransformationFilter m) const override {
+    pm.addPass(createLinalgStrategyDecomposePass(m));
+  }
+};
+
 /// Represent one application of createLinalgStrategyVectorizePass.
 struct Vectorize : public Transformation {
   explicit Vectorize(linalg::LinalgVectorizationOptions options,
                      LinalgTransformationFilter::FilterFunction f = nullptr,
                      bool padVectorize = false)
       : Transformation(f), opName(), options(options),
         vectorizePadding(padVectorize) {}
 
struct CodegenStrategy {
   }
   /// 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 patterns to decompose convolutions.
+  CodegenStrategy &
+  decompose(LinalgTransformationFilter::FilterFunction f = nullptr) {
+    transformationSequence.emplace_back(std::make_unique<Decompose>(f));
+    return *this;
+  }
+  /// Conditionally append patterns to decompose convolutions.
+  CodegenStrategy &
+  decomposeIf(bool b, LinalgTransformationFilter::FilterFunction f = nullptr) {
+    return b ? decompose(f) : *this;
+    return *this;
+  }
   /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
   CodegenStrategy &
   vectorize(StringRef opName,
             LinalgTransformationFilter::FilterFunction f = nullptr,
             bool vectorizePadding = false) {
     transformationSequence.emplace_back(std::make_unique<Vectorize>(
         opName, linalg::LinalgVectorizationOptions(), f, vectorizePadding));
     return *this;

mlir/test/Dialect/Linalg/decompose-convolution.mlir

-// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-decompose-convolution-patterns %s | FileCheck %s
+// RUN: mlir-opt -test-linalg-codegen-strategy="decompose" -split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: func @conv2d_nhwc_4x1x2x8_tensor
 //  CHECK-SAME: (%[[INPUT:.+]]: tensor<4x1x6x3xf32>, %[[FILTER:.+]]: tensor<1x2x3x8xf32>, %[[INIT:.+]]: tensor<4x1x2x8xf32>)
 func @conv2d_nhwc_4x1x2x8_tensor(%input: tensor<4x1x6x3xf32>, %filter: tensor<1x2x3x8xf32>, %init: tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32> {
   %0 = 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(%init : tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32>

mlir/test/Dialect/Linalg/interchange.mlir

-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.generic iterator-interchange=4,0,3,1,2" | FileCheck %s
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.generic iterator-interchange=4,0,3,1,2" -test-linalg-codegen-strategy="anchor-op=linalg.generic iterator-interchange=1,3,4,2,0" | FileCheck --check-prefix=CANCEL-OUT %s
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="iterator-interchange=4,0,3,1,2" | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="iterator-interchange=4,0,3,1,2" -test-linalg-codegen-strategy="iterator-interchange=1,3,4,2,0" | FileCheck --check-prefix=CANCEL-OUT %s
 
 #map0 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
 #map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3)>
 
 func @interchange_generic_op(%arg0 : memref<1x2x3x4x5xindex>, %arg1 : memref<1x2x4xindex>) {
   linalg.generic {
     indexing_maps = [#map0, #map1],
     iterator_types = ["parallel", "parallel", "reduction", "parallel", "reduction"]}

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

ListOption<int64_t> hoistPaddings{
       llvm::cl::desc("Operand hoisting depths when test-pad-pattern."),
       llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
   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> decompose{
+      *this, "decompose",
+      llvm::cl::desc("Decompose convolutions to lower dimensional ones."),
+      llvm::cl::init(false)};
   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 "

 
 void TestLinalgCodegenStrategy::runStrategy(
     LinalgTilingAndFusionOptions tilingAndFusionOptions,
     LinalgTilingOptions tilingOptions,
     LinalgTilingOptions registerTilingOptions,
     LinalgPaddingOptions paddingOptions,
     vector::VectorContractLowering vectorContractLowering,
     vector::VectorTransferSplit vectorTransferSplit) {
-  assert(!anchorOpName.empty());
   CodegenStrategy strategy;
   strategy
       .tileAndFuseIf(fuse && !tileSizes.empty(), anchorOpName,
                      tilingAndFusionOptions)
       .tileIf(!fuse && !tileSizes.empty(), anchorOpName, tilingOptions)
       .promoteIf(!fuse && promote, anchorOpName,
                  LinalgPromotionOptions()
                      .setAlignment(16)
                      .setUseFullTileBuffersByDefault(promoteFullTile))
       .tileIf(!fuse && !registerTileSizes.empty(), anchorOpName,
               registerTilingOptions)
       .promoteIf(!fuse && registerPromote, anchorOpName,
                  LinalgPromotionOptions()
                      .setAlignment(16)
                      .setUseFullTileBuffersByDefault(registerPromoteFullTile))
       .padIf(pad, "", paddingOptions)
+      .decomposeIf(decompose)
       .generalizeIf(generalize, "")
       .interchangeIf(!iteratorInterchange.empty(), iteratorInterchange)
       .vectorizeIf(vectorize, "")
       .vectorLowering(
           LinalgVectorLoweringOptions()
               .setVectorTransformsOptions(
                   vector::VectorTransformsOptions()
                       .setVectorTransformsOptions(vectorContractLowering)

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

Option<bool> skipPartial{
       *this, "skip-partial",
       llvm::cl::desc("Skip loops inside partial iterations during peeling"),
       llvm::cl::init(false)};
   Option<std::string> loopType{
       *this, "loop-type",
       llvm::cl::desc("Specify the type of loops to generate: for, parallel or "
                      "tiled_loop"),
       llvm::cl::init("for")};
-  Option<bool> testDecomposeConvolutionPattern{
-      *this, "test-decompose-convolution-patterns",
-      llvm::cl::desc("Test a set of patterns to rewrite high-D convolution ops "
-                     "into low-D ones"),
-      llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//
if (testTiledLoopPeeling.hasValue())
     return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling,
                                         skipPartial);
   if (testTilePattern)
     return applyTilePattern(getFunction(), loopType, tileSizes, peeledLoops,
                             /*scalarizeDynamicDims=*/false);
   if (testTileScalarizeDynamicDims)
     return applyTilePattern(getFunction(), loopType, tileSizes,
                             /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
-  if (testDecomposeConvolutionPattern) {
-    // TODO: thread all tests through LinalgStrategy passes.
-    OpPassManager dynamicPM("builtin.func");
-    dynamicPM.addPass(createLinalgStrategyDecomposePass());
-    if (failed(runPipeline(dynamicPM, getFunction())))
-      return signalPassFailure();
-  }
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir