[mlir][Linalg] Make a LinalgStrategyDecomposePass available.

[nicolasvasilache]

Rebase.

[nicolasvasilache]

Rebase properly.

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

 /// on SSA use-def chains starting from function operands that are annotated
 /// with the 'inplaceable' attribute.
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass();
 
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgBufferizePass();
 
-/// Create a pass to conver named Linalg operations to Linalg generic
+/// Create a pass to convert named Linalg operations to Linalg generic
 /// operations.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgGeneralizationPass();
 
 /// Create a pass to convert Linalg operations to equivalent operations that
 /// work on primitive types, if possible.
 std::unique_ptr<Pass> createLinalgDetensorizePass();
 
 /// Create a pass to tile a LinalgOp and fuse its producers.
linalg::LinalgTransformationFilter filter =
         linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyGeneralizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgStrategyGeneralizePass(StringRef opName = "",
                                    linalg::LinalgTransformationFilter filter =
                                        linalg::LinalgTransformationFilter());
 
+/// Create a LinalgStrategyDecomposePass.
+// TODO: atm this is applied to all supported ops. If/when we need finer control
+// this should be exposed with an opName + filter and a proper pattern.
+std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategyDecomposePass();
+
 /// 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>>

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.">,
   ];
 }
 
+// TODO: atm this is applied to all supported ops. If/when we need finer control
+// this should be exposed with an opName + filter and a proper pattern.
+def LinalgStrategyDecomposePass
+    : FunctionPass<"linalg-strategy-decompose-pass"> {
+  let summary = "Configurable pass to apply pattern-based generalization.";
+  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
     : 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.">,

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

void runOnFunction() override {
     if (failed(applyPatternsAndFoldGreedily(funcOp,
                                             std::move(generalizationPattern))))
       signalPassFailure();
   }
 
   LinalgTransformationFilter filter;
 };
 
+/// Configurable pass to apply lowering of coarser-grained named linalg ops into
+/// finer-grained named versions.
+struct LinalgStrategyDecomposePass
+    : public LinalgStrategyDecomposePassBase<LinalgStrategyDecomposePass> {
+
+  LinalgStrategyDecomposePass() = default;
+
+  void runOnFunction() override {
+    auto funcOp = getFunction();
+    if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
+      return;
+    RewritePatternSet decompositionPattern(funcOp.getContext());
+    populateDecomposeConvolutionPatterns(decompositionPattern);
+    if (failed(applyPatternsAndFoldGreedily(funcOp,
+                                            std::move(decompositionPattern))))
+      signalPassFailure();
+  }
+};
+
 /// Configurable pass to apply pattern-based linalg generalization.
 struct LinalgStrategyInterchangePass
     : public LinalgStrategyInterchangePassBase<LinalgStrategyInterchangePass> {
 
   LinalgStrategyInterchangePass() = default;
 
   LinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
                                 LinalgTransformationFilter filter)

 }
 
 /// Create a LinalgStrategyGeneralizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyGeneralizePass(StringRef opName,
                                          LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
 }
+/// Create a LinalgStrategyDecomposePass.
+// TODO: atm this is applied to all supported ops. If/when we need finer control
+// this should be exposed with an opName + filter and a proper pattern.
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgStrategyDecomposePass() {
+  return std::make_unique<LinalgStrategyDecomposePass>();
+}
 
 /// Create a LinalgStrategyInterchangePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
                                           LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
                                                          filter);
 }

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

 //===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing Linalg transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
-#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
patterns.add<LinalgVectorizationPattern>(
       funcOp.getContext(),
       LinalgTransformationFilter()
           .addOpFilter<ContractionOpInterface, FillOp, CopyOp, GenericOp>());
   populatePadTensorOpVectorizationPatterns(patterns);
   populateConvolutionVectorizationPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
-static void applyDecomposeConvolutionPatterns(FuncOp funcOp) {
-  RewritePatternSet patterns(funcOp.getContext());
-  populateDecomposeConvolutionPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
-}
-
 static void applyPadTensorToGenericPatterns(FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   patterns.add<PadTensorOpTransformationPattern>(funcOp.getContext());
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
 static void applyGeneralizePadTensorPatterns(FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
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)
-    return applyDecomposeConvolutionPatterns(getFunction());
+  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