diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp --- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp +++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp @@ -590,6 +590,12 @@ return success(); } +static SmallVector convertToMlirShape(ArrayRef shape) { + return to_vector(llvm::map_range(shape, [](int64_t dim) { + return dim == -1 ? ShapedType::kDynamicSize : dim; + })); +} + LogicalResult tosa::SliceOp::inferReturnTypeComponents( MLIRContext *context, ::llvm::Optional location, ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions, @@ -601,7 +607,8 @@ outputShape.push_back(val.cast().getValue().getSExtValue()); } - inferredReturnShapes.push_back(ShapedTypeComponents(outputShape)); + inferredReturnShapes.push_back(ShapedTypeComponents( + convertToMlirShape(outputShape))); return success(); } @@ -655,11 +662,6 @@ return success(); } -static SmallVector convertToMlirShape(ArrayRef shape) { - return to_vector(llvm::map_range(shape, [](int64_t dim) { - return dim == -1 ? ShapedType::kDynamicSize : dim; - })); -} LogicalResult tosa::ReshapeOp::inferReturnTypeComponents( MLIRContext *context, ::llvm::Optional location, diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir --- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir +++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir @@ -539,6 +539,15 @@ // ----- +// CHECK-LABEL: @test_slice_dynamic +func.func @test_slice_dynamic(%arg0 : tensor<10x?x2xf32>) -> () { + // CHECK: "tosa.slice"(%arg0) {size = [7, -1, 1], start = [1, 0, 0]} : (tensor<10x?x2xf32>) -> tensor<7x?x1xf32> + %0 = "tosa.slice"(%arg0) {size = [7, -1, 1], start = [1, 0, 0]} : (tensor<10x?x2xf32>) -> tensor + return +} + +// ----- + // CHECK-LABEL: @test_tile func.func @test_tile(%arg0 : tensor<2x3x?xi32>) -> () { // CHECK: "tosa.tile"(%arg0) {multiples = [2, 1, 5]} : (tensor<2x3x?xi32>) -> tensor<4x3x?xi32>