diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td --- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td @@ -260,8 +260,11 @@ let description = [{ The `tensor.extract` op reads a ranked tensor and returns one element as specified by the given indices. The result of the op is a value with the - same type as the elements of the tensor. The arity of indices must match - the rank of the accessed value. All indices should all be of `index` type. + same type as the elements of the tensor. + + The arity of indices must match the rank of the accessed value. All indices + should all be of `index` type. They must be non-negative and within the + bounds of the tensor. Example: @@ -301,9 +304,11 @@ * source: the "base" tensor from which to extract a slice. * offsets: tensor-rank number of offsets into the "base" tensor from which - to extract the slice. + to extract the slice. Offsets must be non-negative and within the + bounds of the source tensor. * sizes: tensor-rank number of sizes which specify the sizes of the result - tensor type. + tensor type. "offset + size" must not exceed the respective + source dimension size. * strides: tensor-rank number of strides specifying subsampling in each dimension. @@ -731,7 +736,8 @@ of `scalar`. The arity of `indices `must match the rank of the tensor `dest`. All - indices should be of `index` type. + indices should be of `index` type. They must be non-negative and within the + bounds of the tensor. Example: @@ -787,9 +793,11 @@ * source: the tensor that is inserted. * dest: the tensor into which the source tensor is inserted. * offsets: tensor-rank number of offsets into the `dest` tensor into which - the slice is inserted. + the slice is inserted. Offsets must be non-negative and within + the bounds of the destination tensor. * sizes: tensor-rank number of sizes which specify the sizes of the source - tensor type. + tensor type. "offset + size" must not exceed the respective + destination dimension size. * strides: tensor-rank number of strides that specify subsampling in each dimension. @@ -1397,9 +1405,11 @@ * source: the tensor that is inserted. * dest: the tensor into which the source tensor is inserted. * offsets: tensor-rank number of offsets into the `dest` tensor into which - the slice is inserted. + the slice is inserted. Offsets must be non-negative and within + the bounds of the destination tensor. * sizes: tensor-rank number of sizes which specify the sizes of the source - tensor type. + tensor type. "offset + size" must not exceed the respective + destination dimension size. * strides: tensor-rank number of strides that specify subsampling in each dimension. diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp --- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -166,6 +166,31 @@ return droppedDims; } +LogicalResult verifySliceOffsetsAndSizes(Operation *op, RankedTensorType type, + ArrayRef mixedOffsets, + ArrayRef mixedSizes) { + // No out-of-bounds accesses. + for (int64_t i = 0, e = type.getRank(); i < e; ++i) { + int64_t offset = mixedOffsets[i]; + int64_t size = mixedSizes[i]; + if (!type.isDynamicDim(i)) { + if (!ShapedType::isDynamic(offset) && offset >= type.getDimSize(i)) + return op->emitOpError("offset #") << i << " is out of bounds"; + if (!ShapedType::isDynamic(offset) && !ShapedType::isDynamic(size) && + offset + size > type.getDimSize(i)) + return op->emitOpError("dimension #") << i << " runs out of bounds"; + } + } + + // No negative offsets. + for (const auto &it : llvm::enumerate(mixedOffsets)) + if (!ShapedType::isDynamic(it.value())) + if (it.value() < 0) + return op->emitOpError("offset #") << it.index() << " is negative"; + + return success(); +} + //===----------------------------------------------------------------------===// // BitcastOp //===----------------------------------------------------------------------===// @@ -1839,7 +1864,10 @@ RankedTensorType expectedType = ExtractSliceOp::inferResultType( getSourceType(), getMixedOffsets(), getMixedSizes(), getMixedStrides()); SliceVerificationResult result = isRankReducedType(expectedType, getType()); - return produceSliceErrorMsg(result, *this, expectedType); + if (failed(produceSliceErrorMsg(result, *this, expectedType))) + return failure(); + return verifySliceOffsetsAndSizes(*this, getSourceType(), getStaticOffsets(), + getStaticSizes()); } llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { @@ -2221,7 +2249,10 @@ SliceVerificationResult result = verifyInsertSliceOp(getSourceType(), getType(), getStaticOffsets(), getStaticSizes(), getStaticStrides(), &expectedType); - return produceSliceErrorMsg(result, *this, expectedType); + if (failed(produceSliceErrorMsg(result, *this, expectedType))) + return failure(); + return verifySliceOffsetsAndSizes(*this, getDestType(), getStaticOffsets(), + getStaticSizes()); } /// If we have two consecutive InsertSliceOp writing to the same slice, we @@ -3126,7 +3157,10 @@ SliceVerificationResult result = verifyInsertSliceOp(getSourceType(), getDestType(), getStaticOffsets(), getStaticSizes(), getStaticStrides(), &expectedType); - return produceSliceErrorMsg(result, *this, expectedType); + if (failed(produceSliceErrorMsg(result, *this, expectedType))) + return failure(); + return verifySliceOffsetsAndSizes(*this, getDestType(), getStaticOffsets(), + getStaticSizes()); } void ParallelInsertSliceOp::getCanonicalizationPatterns( diff --git a/mlir/test/Dialect/Linalg/vectorize-tensor-extract-masked.mlir b/mlir/test/Dialect/Linalg/vectorize-tensor-extract-masked.mlir --- a/mlir/test/Dialect/Linalg/vectorize-tensor-extract-masked.mlir +++ b/mlir/test/Dialect/Linalg/vectorize-tensor-extract-masked.mlir @@ -89,7 +89,7 @@ // ----- func.func @masked_vectorize_nd_tensor_extract_with_affine_apply_gather(%6: tensor<80x16xf32>, %arg0: index, %extracted_slice : tensor<1x3xf32>) -> tensor<1x3xf32> { - %c16 = arith.constant 16 : index + %c15 = arith.constant 15 : index %1 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"] @@ -97,7 +97,7 @@ ^bb0(%out: f32): %2 = linalg.index 1 : index %3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%2, %arg0) - %extracted = tensor.extract %6[%3, %c16] : tensor<80x16xf32> + %extracted = tensor.extract %6[%3, %c15] : tensor<80x16xf32> linalg.yield %extracted : f32 } -> tensor<1x3xf32> return %1 : tensor<1x3xf32> diff --git a/mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir b/mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir --- a/mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir +++ b/mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir @@ -341,7 +341,7 @@ // The vectorizer converts `affine.apply` so that the subsequent Ops can be vectorised based on the converted ops. Gather load. func.func @vectorize_nd_tensor_extract_with_affine_apply_gather(%6: tensor<80x16xf32>, %arg0: index, %extracted_slice : tensor<1x4xf32>) -> tensor<1x4xf32> { - %c16 = arith.constant 16 : index + %c15 = arith.constant 15 : index %1 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"] @@ -349,7 +349,7 @@ ^bb0(%out: f32): %2 = linalg.index 1 : index %3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%2, %arg0) - %extracted = tensor.extract %6[%3, %c16] : tensor<80x16xf32> + %extracted = tensor.extract %6[%3, %c15] : tensor<80x16xf32> linalg.yield %extracted : f32 } -> tensor<1x4xf32> return %1 : tensor<1x4xf32> @@ -363,12 +363,13 @@ // CHECK: %[[VAL_4:.*]] = arith.constant dense : vector<1x4xi1> // CHECK: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<1x4xf32> // CHECK: %[[VAL_6:.*]] = arith.constant 0 : index -// CHECK: %[[VAL_7:.*]] = arith.constant dense<16> : vector<1x4xindex> +// CHECK: %[[VAL_7a:.*]] = arith.constant dense<16> : vector<1x4xindex> +// CHECK: %[[VAL_7b:.*]] = arith.constant dense<15> : vector<1x4xindex> // CHECK: %[[VAL_8:.*]] = vector.broadcast %[[VAL_1]] : index to vector<4xindex> // CHECK: %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_3]] : vector<4xindex> // CHECK: %[[VAL_10:.*]] = vector.broadcast %[[VAL_9]] : vector<4xindex> to vector<1x4xindex> -// CHECK: %[[VAL_11:.*]] = arith.muli %[[VAL_10]], %[[VAL_7]] : vector<1x4xindex> -// CHECK: %[[VAL_12:.*]] = arith.addi %[[VAL_11]], %[[VAL_7]] : vector<1x4xindex> +// CHECK: %[[VAL_11:.*]] = arith.muli %[[VAL_10]], %[[VAL_7a]] : vector<1x4xindex> +// CHECK: %[[VAL_12:.*]] = arith.addi %[[VAL_11]], %[[VAL_7b]] : vector<1x4xindex> // CHECK: %[[VAL_13:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {{\[}}%[[VAL_12]]], %[[VAL_4]], %[[VAL_5]] : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32> // CHECK: %[[VAL_14:.*]] = vector.transfer_write %[[VAL_13]], %[[VAL_2]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32> // CHECK: return %[[VAL_14]] : tensor<1x4xf32> diff --git a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir --- a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir +++ b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir @@ -263,19 +263,17 @@ // ----- -// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 2)> // CHECK-LABEL: func @insert_slice_of_insert_slice( // CHECK-SAME: %[[t:[0-9a-z]*]]: tensor // CHECK-SAME: %[[r1:[0-9a-z]*]]: tensor<1x14xf32> // CHECK-SAME: %[[pos:[0-9a-z]*]]: index -// CHECK: %[[add:.*]] = affine.apply #[[$map]]()[%[[pos]]] -// CHECK: tensor.insert_slice %[[t]] into %[[r1]][4, %[[add]]] [1, 1] [1, 1] : tensor into tensor<1x14xf32> +// CHECK: tensor.insert_slice %[[t]] into %[[r1]][0, %[[pos]]] [1, 1] [1, 1] : tensor into tensor<1x14xf32> func.func @insert_slice_of_insert_slice(%t: tensor, %r0: tensor<1x1xf32>, %r1: tensor<1x14xf32>, %pos: index) -> tensor<1x14xf32> { - %0 = tensor.insert_slice %t into %r0[1, 2] [1, 1] [1, 1] + %0 = tensor.insert_slice %t into %r0[0, 0] [1, 1] [1, 1] : tensor into tensor<1x1xf32> - %1 = tensor.insert_slice %0 into %r1[3, %pos] [1, 1] [1, 1] + %1 = tensor.insert_slice %0 into %r1[0, %pos] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x14xf32> return %1 : tensor<1x14xf32> } @@ -286,13 +284,13 @@ // CHECK-SAME: %[[t:[0-9a-z]*]]: tensor // CHECK-SAME: %[[r1:[0-9a-z]*]]: tensor<1x14xf32> // CHECK-SAME: %[[pos:[0-9a-z]*]]: index -// CHECK: tensor.insert_slice %[[t]] into %[[r1]][5, %[[pos]]] [1, 1] [1, 1] : tensor into tensor<1x14xf32> +// CHECK: tensor.insert_slice %[[t]] into %[[r1]][0, %[[pos]]] [1, 1] [1, 1] : tensor into tensor<1x14xf32> func.func @insert_slice_of_insert_slice(%t: tensor, %r0: tensor<1xf32>, %r1: tensor<1x14xf32>, %pos: index) -> tensor<1x14xf32> { - %0 = tensor.insert_slice %t into %r0[2] [1] [1] + %0 = tensor.insert_slice %t into %r0[0] [1] [1] : tensor into tensor<1xf32> - %1 = tensor.insert_slice %0 into %r1[3, %pos] [1, 1] [1, 1] + %1 = tensor.insert_slice %0 into %r1[0, %pos] [1, 1] [1, 1] : tensor<1xf32> into tensor<1x14xf32> return %1 : tensor<1x14xf32> } diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir --- a/mlir/test/Dialect/Tensor/invalid.mlir +++ b/mlir/test/Dialect/Tensor/invalid.mlir @@ -200,6 +200,33 @@ // ----- +func.func @extract_slice_out_of_bounds(%t: tensor<5xf32>) { + // expected-error @+1 {{offset #0 is out of bounds}} + %0 = tensor.extract_slice %t[5][4][1] : tensor<5xf32> to tensor<4xf32> + + return +} + +// ----- + +func.func @extract_slice_out_of_bounds_2(%t: tensor<5xf32>) { + // expected-error @+1 {{dimension #0 runs out of bounds}} + %0 = tensor.extract_slice %t[3][4][1] : tensor<5xf32> to tensor<4xf32> + + return +} + +// ----- + +func.func @extract_slice_negative(%t: tensor<5xf32>) { + // expected-error @+1 {{offset #0 is negative}} + %0 = tensor.extract_slice %t[-1][4][1] : tensor<5xf32> to tensor<4xf32> + + return +} + +// ----- + func.func @insert_slice_wrong_result_rank(%t1: tensor, %t2: tensor, %idx : index) { // expected-error @+1 {{expected rank to be smaller or equal to the other rank.}} %0 = tensor.insert_slice %t2 into %t1[0][4][1] : tensor into tensor @@ -238,6 +265,36 @@ // ----- +func.func @insert_slice_out_of_bounds(%t1: tensor<4x4xf32>, %t2: tensor<8x16x2xf32>) { + // expected-error @+1 {{offset #2 is out of bounds}} + %0 = tensor.insert_slice %t1 into %t2[0, 0, 3][4, 4, 1][1, 1, 1] + : tensor<4x4xf32> into tensor<8x16x2xf32> + + return +} + +// ----- + +func.func @insert_slice_out_of_bounds_2(%t1: tensor<4x4xf32>, %t2: tensor<8x16x2xf32>) { + // expected-error @+1 {{dimension #1 runs out of bounds}} + %0 = tensor.insert_slice %t1 into %t2[0, 15, 0][4, 4, 1][1, 1, 1] + : tensor<4x4xf32> into tensor<8x16x2xf32> + + return +} + +// ----- + +func.func @insert_slice_negative(%t1: tensor<4x4xf32>, %t2: tensor<8x16x2xf32>) { + // expected-error @+1 {{offset #2 is negative}} + %0 = tensor.insert_slice %t1 into %t2[0, 0, -1][4, 4, 1][1, 1, 1] + : tensor<4x4xf32> into tensor<8x16x2xf32> + + return +} + +// ----- + func.func @illegal_expanding_reshape_dynamic_tensor (%arg0: tensor) -> tensor { // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}