[mlir][Linalg] Retire CodegenStrategy::transform

[nicolasvasilache]

Summary

Instead each pass should constructed a nested OpPassManager and runPipeline on that.

[mehdi_amini]

Thanks!

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

175

Could this be automatically always added to the pass pipeline by the strategy? (a "cleanupTransformMarkersPass")

(also you introduced a cleanup() method but you're not using it, is this intended?)

[nicolasvasilache]

Add a tail cleanup pass.

[nicolasvasilache]

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

175

I was afraid it would be too aggressive to introduce a pass for this, I also played with RAII but it was not good.
Since you propose it yourself, cleanup pass it is then!

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

linalg::LinalgTransformationFilter filter =
         linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyLowerVectorsPass.
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgStrategyLowerVectorsPass(linalg::LinalgVectorLoweringOptions opt =
                                          linalg::LinalgVectorLoweringOptions(),
                                      linalg::LinalgTransformationFilter filter =
                                          linalg::LinalgTransformationFilter());
+
+/// Create a LinalgStrategyRemoveMarkersPass.
+std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategyRemoveMarkersPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Linalg/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_PASSES_H_

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

def LinalgStrategyLowerVectorsPass
   let constructor = "mlir::createLinalgStrategyLowerVectorsPass()";
   let dependentDialects = ["linalg::LinalgDialect"];
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,
   ];
 }
 
+def LinalgStrategyRemoveMarkersPass
+    : FunctionPass<"linalg-strategy-remove-markers-pass"> {
+  let summary = "Cleanup pass that drops markers.";
+  let constructor = "mlir::createLinalgStrategyRemoveMarkersPass()";
+  let dependentDialects = ["linalg::LinalgDialect"];
+  let options = [
+    Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
+      "Which func op is the anchor to latch on.">,
+  ];
+}
+
 #endif // MLIR_DIALECT_LINALG_PASSES

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

struct CodegenStrategy {
   }
   CodegenStrategy &setEnableVectorToSCFConversion(bool val) {
     this->lateCodegenStrategyOptions.enableVectorToSCFConversion = val;
     return *this;
   }
 
   /// Apply the transformation patterns in sequence with cleanup
   /// transformations interleaved.
-  LogicalResult transform(FuncOp func) const;
   void configurePassPipeline(OpPassManager &pm, MLIRContext *context) const;
 
 private:
   LogicalResult postPatternTransforms(Operation *func) const;
 
   vector::VectorTransformsOptions vectorTransformOptions;
   VectorTransferToSCFOptions vectorToSCFOptions;
   SmallVector<std::unique_ptr<Transformation>, 4> transformationSequence;
   LateCodegenStrategyOptions lateCodegenStrategyOptions;
 };
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_CODEGENSTRATEGY_H_

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

vectorLoweringOptions.enableVectorTransferPartialRewrite =
       lateCodegenStrategyOptions.enableVectorTransferPartialRewrite;
   vectorLoweringOptions.enableVectorContractLowering =
       lateCodegenStrategyOptions.enableVectorContractLowering;
   vectorLoweringOptions.enableVectorToSCFConversion =
       lateCodegenStrategyOptions.enableVectorToSCFConversion;
   vectorLoweringOptions.vectorTransformOptions = vectorTransformOptions;
   vectorLoweringOptions.vectorTransferToSCFOptions = vectorToSCFOptions;
   pm.addPass(createLinalgStrategyLowerVectorsPass(vectorLoweringOptions));
-}
-
-LogicalResult mlir::linalg::CodegenStrategy::transform(FuncOp funcOp) const {
-  PassManager pm(funcOp.getContext(), funcOp.getOperationName());
-  configurePassPipeline(pm, funcOp.getContext());
-  LogicalResult res = pm.run(funcOp);
-  // Ensure we drop the marker in the end.
-  funcOp.walk([](LinalgOp op) {
-    op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
-  });
-  return res;
+  pm.addPass(createLinalgStrategyRemoveMarkersPass());
 }

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

if (options.enableVectorToSCFConversion) {
                                             options.vectorTransferToSCFOptions);
     }
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   }
 
   LinalgVectorLoweringOptions options;
   LinalgTransformationFilter filter;
 };
+
+/// Configurable pass to lower vector operations.
+struct LinalgStrategyRemoveMarkersPass
+    : public LinalgStrategyRemoveMarkersPassBase<
+          LinalgStrategyRemoveMarkersPass> {
+
+  void runOnFunction() override {
+    auto funcOp = getFunction();
+    if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
+      return;
+    funcOp.walk([](LinalgOp op) {
+      op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
+    });
+  }
+};
 } // namespace
 
 /// Create a LinalgStrategyTilePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
                                    LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyTilePass>(opName, opt, filter);
 }

 }
 
 /// Create a LinalgStrategyLowerVectorsPass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
                                            LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyLowerVectorsPass>(opt, filter);
 }
+
+/// Create a LinalgStrategyRemoveMarkersPass.
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgStrategyRemoveMarkersPass() {
+  return std::make_unique<LinalgStrategyRemoveMarkersPass>();
+}
 No newline at end of file

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

strategy.tileIf(!tileSizes.empty(), anchorOpName, tilingOptions)
       .setEnableVectorContractLowering(true)
       .setEnableVectorToSCFConversion(true)
       .setVectorTransformsOptions(
           vector::VectorTransformsOptions()
               .setVectorTransformsOptions(vectorContractLowering)
               .setVectorTransferSplit(vectorTransferSplit))
       .setVectorTransferToSCFOptions(
           VectorTransferToSCFOptions().setUnroll(unrollVectorTransfers));
-  (void)strategy.transform(getFunction());
+
+  // Created a nested OpPassManager and run.
+  FuncOp funcOp = getFunction();
+  OpPassManager dynamicPM("builtin.func");
+  strategy.configurePassPipeline(dynamicPM, funcOp.getContext());
+  if (failed(runPipeline(dynamicPM, funcOp)))
+    return signalPassFailure();
 }
 } // end anonymous namespace
 
 /// Apply transformations specified as patterns.
 void TestLinalgCodegenStrategy::runOnFunction() {
   if (!anchorFuncOpName.empty() && anchorFuncOpName != getFunction().getName())
     return;
 
   LinalgTilingOptions tilingOptions;
   if (!tileSizes.empty())
     tilingOptions = tilingOptions.setTileSizes(tileSizes);
   if (!tileInterchange.empty())