aboutsummaryrefslogtreecommitdiff
path: root/mlir
diff options
context:
space:
mode:
authorJianbang Yang <yangjianbang112@gmail.com>2024-03-25 21:08:18 +0800
committerGitHub <noreply@github.com>2024-03-25 21:08:18 +0800
commit4bb9f918ff4e2a932ce1e3da5f14103778b04136 (patch)
tree2167376d1bb4916fa7f28243227911824700e916 /mlir
parent46a737cef2a709d5c121ab616fc80088ce2eed46 (diff)
downloadllvm-4bb9f918ff4e2a932ce1e3da5f14103778b04136.zip
llvm-4bb9f918ff4e2a932ce1e3da5f14103778b04136.tar.gz
llvm-4bb9f918ff4e2a932ce1e3da5f14103778b04136.tar.bz2
[mlir][tensor] fix out-of-bound index in tensor.dim (#85901)
fix a crash when fold tensor.dim with out-of-bound index. Fixes: https://github.com/llvm/llvm-project/issues/70183
Diffstat (limited to 'mlir')
-rw-r--r--mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp3
-rw-r--r--mlir/test/Dialect/Tensor/canonicalize.mlir23
2 files changed, 26 insertions, 0 deletions
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 2b226c7..a656c81 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -333,6 +333,9 @@ struct FoldDimOfAllocTensorOp : public OpRewritePattern<tensor::DimOp> {
auto allocTensorOp = dimOp.getSource().getDefiningOp<AllocTensorOp>();
if (!allocTensorOp || !maybeConstantIndex)
return failure();
+ if (*maybeConstantIndex < 0 ||
+ *maybeConstantIndex >= allocTensorOp.getType().getRank())
+ return failure();
if (!allocTensorOp.getType().isDynamicDim(*maybeConstantIndex))
return failure();
rewriter.replaceOp(
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index e5374f0..9ab54fe 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2367,3 +2367,26 @@ func.func @dim_of_reshape_undominated(%arg0: tensor<*xf32>, %arg1: tensor<?xinde
%dim = tensor.dim %reshape, %0 : tensor<*xf32>
return %dim : index
}
+
+// -----
+
+// Test case: This test fails to fold because the index of tensor.dim is out_of_bounds
+// CHECK-LABEL: func @dim_out_of_bounds(
+// CHECK: %[[IDX:.*]] = index.constant 28
+// CHECK-NEXT: bufferization.alloc_tensor
+// CHECK-NEXT: %[[DIM:.*]] = tensor.dim %{{.*}}, %[[IDX]]
+// CHECK-NEXT: memref.alloc
+// CHECK-NEXT: memref.cast
+// CHECK-NEXT: affine.vector_load %{{.*}}[{{.*}}, {{.*}}, symbol(%[[DIM]])]
+// CHECK-NEXT: return
+func.func @dim_out_of_bounds() -> vector<7xi32> {
+ %c1 = arith.constant 1 : index
+ %idx28 = index.constant 28
+ %c29 = arith.constant 29 : index
+ %3 = bufferization.alloc_tensor(%c29) : tensor<?xi16>
+ %dim = tensor.dim %3, %idx28 : tensor<?xi16>
+ %alloc_21 = memref.alloc(%c29) : memref<?x26x2xi32>
+ %16 = affine.vector_load %alloc_21[%c1, %c1, %dim] : memref<?x26x2xi32>, vector<7xi32>
+ return %16 : vector<7xi32>
+}
+