[Mlir-commits] [mlir] [mlir][tensor] Verify that `tensor.pad` padding is non-negative (PR #231389)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tensor

Author: Federico Bruzzone (FedericoBruzzone)

<details>
<summary>Changes</summary>

`tensor.pad` did not reject negative static `low` or `high` padding, so an op like `tensor.pad %t low[-2] high[0]` was verified as long as the result type was consistent with it.

With this PR, we reject negative static padding in `PadOp::verify`. Dynamic padding values are, of course, not checked.

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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td (+2) 
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+5) 
- (modified) mlir/test/Dialect/Tensor/invalid.mlir (+21) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index beac1e91ee6e3..7744f2f4dc44e 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1275,6 +1275,8 @@ def Tensor_PadOp : Tensor_Op<"pad", [
     The result tensor dimensions are `low[i]` + `dim[i]` + `high[i]` for each
     dimension `i`. The number of elements of `low` and `high` must match the
     rank of the input tensor. They can be either a constant or a dynamic value.
+    The padding amounts must be non-negative. The verifier checks this for
+    constant values only.
 
     The region of the `tensor.pad` operation returns the value to use
     for the padding. The arguments of the region represent the index
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 53ee57fa72a17..4b99c9d3ce7ce 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3252,6 +3252,11 @@ void PadOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
 }
 
 LogicalResult PadOp::verify() {
+  for (ArrayRef<int64_t> padding : {getStaticLow(), getStaticHigh()})
+    for (int64_t size : padding)
+      if (!ShapedType::isDynamic(size) && size < 0)
+        return emitOpError("padding must be non-negative, got ") << size;
+
   auto sourceType = llvm::cast<RankedTensorType>(getSource().getType());
   auto resultType = llvm::cast<RankedTensorType>(getResult().getType());
   auto expectedType =
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index 19db4680b6a63..d13290f141e64 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -731,3 +731,24 @@ func.func @no_fold_invalid_collapse() -> tensor<i64> {
     return %0 : tensor<i64>
 }
 
+// -----
+
+func.func @pad_negative_low(%arg0: tensor<10xf32>, %arg1: f32) -> tensor<8xf32> {
+  // expected-error @+1 {{padding must be non-negative, got -2}}
+  %0 = tensor.pad %arg0 low[-2] high[0] {
+  ^bb0(%arg2: index):
+    tensor.yield %arg1 : f32
+  } : tensor<10xf32> to tensor<8xf32>
+  return %0 : tensor<8xf32>
+}
+
+// -----
+
+func.func @pad_negative_high(%arg0: tensor<?x4xf32>, %arg1: f32) -> tensor<?x3xf32> {
+  // expected-error @+1 {{padding must be non-negative, got -1}}
+  %0 = tensor.pad %arg0 low[0, 0] high[0, -1] {
+  ^bb0(%arg2: index, %arg3: index):
+    tensor.yield %arg1 : f32
+  } : tensor<?x4xf32> to tensor<?x3xf32>
+  return %0 : tensor<?x3xf32>
+}

``````````

</details>


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