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

[Federico Bruzzone]

https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/223712

>From 0e677828d5857daf830112984e08cdb4c4917d39 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 15 Sep 2026 16:11:04 +0200
Subject: [PATCH 1/2] [mlir][linalg] Handle `DropPadUnitDims` when padding
 value is defined inside the pad body

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../Linalg/Transforms/DropUnitDims.cpp        |  4 ++++
 .../Dialect/Linalg/drop-unit-extent-dims.mlir | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index dac2891ce65e56..1b99a63ebc7bba 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -646,6 +646,10 @@ struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {
           padOp, "unimplemented: non-constant padding value");
     }
 
+    // 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);
+
     ArrayRef<int64_t> sourceShape = padOp.getSourceType().getShape();
     ArrayRef<int64_t> resultShape = padOp.getResultType().getShape();
     int64_t padRank = sourceShape.size();
diff --git a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
index 841d0e5f565125..ed804587c88fbc 100644
--- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
+++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
@@ -1032,6 +1032,28 @@ 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>
+
+// -----
+
 func.func @do_not_drop_non_constant_padding(%arg0: tensor<1x1x3x1x1xf32>, %pad: f32) -> tensor<1x2x3x1x3xf32>
 {
   %c0 = arith.constant 0 : index

>From 5f1c9d1331c4107816ef5ed26a5d76cf5dd2479b Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Sat, 24 Oct 2026 10:27:54 +0200
Subject: [PATCH 2/2] Address comments

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../Linalg/Transforms/DropUnitDims.cpp        |  8 +++----
 .../Dialect/Linalg/drop-unit-extent-dims.mlir | 21 +++++++++++++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index 1b99a63ebc7bba..d7cedd7369b296 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -646,10 +646,6 @@ struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {
           padOp, "unimplemented: non-constant padding value");
     }
 
-    // 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);
-
     ArrayRef<int64_t> sourceShape = padOp.getSourceType().getShape();
     ArrayRef<int64_t> resultShape = padOp.getResultType().getShape();
     int64_t padRank = sourceShape.size();
@@ -706,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 ed804587c88fbc..0c289823557fd5 100644
--- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
+++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
@@ -1054,6 +1054,27 @@ func.func @drop_unit_pad_dims_constant_inside_body(%arg0: tensor<1x28x28x1xf32>)
 
 // -----
 
+// 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