From f1dd6b3cf8d957f28e06f475c8a652f201fc830b Mon Sep 17 00:00:00 2001 From: Longsheng Mou Date: Tue, 15 Apr 2025 09:29:44 +0800 Subject: [PATCH] [mlir][tensor] Fix `createFillOrGenerateOp` (#121205) This PR clones the padding value defined inside the PadOp block to outside to prevent a crash. Fixes #120947. --- .../Dialect/Linalg/Transforms/Transforms.cpp | 6 +++++- .../TensorToLinalg/tensor-ops-to-linalg.mlir | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp index dcd50cc44f81..4c2a7c36d8b5 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -927,8 +927,12 @@ Value DecomposePadOpPattern::createFillOrGenerateOp( RewriterBase &rewriter, tensor::PadOp padOp, Value dest, const SmallVector &dynSizes) const { auto padValue = padOp.getConstantPaddingValue(); - if (padValue) + if (padValue) { + // Move the padding value defined inside the PadOp block to outside. + if (padValue.getParentBlock() == &padOp.getRegion().front()) + rewriter.moveOpBefore(padValue.getDefiningOp(), padOp); return rewriter.create(padOp.getLoc(), padValue, dest).result(); + } // Fill could not be optimized: Lower to tensor::GenerateOp with region. auto generateOp = rewriter.create( diff --git a/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir b/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir index a0a676edceb7..b58d407733f2 100644 --- a/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir +++ b/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir @@ -44,3 +44,22 @@ func.func @generalize_pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1 } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32> return %out : tensor<4x?x?x?xf32> } + +// ----- + +// CHECK-LABEL: func.func @generalize_pad_tensor_constant_inside( +// CHECK-SAME: %[[SRC:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> { +// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<1x32x32x1xf32> +// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32 +// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<1x32x32x1xf32>) -> tensor<1x32x32x1xf32> +// CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[SRC]] into %[[FILL]][0, 2, 2, 0] [1, 28, 28, 1] [1, 1, 1, 1] : tensor<1x28x28x1xf32> into tensor<1x32x32x1xf32> +// CHECK: return %[[PADDED]] : tensor<1x32x32x1xf32> +// CHECK: } +func.func @generalize_pad_tensor_constant_inside(%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.000000e+00 : f32 + tensor.yield %cst : f32 + } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32> + return %0 : tensor<1x32x32x1xf32> +}