diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td --- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td +++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td @@ -1562,7 +1562,10 @@ //===----------------------------------------------------------------------===// // Operator: resize //===----------------------------------------------------------------------===// -def Tosa_ResizeOp : Tosa_Op<"resize", [NoSideEffect]> { +def Tosa_ResizeOp : Tosa_Op<"resize", [ + DeclareOpInterfaceMethods, + NoSideEffect]> { let summary = "Resize operation, supports various resize/upsample modes"; @@ -1597,7 +1600,9 @@ //===----------------------------------------------------------------------===// // Operator: cast //===----------------------------------------------------------------------===// -def Tosa_CastOp: Tosa_Op<"cast", [NoSideEffect, SameOperandsAndResultShape]> { +def Tosa_CastOp: Tosa_Op<"cast", [NoSideEffect, SameOperandsAndResultShape, + DeclareOpInterfaceMethods]> { let summary = "Cast operation"; @@ -1635,7 +1640,9 @@ //===----------------------------------------------------------------------===// // Operator: rescale //===----------------------------------------------------------------------===// -def Tosa_RescaleOp: Tosa_Op<"rescale", [NoSideEffect]> { +def Tosa_RescaleOp: Tosa_Op<"rescale", [NoSideEffect, + DeclareOpInterfaceMethods]> { let summary = "Tosa rescale operator"; let description = [{ @@ -1703,7 +1710,9 @@ //===----------------------------------------------------------------------===// // Operator: identity //===----------------------------------------------------------------------===// -def Tosa_IdentityOp: Tosa_Op<"identity", [NoSideEffect]> { +def Tosa_IdentityOp: Tosa_Op<"identity", [NoSideEffect, + DeclareOpInterfaceMethods]> { let summary = "Identity operator"; let description = [{ Returns a tensor with the same shape, size, type 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 @@ -302,6 +302,12 @@ } } +static void getF64Values(ArrayAttr arrayAttr, SmallVector &values) { + for (auto it : arrayAttr) { + values.push_back(it.cast().getValueAsDouble()); + } +} + LogicalResult tosa::ArgMaxOp::inferReturnTypeComponents( MLIRContext *context, ::llvm::Optional location, ValueRange operands, DictionaryAttr attributes, RegionRange regions, @@ -683,6 +689,73 @@ return success(); } +LogicalResult tosa::ResizeOp::inferReturnTypeComponents( + MLIRContext *context, ::llvm::Optional location, + ValueRange operands, DictionaryAttr attributes, RegionRange regions, + SmallVectorImpl &inferredReturnShapes) { + llvm::SmallVector outputShape; + outputShape.resize(4, -1); + + int32_t inHeight = -1; + int32_t inWidth = -1; + + if (auto ty = operands[0].getType().dyn_cast()) { + outputShape[0] = ty.getDimSize(0); + outputShape[3] = ty.getDimSize(3); + + inHeight = ty.getDimSize(1); + inWidth = ty.getDimSize(2); + } + + int32_t shift = + attributes.get("shift").cast().getValue().getSExtValue(); + llvm::SmallVector newShape; + getI64Values(attributes.get("output_size").cast(), newShape); + outputShape[1] = newShape[0]; + outputShape[2] = newShape[1]; + + llvm::SmallVector strideInt; + llvm::SmallVector offsetInt; + llvm::SmallVector strideFp; + llvm::SmallVector offsetFp; + getI64Values(attributes.get("offset").cast(), offsetInt); + getF64Values(attributes.get("offset_fp").cast(), offsetFp); + getI64Values(attributes.get("stride").cast(), strideInt); + getF64Values(attributes.get("stride_fp").cast(), strideFp); + + bool fpMode = strideInt[0] == 0; + + // We can compute the output shape if attribute specifies unknown dimensions + // based on the offset and stride. If we perfectly line up to the last index + // we need to round up the size to include it. + if (outputShape[1] == -1 && inHeight >= 0 && fpMode) { + float sizeFp = (inHeight - offsetFp[0] - 1) / strideFp[0]; + float round = std::floor(sizeFp) == sizeFp ? 1 : 0; + outputShape[1] = std::ceil(sizeFp) + round; + } + + if (outputShape[2] == -1 && inWidth >= 0 && fpMode) { + float sizeFp = (inWidth - offsetFp[1] - 1) / strideFp[1]; + float round = std::floor(sizeFp) == sizeFp ? 1 : 0; + outputShape[2] = std::ceil(sizeFp) + round; + } + + if (outputShape[1] == -1 && inHeight >= 0 && !fpMode) { + int64_t size = (inHeight - 1); + size = ((size << shift) - offsetInt[0]) / strideInt[0]; + outputShape[1] = size + 1; + } + + if (outputShape[2] == -1 && inWidth >= 0 && !fpMode) { + int64_t size = (inWidth - 1); + size = ((size << shift) - offsetInt[1]) / strideInt[1]; + outputShape[2] = size + 1; + } + + inferredReturnShapes.push_back(ShapedTypeComponents(outputShape)); + return success(); +} + LogicalResult tosa::ScatterOp::inferReturnTypeComponents( MLIRContext *context, ::llvm::Optional location, ValueRange operands, DictionaryAttr attributes, RegionRange regions, @@ -814,6 +887,7 @@ NARY_SHAPE_INFER(tosa::BitwiseOrOp) NARY_SHAPE_INFER(tosa::BitwiseXorOp) NARY_SHAPE_INFER(tosa::BitwiseNotOp) +NARY_SHAPE_INFER(tosa::CastOp) NARY_SHAPE_INFER(tosa::CeilOp) NARY_SHAPE_INFER(tosa::ClampOp) NARY_SHAPE_INFER(tosa::ClzOp) @@ -823,6 +897,7 @@ NARY_SHAPE_INFER(tosa::FloorOp) NARY_SHAPE_INFER(tosa::GreaterEqualOp) NARY_SHAPE_INFER(tosa::GreaterOp) +NARY_SHAPE_INFER(tosa::IdentityOp) NARY_SHAPE_INFER(tosa::LogOp) NARY_SHAPE_INFER(tosa::LogicalAndOp) NARY_SHAPE_INFER(tosa::LogicalLeftShiftOp) @@ -837,6 +912,7 @@ NARY_SHAPE_INFER(tosa::PowOp) NARY_SHAPE_INFER(tosa::ReciprocalOp) NARY_SHAPE_INFER(tosa::ReluNOp) +NARY_SHAPE_INFER(tosa::RescaleOp) NARY_SHAPE_INFER(tosa::ReverseOp) NARY_SHAPE_INFER(tosa::RsqrtOp) NARY_SHAPE_INFER(tosa::SelectOp) 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 @@ -65,6 +65,9 @@ // CHECK: "tosa.sigmoid"(%arg0) : (tensor<4xf32>) -> tensor<4xf32> %12 = "tosa.sigmoid"(%arg0) : (tensor<4xf32>) -> tensor<*xf32> + + // CHECK: "tosa.cast"(%arg0) : (tensor<4xf32>) -> tensor<4xi32> + %13 = "tosa.cast"(%arg0) : (tensor<4xf32>) -> tensor<*xi32> return } @@ -92,6 +95,12 @@ // CHECK: "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<4xi32>) -> tensor<4xi32> %6 = "tosa.reverse"(%arg0) { axis = 0 : i64 } : (tensor<4xi32>) -> tensor + + // CHECK: "tosa.rescale"(%arg0) {{.+}} : (tensor<4xi32>) -> tensor<4xi16> + %7 = "tosa.rescale"(%arg0) {input_zp = 243 : i32, output_zp = 252 : i32, multiplier = [42 : i32, 43 : i32], shift = [14 : i32, 15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<4xi32>) -> (tensor<*xi16>) + + // CHECK: "tosa.identity"(%arg0) : (tensor<4xi32>) -> tensor<4xi32> + %8 = "tosa.identity"(%arg0) : (tensor<4xi32>) -> tensor return } @@ -660,3 +669,66 @@ %0 = "tosa.scatter"(%arg0, %arg1, %arg2) : (tensor, tensor<3x?xi32>, tensor) -> (tensor) return } + +// ----- + +// CHECK-LABEL: @resize_output_size +func @resize_output_size(%arg0: tensor<2x?x?x3xi32>) { + // CHECK: -> tensor<2x4x5x3xi32> + %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [0, 1], offset_fp = [0.000000e+00 : f32, 0.000000e+00 : f32], output_size = [4, 5], shift = 8 : i32, stride = [1, 1], stride_fp = [0.000000e+00 : f32, 0.000000e+00 : f32]} : (tensor<2x?x?x3xi32>) -> tensor + return +} + +// ----- + +// CHECK-LABEL: @resize_int_horizontal +func @resize_int_horizontal(%arg0: tensor<1x2x4x1xi32>) { + // CHECK: -> tensor<1x2x7x1xi32> + %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [0, 0], offset_fp = [0.000000e+00 : f32, 0.000000e+00 : f32], output_size = [-1, -1], shift = 8 : i32, stride = [256, 128], stride_fp = [0.000000e+00 : f32, 0.000000e+00 : f32]} : (tensor<1x2x4x1xi32>) -> tensor + return +} + +// ----- + +// CHECK-LABEL: @resize_int_vertical +func @resize_int_vertical(%arg0: tensor<1x2x4x1xi32>) { + // CHECK: -> tensor<1x3x4x1xi32> + %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [0, 0], offset_fp = [0.000000e+00 : f32, 0.000000e+00 : f32], output_size = [-1, -1], shift = 8 : i32, stride = [128, 256], stride_fp = [0.000000e+00 : f32, 0.000000e+00 : f32]} : (tensor<1x2x4x1xi32>) -> tensor + return +} + +// ----- + +// CHECK-LABEL: @resize_int_offsetted +func @resize_int_offsetted(%arg0: tensor<1x2x4x1xi32>) { + // CHECK: -> tensor<1x4x6x1xi32> + %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [64, 64], offset_fp = [0.000000e+00 : f32, 0.000000e+00 : f32], output_size = [-1, -1], shift = 8 : i32, stride = [64, 128], stride_fp = [0.000000e+00 : f32, 0.000000e+00 : f32]} : (tensor<1x2x4x1xi32>) -> tensor + return +} + +// ----- + +// CHECK-LABEL: @resize_fp_horizontal +func @resize_fp_horizontal(%arg0: tensor<1x2x4x1xi32>) { + // CHECK: -> tensor<1x2x7x1xi32> + %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [0, 0], offset_fp = [0.000000e+00 : f32, 0.000000e+00 : f32], output_size = [-1, -1], shift = 0 : i32, stride = [0, 0], stride_fp = [1.000000e+00 : f32, 5.000000e-01 : f32]} : (tensor<1x2x4x1xi32>) -> tensor + return +} + +// ----- + +// CHECK-LABEL: @resize_fp_vertical +func @resize_fp_vertical(%arg0: tensor<1x2x4x1xi32>) { + // CHECK: -> tensor<1x3x4x1xi32> + %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [0, 0], offset_fp = [0.000000e+00 : f32, 0.000000e+00 : f32], output_size = [-1, -1], shift = 0 : i32, stride = [0, 0], stride_fp = [5.000000e-01 : f32, 1.000000e+00 : f32]} : (tensor<1x2x4x1xi32>) -> tensor + return +} + +// ----- + +// CHECK-LABEL: @resize_fp_offsetted +func @resize_fp_offsetted(%arg0: tensor<1x2x4x1xi32>) { + // CHECK: -> tensor<1x4x6x1xi32> + %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [0, 0], offset_fp = [2.500000e-01 : f32, 2.500000e-01 : f32], output_size = [-1, -1], shift = 0 : i32, stride = [0, 0], stride_fp = [2.500000e-01 : f32, 5.000000e-01 : f32]} : (tensor<1x2x4x1xi32>) -> tensor + return +}