[mlir][linalg] Expose SwapSliceOfPad as strategy pass

[nicolasvasilache]

Should this be run as part of the enabling passes rather than a single pattern pass that we have to not forget to call?
Putting a blocker until we have a consensus.

[antiagainst]

This is not needed anymore.

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

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 pass to apply patterns to convert extract_slice(pad_tensor) into
+/// pad_tensor(extract_slice).
+std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategySwapSliceOfPadPass();
+
 /// 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 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 LinalgStrategySwapSliceOfPadPass
+    : FunctionPass<"linalg-strategy-swap-slice-of-pad-pass"> {
+  let summary = "Configurable pass to convert extract_slice(pad_tensor) into "
+                "pad_tensor(extract_slice).";
+  let constructor = "mlir::createLinalgStrategySwapSliceOfPadPass()";
+  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 {
     RewritePatternSet decompositionPattern(funcOp.getContext());
     populateDecomposeConvolutionPatterns(decompositionPattern);
     if (failed(applyPatternsAndFoldGreedily(funcOp,
                                             std::move(decompositionPattern))))
       signalPassFailure();
   }
 };
 
+/// Configurable pass to apply patterns to convert extract_slice(pad_tensor)
+/// into pad_tensor(extract_slice).
+struct LinalgStrategySwapSliceOfPadPass
+    : public LinalgStrategySwapSliceOfPadPassBase<
+          LinalgStrategySwapSliceOfPadPass> {
+  void runOnFunction() override {
+    auto funcOp = getFunction();
+    if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
+      return;
+    MLIRContext *context = funcOp.getContext();
+    RewritePatternSet patterns(context);
+    patterns.insert<linalg::ExtractSliceOfPadTensorSwapPattern>(context);
+    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
+      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 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>();
 }
 
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgStrategySwapSliceOfPadPass() {
+  return std::make_unique<LinalgStrategySwapSliceOfPadPass>();
+}
+
 /// Create a LinalgStrategyInterchangePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
                                           LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
                                                          filter);
 }