[mlir][linalg] Do not set insertion point inside padding function (#165420)

Remove insertion point in rewriteAsPaddedOp. There is no gurantee that
the sizes provided by the user are before the operation to pad. It's
better to let the user handle where to insert the newly created
operations, as long as they are after the origin operation to pad.

3 files changed, 5 insertions, 4 deletions

diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index c89fc59..d00183a 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -653,6 +653,9 @@ struct PadTilingInterfaceResult {
 //    interpreted as the bounding box (dynamic) value to pad to.
 /// * Use "options.paddingValues" to set the padding value of the created
 //    tensor::PadOp.
+//
+// The transformation assumes that the insertion point is set after the
+// operation to pad.
 FailureOr<PadTilingInterfaceResult>
 rewriteAsPaddedOp(OpBuilder &, TilingInterface toPad,
                   PadTilingInterfaceOptions options,
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 794dda9..8b89244 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2464,6 +2464,8 @@ transform::PadTilingInterfaceOp::apply(transform::TransformRewriter &rewriter,
         .setPaddingSizes(getMixedPaddingSizes())
         .setPadToMultipleOf(getPadToMultipleOf());
 
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointAfter(targetOp);
     auto maybePadOps = rewriteAsPaddedOp(
         rewriter, cast<TilingInterface>(targetOp.getOperation()), options);
     if (failed(maybePadOps)) {
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index 3e787a2..52ab92f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -288,10 +288,6 @@ FailureOr<PadTilingInterfaceResult> linalg::rewriteAsPaddedOp(
     return failure();
   }
 
-  OpBuilder::InsertionGuard g(builder);
-  // Set IP after toPad because we also take the dims of toPad's output.
-  builder.setInsertionPointAfter(toPad);
-
   // 1. Get the loopUpperBounds from the TilingInterface.
   SmallVector<Range> iterationDomain = toPad.getIterationDomain(builder);