[Mlir-commits] [mlir] bfe0941 - [mlir][linalg] Handle `DropPadUnitDims` when padding value is defined inside the pad body (#223712)

[llvmlistbot at llvm.org]

Author: Federico Bruzzone
Date: 2026-10-24T17:40:41+02:00
New Revision: bfe094117ba828cd338ac415bf0f7e98f8b8d4c8

URL: https://github.com/llvm/llvm-project/commit/bfe094117ba828cd338ac415bf0f7e98f8b8d4c8
DIFF: https://github.com/llvm/llvm-project/commit/bfe094117ba828cd338ac415bf0f7e98f8b8d4c8.diff

LOG: [mlir][linalg] Handle `DropPadUnitDims` when padding value is defined inside the pad body (#223712)

`DropPadUnitDims` reused the pad's padding value as-is when building the
replacement `tensor.pad`.
If that value was defined inside the original pad's own body, the new op
ended up referencing it from outside its region, and erasing the
original pad then hit an assertion.

This PR clone the defining op when the padding value lives inside the
pad body.

Fixes #203813

---------

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
    mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index 75e6fc9491129e..40dc35ce6babbf 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -702,6 +702,10 @@ struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {
       return rewriter.notifyMatchFailure(padOp, "Failed to collapse source");
     }
 
+    // Clone the padding value if it is defined in the region of the pad op.
+    if (paddingVal.getParentBlock() == &padOp.getRegion().front())
+      paddingVal = rewriter.clone(*paddingVal.getDefiningOp())->getResult(0);
+
     auto newResultType = RankedTensorType::get(
         newResultShape, padOp.getResultType().getElementType());
     auto newPadOp = tensor::PadOp::create(

diff  --git a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
index 841d0e5f565125..0c289823557fd5 100644
--- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
+++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
@@ -1032,6 +1032,49 @@ func.func @drop_unit_pad_dynamic_dims(%arg0: tensor<1x?xf32>) -> tensor<1x?xf32>
 
 // -----
 
+func.func @drop_unit_pad_dims_constant_inside_body(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32>
+{
+  %0 = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] {
+    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
+      %cst = arith.constant 0.0 : f32
+      tensor.yield %cst : f32
+  } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>
+  return %0 : tensor<1x32x32x1xf32>
+}
+
+// CHECK-LABEL: func @drop_unit_pad_dims_constant_inside_body
+//       CHECK:   %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape
+//  CHECK-SAME:     {{\[}}[0, 1], [2, 3]{{\]}} : tensor<1x28x28x1xf32> into tensor<28x28xf32>
+//       CHECK:   %[[PADDED:.+]] = tensor.pad %[[COLLAPSE]] low[2, 2] high[2, 2]
+//       CHECK:     tensor.yield %[[CST]]
+//       CHECK:   } : tensor<28x28xf32> to tensor<32x32xf32>
+//       CHECK:   tensor.expand_shape %[[PADDED]]
+//  CHECK-SAME:     {{\[}}[0, 1], [2, 3]{{\]}} output_shape [1, 32, 32, 1] : tensor<32x32xf32> into tensor<1x32x32x1xf32>
+
+// -----
+
+// The pattern must not modify the IR when there are no unit dims to drop.
+
+func.func @negative_drop_pad_without_unit_dims_constant_inside_body(%arg0: tensor<2xf32>) -> tensor<3xf32>
+{
+  %0 = tensor.pad %arg0 low[0] high[1] {
+    ^bb0(%arg1: index):
+      %cst = arith.constant 0.0 : f32
+      tensor.yield %cst : f32
+  } : tensor<2xf32> to tensor<3xf32>
+  return %0 : tensor<3xf32>
+}
+
+// CHECK-LABEL: func @negative_drop_pad_without_unit_dims_constant_inside_body
+//   CHECK-NOT:   tensor.collapse_shape
+//       CHECK:   %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+//       CHECK:   tensor.pad %{{.*}} low[0] high[1]
+//       CHECK:     tensor.yield %[[CST]]
+//   CHECK-NOT:   tensor.expand_shape
+
+// -----
+
 func.func @do_not_drop_non_constant_padding(%arg0: tensor<1x1x3x1x1xf32>, %pad: f32) -> tensor<1x2x3x1x3xf32>
 {
   %c0 = arith.constant 0 : index