[mlir][tensor] Fix createFillOrGenerateOp

[CoTinker]

This PR clones the padding value defined inside the PadOp block to outside to prevent a crash. Fixes #120947.

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-tensor

@llvm/pr-subscribers-mlir-linalg

Author: Longsheng Mou (CoTinker)

Changes

In this context, 'constant' refers to a value defined outside the PadOp block. Therefore, we should perform the 'inside block check' before the 'constant check' to avoid a crash. Fixes #120947.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+3-3)
• (modified) mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir (+25)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index f79c774ceb3e9a..be846a11dcef18 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3638,12 +3638,12 @@ Value PadOp::getConstantPaddingValue() {
   if (!yieldOp)
     return {};
   Value padValue = yieldOp.getValue();
-  // Check if yield value is a constant.
-  if (matchPattern(padValue, m_Constant()))
-    return padValue;
   // Check if yield value is defined inside the PadOp block.
   if (padValue.getParentBlock() == &getRegion().front())
     return {};
+  // Check if yield value is a constant.
+  if (matchPattern(padValue, m_Constant()))
+    return padValue;
   // Else: Yield value defined outside of the PadOp block.
   return padValue;
 }
diff --git a/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir b/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir
index a0a676edceb745..c4a5916aaf9a61 100644
--- a/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir
+++ b/mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir
@@ -19,6 +19,8 @@ func.func @generalize_pad_tensor_static_shape(%arg0: tensor<1x28x28x1xf32>) -> t
   return %0 : tensor<1x32x32x1xf32>
 }
 
+// -----
+
 // CHECK-LABEL:   func @generalize_pad_tensor_dynamic_shape(
 // CHECK-SAME:                                              %[[IN:.*]]: tensor<4x?x2x?xf32>,
 // CHECK-SAME:                                              %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> {
@@ -44,3 +46,26 @@ 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>
 }
+
+// -----
+
+// Ensure that the constant value inside the PadOp block does not cause a crash.
+
+// CHECK-LABEL:   func.func @generalize_pad_tensor_constant_inside(
+// CHECK-SAME:                                                     %[[VAL_0:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
+// CHECK:           %[[VAL_1:.*]] = tensor.generate  {
+// CHECK:           ^bb0(%[[VAL_2:.*]]: index, %[[VAL_3:.*]]: index, %[[VAL_4:.*]]: index, %[[VAL_5:.*]]: index):
+// CHECK:             %[[VAL_6:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:             tensor.yield %[[VAL_6]] : f32
+// CHECK:           } : tensor<1x32x32x1xf32>
+// CHECK:           %[[VAL_7:.*]] = tensor.insert_slice %[[VAL_0]] into %[[VAL_8:.*]][0, 2, 2, 0] [1, 28, 28, 1] [1, 1, 1, 1] : tensor<1x28x28x1xf32> into tensor<1x32x32x1xf32>
+// CHECK:           return %[[VAL_7]] : 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>
+}

[banach-space]

Makes sense, thanks. I've left a couple of minor comments.

Side comment - the test that's been added could quite easily be lowered to linalg.fill. Unless I am missing sth? If that's the case, lets leave a TODO somewhere. Probably near the test.

[CoTinker]

Makes sense, thanks. I've left a couple of minor comments.

Side comment - the test that's been added could quite easily be lowered to linalg.fill. Unless I am missing sth? If that's the case, lets leave a TODO somewhere. Probably near the test.

Actually, the constant value inside the PadOp can canonicalize to the outside, such as:

// mlir-opt --canonicalize
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>
}

--->

func.func @generalize_pad_tensor_constant_inside(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
    %cst = arith.constant 0.000000e+00 : f32
    %padded = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst : f32
    } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>
    return %padded : tensor<1x32x32x1xf32>
  }

[banach-space]

Actually, the constant value inside the PadOp can canonicalize to the outside

That's what I expected. It makes me think - wouldn't it be better to just bail out in this case (as opposed to creating tensor.generate)?

[hanhanW]

Actually, the constant value inside the PadOp can canonicalize to the outside

That's what I expected. It makes me think - wouldn't it be better to just bail out in this case (as opposed to creating tensor.generate)?

I'm not sure how the conversion works. However, shouldn't the generated tensor.generate op be replaced with a constant? It can happen either in the conversion or the tensor::GenerateOp's folder.

// CHECK:           %[[VAL_1:.*]] = tensor.generate  {
// CHECK:           ^bb0(%[[VAL_2:.*]]: index, %[[VAL_3:.*]]: index, %[[VAL_4:.*]]: index, %[[VAL_5:.*]]: index):
// CHECK:             %[[VAL_6:.*]] = arith.constant 0.000000e+00 : f32
// CHECK:             tensor.yield %[[VAL_6]] : f32
// CHECK:           } : tensor<1x32x32x1xf32>

[CoTinker]

I'm sorry for the delay. Could you review it again? Thanks. @banach-space @hanhanW

[CoTinker]

Ping~

[nicolasvasilache]

thanks!

[banach-space]

Thanks!

[CoTinker]

Thanks for your review.