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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hocky Yudhiono (hockyy)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/192424.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp (+5) 
- (modified) mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir (+16) 


``````````diff
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

``````````

</details>


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