[Mlir-commits] [mlir] [mlir][linalg] Fix tensor.pad drop unit dims crash for region-local padding constants (PR #192424)

[Hocky Yudhiono]

https://github.com/hockyy created https://github.com/llvm/llvm-project/pull/192424

Fix a crash in `linalg-fold-unit-extent-dims` when rewriting `tensor.pad` ops whose padding value is defined inside the pad region.

>From 38889131c80cc2033034a620ed6c2e543e53dbc4 Mon Sep 17 00:00:00 2001
From: Hocky Yudhiono <hocky.yudhiono at gmail.com>
Date: Thu, 16 Apr 2026 18:25:00 +0800
Subject: [PATCH] [mlir][linalg] Fix tensor.pad drop unit dims crash when pad
 value inside body

---
 .../Dialect/Linalg/Transforms/DropUnitDims.cpp   |  5 +++++
 .../Dialect/Linalg/drop-unit-extent-dims.mlir    | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index c3dca148b7f94..933c7fa267ac8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -644,6 +644,11 @@ struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {
           padOp, "unimplemented: non-constant padding value");
     }
 
+    // `getConstantPaddingValue` may return a constant defined inside the
+    // `tensor.pad` region.
+    if (paddingVal.getParentBlock() == &padOp.getRegion().front())
+      rewriter.moveOpBefore(paddingVal.getDefiningOp(), padOp);
+
     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 841d0e5f56512..a082d590fc36c 100644
--- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
+++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
@@ -993,6 +993,22 @@ func.func @drop_unit_pad_dims(%arg0: tensor<1x1x3x1x1xf32>) -> tensor<1x2x3x1x3x
 
 // -----
 
+func.func @drop_unit_pad_dims_constant_inside(%arg0: tensor<1x1xf32>) -> tensor<1x2xf32> {
+  %0 = tensor.pad %arg0 low[0, 0] high[0, 1] {
+    ^bb0(%arg1: index, %arg2: index):
+      %cst2 = arith.constant 2.0 : f32
+      tensor.yield %cst2 : f32
+  } : tensor<1x1xf32> to tensor<1x2xf32>
+  return %0 : tensor<1x2xf32>
+}
+
+// CHECK-LABEL: func @drop_unit_pad_dims_constant_inside
+//       CHECK:   %[[CST:.*]] = arith.constant 2.000000e+00 : f32
+//       CHECK:   tensor.pad
+//       CHECK:     tensor.yield %[[CST]] : f32
+
+// -----
+
 func.func @drop_unit_pad_dynamic_dims(%arg0: tensor<1x?xf32>) -> tensor<1x?xf32>
 {
   %c0 = arith.constant 0 : index