diff options
author | Jianbang Yang <yangjianbang112@gmail.com> | 2024-03-25 21:08:18 +0800 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-03-25 21:08:18 +0800 |
commit | 4bb9f918ff4e2a932ce1e3da5f14103778b04136 (patch) | |
tree | 2167376d1bb4916fa7f28243227911824700e916 /mlir | |
parent | 46a737cef2a709d5c121ab616fc80088ce2eed46 (diff) | |
download | llvm-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.cpp | 3 | ||||
-rw-r--r-- | mlir/test/Dialect/Tensor/canonicalize.mlir | 23 |
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> +} + |