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, @@ -343,13 +349,13 @@ // Copy the Operand's rank. if (!hasRankedInput) - outputShape.resize(operandTy.getRank(), -1); + outputShape.resize(operandTy.getRank(), ShapedType::kDynamicSize); // Copy shapes until the dim is non-dynamic. for (int i = 0, s = operandTy.getRank(); i < s; i++) { if (i == axis || operandTy.isDynamicDim(i)) continue; - if (outputShape[i] == -1) + if (outputShape[i] == ShapedType::kDynamicSize) outputShape[i] = operandTy.getDimSize(i); if (outputShape[i] != operandTy.getDimSize(i)) return failure(); @@ -371,7 +377,7 @@ // We need to know the length of the concatenation axis of all inputs to // determine the dimension size of the output shape. if (!operandTy.hasRank() || operandTy.isDynamicDim(axis)) { - concatDimSize = -1; + concatDimSize = ShapedType::kDynamicSize; break; } @@ -394,7 +400,7 @@ // All shapes are dynamic. SmallVector outShape; - outShape.resize(2, -1); + outShape.resize(2, ShapedType::kDynamicSize); if (inputTy.hasRank()) { outShape[0] = inputTy.getDimSize(0); @@ -405,7 +411,8 @@ } if (biasTy.hasRank()) { - outShape[1] = outShape[1] == -1 ? biasTy.getDimSize(0) : outShape[1]; + outShape[1] = outShape[1] == ShapedType::kDynamicSize ? biasTy.getDimSize(0) + : outShape[1]; } inferredReturnShapes.push_back(ShapedTypeComponents(outShape)); @@ -421,7 +428,7 @@ // All shapes are dynamic. SmallVector outShape; - outShape.resize(3, -1); + outShape.resize(3, ShapedType::kDynamicSize); if (lhsTy.hasRank()) { outShape[0] = lhsTy.getDimSize(0); @@ -429,7 +436,8 @@ } if (rhsTy.hasRank()) { - outShape[0] = outShape[0] == -1 ? rhsTy.getDimSize(0) : outShape[0]; + outShape[0] = outShape[0] == ShapedType::kDynamicSize ? rhsTy.getDimSize(0) + : outShape[0]; outShape[2] = rhsTy.getDimSize(2); } @@ -460,7 +468,7 @@ return success(); } - outputShape.resize(paddingTy.getDimSize(0), -1); + outputShape.resize(paddingTy.getDimSize(0), ShapedType::kDynamicSize); inferredReturnShapes.push_back(ShapedTypeComponents(outputShape)); return success(); } @@ -468,7 +476,7 @@ DenseIntElementsAttr paddings; // If the paddings value is not a constant, all dimensions must be dynamic. if (!matchPattern(operands[1], m_Constant(&paddings))) { - outputShape.resize(inputTy.getRank(), -1); + outputShape.resize(inputTy.getRank(), ShapedType::kDynamicSize); inferredReturnShapes.push_back(ShapedTypeComponents(outputShape)); return success(); } @@ -481,7 +489,7 @@ outputShape.reserve(inputTy.getRank()); for (int i = 0, s = inputTy.getRank(); i < s; i++) { if (inputTy.isDynamicDim(i)) { - outputShape.push_back(-1); + outputShape.push_back(ShapedType::kDynamicSize); continue; } @@ -531,7 +539,7 @@ ShapedType inputTy = operands[0].getType().cast(); SmallVector outputShape; if (!inputTy.hasRank()) { - outputShape.resize(multiples.size(), -1); + outputShape.resize(multiples.size(), ShapedType::kDynamicSize); inferredReturnShapes.push_back(ShapedTypeComponents(outputShape)); return success(); } @@ -547,7 +555,7 @@ outputShape.reserve(multiples.size()); for (int i = 0, s = inputTy.getRank(); i < s; i++) { int dim = inputTy.getDimSize(i); - if (dim != -1) + if (dim != ShapedType::kDynamicSize) dim *= multipleValues[i]; outputShape.push_back(dim); } @@ -579,14 +587,14 @@ int64_t numElements = type.getNumElements(); int64_t staticMul = 1; for (auto val : newShapeValue) { - if (val != -1) { + if (val != ShapedType::kDynamicSize) { staticMul *= val; } } // Determine the length of the dynamic dimension. for (auto &val : newShapeValue) { - if (val == -1) + if (val == ShapedType::kDynamicSize) val = numElements / staticMul; } @@ -612,7 +620,7 @@ // can determine the output rank. SmallVector outputShape; if (!inputTy.hasRank()) { - outputShape.resize(permsTy.getDimSize(0), -1); + outputShape.resize(permsTy.getDimSize(0), ShapedType::kDynamicSize); inferredReturnShapes.push_back(ShapedTypeComponents(outputShape)); return success(); } @@ -641,7 +649,7 @@ } DenseIntElementsAttr perms; - outputShape.resize(inputTy.getRank(), -1); + outputShape.resize(inputTy.getRank(), ShapedType::kDynamicSize); // If the permuations are a constant we can directly determine the output // shape. if (matchPattern(operands[1], m_Constant(&perms))) { @@ -665,7 +673,7 @@ ValueRange operands, DictionaryAttr attributes, RegionRange regions, SmallVectorImpl &inferredReturnShapes) { llvm::SmallVector outputShape; - outputShape.resize(3, -1); + outputShape.resize(3, ShapedType::kDynamicSize); if (auto ty = operands[0].getType().dyn_cast()) { outputShape[0] = ty.getDimSize(0); @@ -673,9 +681,9 @@ } if (auto ty = operands[1].getType().dyn_cast()) { - if (outputShape[0] == -1) + if (outputShape[0] == ShapedType::kDynamicSize) outputShape[0] = ty.getDimSize(0); - if (outputShape[1] == -1) + if (outputShape[1] == ShapedType::kDynamicSize) outputShape[1] = ty.getDimSize(1); } @@ -683,12 +691,82 @@ 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, ShapedType::kDynamicSize); + + int32_t inHeight = ShapedType::kDynamicSize; + int32_t inWidth = ShapedType::kDynamicSize; + + 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); + + // If we have a 0 zero in integers we know that the resize indexing needs to + // be performed in floating point. Use the floating point varient to compute + // the resize shape. + 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] == ShapedType::kDynamicSize && 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] == ShapedType::kDynamicSize && 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] == ShapedType::kDynamicSize && inHeight >= 0 && !fpMode) { + int64_t size = (inHeight - 1); + size = ((size << shift) - offsetInt[0]) / strideInt[0]; + outputShape[1] = size + 1; + } + + if (outputShape[2] == ShapedType::kDynamicSize && 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, SmallVectorImpl &inferredReturnShapes) { llvm::SmallVector outputShape; - outputShape.resize(3, -1); + outputShape.resize(3, ShapedType::kDynamicSize); if (auto ty = operands[0].getType().dyn_cast()) { outputShape[0] = ty.getDimSize(0); @@ -697,14 +775,14 @@ } if (auto ty = operands[1].getType().dyn_cast()) { - if (outputShape[0] == -1) + if (outputShape[0] == ShapedType::kDynamicSize) outputShape[0] = ty.getDimSize(0); } if (auto ty = operands[2].getType().dyn_cast()) { - if (outputShape[0] == -1) + if (outputShape[0] == ShapedType::kDynamicSize) outputShape[0] = ty.getDimSize(0); - if (outputShape[2] == -1) + if (outputShape[2] == ShapedType::kDynamicSize) outputShape[2] = ty.getDimSize(2); } @@ -814,6 +892,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 +902,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 +917,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 +}