Skip to content

Commit ca8cf90

Browse files
authored
[mlir][tensor] Check the EmptyOp's dynamicSize to be non-negative (#65577)
This patch addresses a crash that occurs when negative dynamic sizes are provided in tensor.emptyOp by adding a check to ensure that dynamic sizes are non-negative. Fixes #64064
1 parent 2873a9a commit ca8cf90

File tree

2 files changed

+16
-0
lines changed

2 files changed

+16
-0
lines changed

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -712,6 +712,9 @@ struct ReplaceEmptyTensorStaticShapeDims : OpRewritePattern<EmptyOp> {
712712
Value dynamicSize = op.getDynamicSizes()[ctr++];
713713
std::optional<int64_t> cst = getConstantIntValue(dynamicSize);
714714
if (cst.has_value()) {
715+
// dynamic size must be non-negative.
716+
if (cst.value() < 0)
717+
return failure();
715718
staticShape[i] = *cst;
716719
changedType = true;
717720
} else {

mlir/test/Dialect/Tensor/canonicalize.mlir

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1848,3 +1848,16 @@ func.func @pack_unpack_dynamic_with_padding(%t: tensor<?x?x?x?xf32>, %dim1: inde
18481848
%packed = tensor.pack %unpacked padding_value(%pad: f32) inner_dims_pos = [0, 1] inner_tiles = [%tile1, %tile2] into %tensor_empty1 : tensor<?x?xf32> -> tensor<?x?x?x?xf32>
18491849
return %packed : tensor<?x?x?x?xf32>
18501850
}
1851+
1852+
// -----
1853+
1854+
// CHECK: func.func @invalid_empty_negative_size
1855+
// CHECK: %[[IDX:.*]] = index.constant
1856+
// CHECK: %[[T:.*]] = tensor.empty(%[[IDX]]) : tensor<4x5x?xf32>
1857+
func.func @invalid_empty_negative_size() -> (tensor<4x5x?xf32>) {
1858+
%c1 = arith.constant 1 : index
1859+
%cn2 = arith.constant 2 : index
1860+
%0 = index.sub %c1, %cn2
1861+
%1 = tensor.empty(%0) : tensor<4x5x?xf32>
1862+
return %1 : tensor<4x5x?xf32>
1863+
}

0 commit comments

Comments
 (0)