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,30 @@ 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) { + std::optional offset = getConstantIntValue(mixedOffsets[i]); + std::optional size = getConstantIntValue(mixedSizes[i]); + if (!type.isDynamicDim(i)) { + if (offset && *offset >= type.getDimSize(i)) + return op->emitOpError("offset #") << i << " is out of bounds"; + if (offset && 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 (std::optional pos = getConstantIntValue(it.value())) + if (pos < 0) + return op->emitOpError("offset #") << it.index() << " is negative"; + + return success(); +} + //===----------------------------------------------------------------------===// // BitcastOp //===----------------------------------------------------------------------===// @@ -851,6 +875,20 @@ auto tensorType = llvm::cast(getTensor().getType()); if (tensorType.getRank() != static_cast(getIndices().size())) return emitOpError("incorrect number of indices for extract_element"); + + // No out-of-bounds indices. + for (const auto &it : llvm::enumerate(getIndices())) + if (!getTensor().getType().isDynamicDim(it.index())) + if (std::optional pos = getConstantIntValue(it.value())) + if (pos >= getTensor().getType().getDimSize(it.index())) + return emitOpError("index #") << it.index() << " is out of bounds"; + + // No negative indices. + for (const auto &it : llvm::enumerate(getIndices())) + if (std::optional pos = getConstantIntValue(it.value())) + if (pos < 0) + return emitOpError("index #") << it.index() << " is negative"; + return success(); } @@ -1078,6 +1116,20 @@ auto destType = llvm::cast(getDest().getType()); if (destType.getRank() != static_cast(getIndices().size())) return emitOpError("incorrect number of indices"); + + // No out-of-bounds indices. + for (const auto &it : llvm::enumerate(getIndices())) + if (!getDest().getType().isDynamicDim(it.index())) + if (std::optional pos = getConstantIntValue(it.value())) + if (pos >= getDest().getType().getDimSize(it.index())) + return emitOpError("index #") << it.index() << " is out of bounds"; + + // No negative indices. + for (const auto &it : llvm::enumerate(getIndices())) + if (std::optional pos = getConstantIntValue(it.value())) + if (pos < 0) + return emitOpError("index #") << it.index() << " is negative"; + return success(); } @@ -1839,7 +1891,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(), getMixedOffsets(), + getMixedSizes()); } llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { @@ -2221,7 +2276,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(), getMixedOffsets(), + getMixedSizes()); } /// If we have two consecutive InsertSliceOp writing to the same slice, we @@ -3126,7 +3184,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(), getMixedOffsets(), + getMixedSizes()); } 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/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir --- a/mlir/test/Dialect/Tensor/canonicalize.mlir +++ b/mlir/test/Dialect/Tensor/canonicalize.mlir @@ -258,51 +258,6 @@ // ----- -// Ensure the optimization doesn't segfault from bad constants -// CHECK-LABEL: func @extract_negative_from_tensor.from_elements -func.func @extract_negative_from_tensor.from_elements(%element : index) -> index { - // CHECK-SAME: ([[ARG:%.*]]: index) - %c-1 = arith.constant -1 : index - %tensor = tensor.from_elements %element : tensor<1xindex> - %extracted_element = tensor.extract %tensor[%c-1] : tensor<1xindex> - // CHECK: tensor.from_elements - // CHECK: %[[RESULT:.*]] = tensor.extract - // CHECK: return %[[RESULT]] - return %extracted_element : index -} - -// ----- - -// Ensure the optimization doesn't segfault from bad constants -// CHECK-LABEL: func @extract_oob_from_tensor.from_elements -func.func @extract_oob_from_tensor.from_elements(%element : index) -> index { - // CHECK-SAME: ([[ARG:%.*]]: index) - %c1 = arith.constant 1 : index - %tensor = tensor.from_elements %element : tensor<1xindex> - %extracted_element = tensor.extract %tensor[%c1] : tensor<1xindex> - // CHECK: tensor.from_elements - // CHECK: %[[RESULT:.*]] = tensor.extract - // CHECK: return %[[RESULT]] - return %extracted_element : index -} - -// ----- - -// Ensure the optimization doesn't segfault from bad constants -// CHECK-LABEL: func @extract_oob_from_tensor.from_elements -func.func @extract_oob_from_tensor.from_elements(%element : index) -> index { - // CHECK-SAME: ([[ARG:%.*]]: index) - %c2 = arith.constant 2 : index - %tensor = tensor.from_elements %element : tensor<1xindex> - %extracted_element = tensor.extract %tensor[%c2] : tensor<1xindex> - // CHECK: tensor.from_elements - // CHECK: %[[RESULT:.*]] = tensor.extract - // CHECK: return %[[RESULT]] - return %extracted_element : index -} - -// ----- - // CHECK-LABEL: func @extract_from_tensor.generate // CHECK-SAME: %[[IDX:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32> func.func @extract_from_tensor.generate(%idx: index, %tensor: tensor<*xf32>) -> index { 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 @@ -24,6 +24,24 @@ // ----- +func.func @extract_out_of_bounds(%arg0: tensor<7xf32>) { + %c7 = arith.constant 7 : index + // expected-error@+1 {{index #0 is out of bounds}} + %0 = tensor.extract %arg0[%c7] : tensor<7xf32> + return +} + +// ----- + +func.func @extract_negative(%arg0: tensor<7xf32>) { + %c1 = arith.constant -1 : index + // expected-error@+1 {{index #0 is negative}} + %0 = tensor.extract %arg0[%c1] : tensor<7xf32> + return +} + +// ----- + func.func @insert_too_many_indices(%arg0: f32, %arg1: tensor) { // expected-error@+1 {{incorrect number of indices}} %0 = tensor.insert %arg0 into %arg1[] : tensor @@ -32,6 +50,24 @@ // ----- +func.func @insert_out_of_bounds(%arg0: tensor<7xf32>, %arg1: f32) { + %c7 = arith.constant 7 : index + // expected-error@+1 {{index #0 is out of bounds}} + %0 = tensor.insert %arg1 into %arg0[%c7] : tensor<7xf32> + return +} + +// ----- + +func.func @insert_negative(%arg0: tensor<7xf32>, %arg1: f32) { + %c1 = arith.constant -1 : index + // expected-error@+1 {{index #0 is negative}} + %0 = tensor.insert %arg1 into %arg0[%c1] : tensor<7xf32> + return +} + +// ----- + func.func @tensor.from_elements_wrong_result_type() { // expected-error@+2 {{'tensor.from_elements' invalid kind of type specified}} %c0 = arith.constant 0 : i32 @@ -200,6 +236,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 +301,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)}}