Skip to content

Commit ab7e8b7

Browse files
authored
[mlir][tosa] fix a crash when sliceOp has invalid attribute (#68486) (#70063)
add a verifier for tosa::sliceOp
1 parent ff67b68 commit ab7e8b7

File tree

3 files changed

+35
-0
lines changed

3 files changed

+35
-0
lines changed

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1621,6 +1621,7 @@ def Tosa_SliceOp : Tosa_InferShapedTypeOp<"slice"> {
16211621

16221622
let hasCanonicalizer = 1;
16231623
let hasFolder = 1;
1624+
let hasVerifier = 1;
16241625
}
16251626

16261627
//===----------------------------------------------------------------------===//

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -820,6 +820,22 @@ LogicalResult tosa::SliceOp::inferReturnTypeComponents(
820820
return success();
821821
}
822822

823+
LogicalResult tosa::SliceOp::verify() {
824+
auto inputType = llvm::dyn_cast<RankedTensorType>(getInput().getType());
825+
if (!inputType)
826+
return success();
827+
828+
if (static_cast<size_t>(inputType.getRank()) != getStart().size())
829+
return emitOpError(
830+
"length of start attribute is not equal rank of input shape");
831+
832+
if (static_cast<size_t>(inputType.getRank()) != getSize().size())
833+
return emitOpError(
834+
"length of size attribute is not equal rank of input shape");
835+
836+
return success();
837+
}
838+
823839
LogicalResult tosa::TableOp::inferReturnTypeComponents(
824840
MLIRContext *context, ::std::optional<Location> location,
825841
TableOp::Adaptor adaptor,

mlir/test/Dialect/Tosa/invalid.mlir

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,3 +303,21 @@ func.func @test_variable_write_shape(%arg0: tensor<1x4x8xi32>) -> () {
303303
tosa.variable.write @stored_var, %arg0 : tensor<1x4x8xi32>
304304
return
305305
}
306+
307+
// -----
308+
309+
func.func @test_slice_invalid_start() {
310+
%0 = tensor.empty() : tensor<4x31x31xf32>
311+
// expected-error@+1 {{'tosa.slice' op length of start attribute is not equal rank of input shape}}
312+
%1 = tosa.slice %0 {size = array<i64: 1, 1, 1>, start = array<i64: 1, 1>} : (tensor<4x31x31xf32>) -> tensor<*xf32>
313+
return
314+
}
315+
316+
// -----
317+
318+
func.func @test_slice_invalid_size() {
319+
%0 = tensor.empty() : tensor<4x31x31xf32>
320+
// expected-error@+1 {{'tosa.slice' op length of size attribute is not equal rank of input shape}}
321+
%1 = tosa.slice %0 {size = array<i64: 1>, start = array<i64: 1, 1, 1>} : (tensor<4x31x31xf32>) -> tensor<*xf32>
322+
return
323+
}

0 commit comments

Comments
 (0)