[mlir][tosa] Shape inference for a few remaining easy cases:

[rsuderman]

Summary

Handles shape inference for identity, cast, and rescale. These were missed
during the initialy elementwise work. This includes resize shape propagation
which includes both attribute and input type based propagation.

[jpienaar]

I don't know the ops that well, I could verify these if interested, but was going to rely on Suraj else :)

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

699

Could we use ShapedType::kDynamic here and below?

728

fpMode makes me think floating point mode, and then I get confused why it is triggered by strideInt. Could you add a comment here?

Looking at the op, does it always have stride and stride_fp? And is only one of them set at a time? Why not have both be optional and then verify that exactly one is set?

736

So this is ceiling of size + 1 if sizeFp was already int sized? Would any of llvm/Support/MathExtras.h work here?

[eric-k256]

In D105845#2883409, @jpienaar wrote:

I don't know the ops that well, I could verify these if interested, but was going to rely on Suraj else :)

Suraj is on sabbatical, so won't be looking at this. (At least he shouldn't be :) )

There's a problem with adding shape inference to the resize operator. With the integer version, you could end up with an off by one error in the output size when trying to calculate it from the integer shift/stride values. I think this is the only spot in the specification that explicitly requires the output shape to be specified. We can try to look and see if there is something we can change in the spec to make sure the inferred shape works.

Identity/Cast/Rescale all seem fine.

[rsuderman]

Added kDynamicSize use instead of explicit -1 with comment
about fpmode.

[rsuderman]

@eric-k256 We should see about removing that ambiguity to the spec. Its pretty valuable to determine output shapes entirely based on attributes / input shapes whenever possible. Specifying output_shape on construction helps with this but for something like resize, we should be able to find a better solution.

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

728

I think the specification should be updated to this. At this point TosaOps.td specifies both by values and branches on set value. I agree it would be worth changing.

736

I couldn't find any that seemed useful. The expectation is to round up, even in cases where the value is an exact integer. Floating point can be a bit ambiguous but hopefully this is not depended on too much (we should depend more on output_shape anyway).

[jpienaar]

Looks good thanks, I wonder if we could in a way fuzz this with the reference implementation

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

1566

You can use an OpTraitList to create a convenient "alias" for this.

[rsuderman]

In D105845#2923843, @jpienaar wrote:

Looks good thanks, I wonder if we could in a way fuzz this with the reference implementation

It should be possible to strip all types excluding the surrounding contract then perform shape propagation. In theory it should produce the same output. The main problem is converting TOSA flatbuffers to sample cases, recent changes in the flatbuffer broke my test generator.

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

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.11
 // Operator Class: Image Frontend Functions.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: resize
 //===----------------------------------------------------------------------===//
-def Tosa_ResizeOp : Tosa_Op<"resize", [NoSideEffect]> {
+def Tosa_ResizeOp : Tosa_Op<"resize", [
+      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                              ["inferReturnTypeComponents"]>,
+      NoSideEffect]> {
 
   let summary = "Resize operation, supports various resize/upsample modes";
 
   let description = [{
     Resizes a tensor. Resize is only allowed in the H and W dimensions. In
     expected use, stride_y is approximately (IH<<shift)/OH and stride_x is
     approximately (IW<<shift)/OW.  OH and OW are also supplied as inputs since
     there may be off by one errors if calculating OH and OW from the strides.

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.12
 // Operator Class: Type Conversion.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: cast
 //===----------------------------------------------------------------------===//
-def Tosa_CastOp: Tosa_Op<"cast", [NoSideEffect, SameOperandsAndResultShape]> {
+def Tosa_CastOp: Tosa_Op<"cast", [NoSideEffect, SameOperandsAndResultShape,
+      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                              ["inferReturnTypeComponents"]>]> {
 
   let summary = "Cast operation";
 
   let description = [{
     Performs a set of permissible cast operations
         Mode                    Input   Output
         ---------------------------------------
         signed 8 to bool        int8    Boolean
def Tosa_CastOp: Tosa_Op<"cast", [NoSideEffect, SameOperandsAndResultShape,
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: rescale
 //===----------------------------------------------------------------------===//
-def Tosa_RescaleOp: Tosa_Op<"rescale", [NoSideEffect]> {
+def Tosa_RescaleOp: Tosa_Op<"rescale", [NoSideEffect, 
+      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                              ["inferReturnTypeComponents"]>]> {
   let summary = "Tosa rescale operator";
 
   let description = [{
     Rescale quantized values into a new domain. Supported rescalings are:
     Mode                    Input   Output
     signed 8 to 8           int8    int8
     signed 8 to 16          int8    int16
     signed 8 to 32          int8    int32
let results = (outs
     Tosa_Tensor:$output
   );
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: identity
 //===----------------------------------------------------------------------===//
-def Tosa_IdentityOp: Tosa_Op<"identity", [NoSideEffect]> {
+def Tosa_IdentityOp: Tosa_Op<"identity", [NoSideEffect,
+      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                              ["inferReturnTypeComponents"]>]> {
   let summary = "Identity operator";
   let description = [{
     Returns a tensor with the same shape, size, type
     and content as the input.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1

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

 //===----------------------------------------------------------------------===//
 
 static void getI64Values(ArrayAttr arrayAttr, SmallVector<int64_t> &values) {
   for (auto it : arrayAttr) {
     values.push_back(it.cast<IntegerAttr>().getValue().getSExtValue());
   }
 }
 
+static void getF64Values(ArrayAttr arrayAttr, SmallVector<double> &values) {
+  for (auto it : arrayAttr) {
+    values.push_back(it.cast<FloatAttr>().getValueAsDouble());
+  }
+}
+
 LogicalResult tosa::ArgMaxOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ShapedType inputTy = operands[0].getType().cast<ShapedType>();
   IntegerAttr axis = attributes.get("axis").cast<IntegerAttr>();
   int32_t axisVal = axis.getValue().getSExtValue();
 
LogicalResult tosa::ConcatOp::inferReturnTypeComponents(
   bool hasRankedInput = false;
   for (auto operand : operands) {
     ShapedType operandTy = operand.getType().cast<ShapedType>();
     if (!operandTy.hasRank())
       continue;
 
     // 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();
     }
 
     hasRankedInput = true;
   }
 
   if (!hasRankedInput) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   // Determine the dimension size along the concatenation axis.
   int concatDimSize = 0;
   for (auto operand : operands) {
     ShapedType operandTy = operand.getType().cast<ShapedType>();
 
     // 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;
     }
 
     concatDimSize += operandTy.getDimSize(axis);
   }
 
   outputShape[axis] = concatDimSize;
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::FullyConnectedOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ShapedType inputTy = operands[0].getType().cast<ShapedType>();
   ShapedType weightTy = operands[1].getType().cast<ShapedType>();
   ShapedType biasTy = operands[2].getType().cast<ShapedType>();
 
   // All shapes are dynamic.
   SmallVector<int64_t> outShape;
-  outShape.resize(2, -1);
+  outShape.resize(2, ShapedType::kDynamicSize);
 
   if (inputTy.hasRank()) {
     outShape[0] = inputTy.getDimSize(0);
   }
 
   if (weightTy.hasRank()) {
     outShape[1] = weightTy.getDimSize(0);
   }
 
   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));
   return success();
 }
 
 LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ShapedType lhsTy = operands[0].getType().cast<ShapedType>();
   ShapedType rhsTy = operands[1].getType().cast<ShapedType>();
 
   // All shapes are dynamic.
   SmallVector<int64_t> outShape;
-  outShape.resize(3, -1);
+  outShape.resize(3, ShapedType::kDynamicSize);
 
   if (lhsTy.hasRank()) {
     outShape[0] = lhsTy.getDimSize(0);
     outShape[1] = lhsTy.getDimSize(1);
   }
 
   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);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
 LogicalResult tosa::PadOp::inferReturnTypeComponents(
LogicalResult tosa::PadOp::inferReturnTypeComponents(
   // If the input rank is unknown we can info the output rank using the padding
   // shape's first dim.
   if (!inputTy.hasRank()) {
     if (paddingTy.isDynamicDim(0)) {
       inferredReturnShapes.push_back(ShapedTypeComponents());
       return success();
     }
 
-    outputShape.resize(paddingTy.getDimSize(0), -1);
+    outputShape.resize(paddingTy.getDimSize(0), ShapedType::kDynamicSize);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   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();
   }
 
   SmallVector<int64_t> paddingValues;
   for (auto val : paddings) {
     paddingValues.push_back(val.getSExtValue());
   }
 
   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;
     }
 
     outputShape.push_back(inputTy.getDimSize(i) + paddingValues[i * 2] +
                           paddingValues[i * 2 + 1]);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));

 LogicalResult tosa::TileOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   auto multiples = attributes.get("multiples").cast<ArrayAttr>().getValue();
   ShapedType inputTy = operands[0].getType().cast<ShapedType>();
   SmallVector<int64_t> outputShape;
   if (!inputTy.hasRank()) {
-    outputShape.resize(multiples.size(), -1);
+    outputShape.resize(multiples.size(), ShapedType::kDynamicSize);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   // We need the multiple values to determine the output shape.
   SmallVector<int64_t> multipleValues;
   multipleValues.reserve(multiples.size());
   for (auto val : multiples) {
     multipleValues.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
   }
 
   // Any non dynamic dimension can be multiplied to a known size.
   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);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
   }
 
   // Determine the number of elements covered by the slice of all static
   // dimensions. This allows us to infer the length of the remaining dynamic
   // dimension.
   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;
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
   return success();
 }
 
 LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
if (!inputTy.hasRank() && (!permsTy.hasRank() || permsTy.isDynamicDim(0))) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   // Without the input dims we cannot determine the output dim sizes but we
   // can determine the output rank.
   SmallVector<int64_t> outputShape;
   if (!inputTy.hasRank()) {
-    outputShape.resize(permsTy.getDimSize(0), -1);
+    outputShape.resize(permsTy.getDimSize(0), ShapedType::kDynamicSize);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   // Rank-0 means no permutations matter.
   if (inputTy.getRank() == 0) {
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
   // permutation.
   if (allTheSame) {
     outputShape.resize(inputTy.getRank(), inputTy.getDimSize(0));
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   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))) {
     llvm::SmallVector<int64_t> permValues;
     for (auto val : perms) {
       permValues.push_back(val.getSExtValue());
     }
 
     outputShape.reserve(inputTy.getRank());
     for (int i = 0, s = inputTy.getRank(); i < s; i++) {
       outputShape[i] = inputTy.getDimSize(permValues[i]);
     }
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::GatherOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape;
-  outputShape.resize(3, -1);
+  outputShape.resize(3, ShapedType::kDynamicSize);
 
   if (auto ty = operands[0].getType().dyn_cast<RankedTensorType>()) {
     outputShape[0] = ty.getDimSize(0);
     outputShape[2] = ty.getDimSize(2);
   }
 
   if (auto ty = operands[1].getType().dyn_cast<RankedTensorType>()) {
-    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);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
+LogicalResult tosa::ResizeOp::inferReturnTypeComponents(
+    MLIRContext *context, ::llvm::Optional<Location> location,
+    ValueRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  llvm::SmallVector<int64_t, 4> outputShape;
+  outputShape.resize(4, ShapedType::kDynamicSize);
+
+  int32_t inHeight = ShapedType::kDynamicSize;
+  int32_t inWidth = ShapedType::kDynamicSize;
+
+  if (auto ty = operands[0].getType().dyn_cast<RankedTensorType>()) {
+    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<IntegerAttr>().getValue().getSExtValue();
+  llvm::SmallVector<int64_t> newShape;
+  getI64Values(attributes.get("output_size").cast<ArrayAttr>(), newShape);
+  outputShape[1] = newShape[0];
+  outputShape[2] = newShape[1];
+
+  llvm::SmallVector<int64_t> strideInt;
+  llvm::SmallVector<int64_t> offsetInt;
+  llvm::SmallVector<double> strideFp;
+  llvm::SmallVector<double> offsetFp;
+  getI64Values(attributes.get("offset").cast<ArrayAttr>(), offsetInt);
+  getF64Values(attributes.get("offset_fp").cast<ArrayAttr>(), offsetFp);
+  getI64Values(attributes.get("stride").cast<ArrayAttr>(), strideInt);
+  getF64Values(attributes.get("stride_fp").cast<ArrayAttr>(), 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> location,
     ValueRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape;
-  outputShape.resize(3, -1);
+  outputShape.resize(3, ShapedType::kDynamicSize);
 
   if (auto ty = operands[0].getType().dyn_cast<RankedTensorType>()) {
     outputShape[0] = ty.getDimSize(0);
     outputShape[1] = ty.getDimSize(1);
     outputShape[2] = ty.getDimSize(2);
   }
 
   if (auto ty = operands[1].getType().dyn_cast<RankedTensorType>()) {
-    if (outputShape[0] == -1)
+    if (outputShape[0] == ShapedType::kDynamicSize)
       outputShape[0] = ty.getDimSize(0);
   }
 
   if (auto ty = operands[2].getType().dyn_cast<RankedTensorType>()) {
-    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);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 static LogicalResult ReduceInferReturnTypes(

 
 NARY_SHAPE_INFER(tosa::AbsOp)
 NARY_SHAPE_INFER(tosa::AddOp)
 NARY_SHAPE_INFER(tosa::ArithmeticRightShiftOp)
 NARY_SHAPE_INFER(tosa::BitwiseAndOp)
 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)
 NARY_SHAPE_INFER(tosa::DivOp)
 NARY_SHAPE_INFER(tosa::EqualOp)
 NARY_SHAPE_INFER(tosa::ExpOp)
 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)
 NARY_SHAPE_INFER(tosa::LogicalNotOp)
 NARY_SHAPE_INFER(tosa::LogicalOrOp)
 NARY_SHAPE_INFER(tosa::LogicalRightShiftOp)
 NARY_SHAPE_INFER(tosa::LogicalXorOp)
 NARY_SHAPE_INFER(tosa::MaximumOp)
 NARY_SHAPE_INFER(tosa::MinimumOp)
 NARY_SHAPE_INFER(tosa::MulOp)
 NARY_SHAPE_INFER(tosa::NegateOp)
 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)
 NARY_SHAPE_INFER(tosa::SubOp)
 NARY_SHAPE_INFER(tosa::TanhOp)
 NARY_SHAPE_INFER(tosa::SigmoidOp)
 #undef PRED_SHAPE_INFER
 
 //===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

func @test_unary_f32(%arg0 : tensor<4xf32>) -> () {
   // CHECK: "tosa.rsqrt"(%arg0) : (tensor<4xf32>) -> tensor<4xf32>
   %10 = "tosa.rsqrt"(%arg0) : (tensor<4xf32>) -> tensor<*xf32>
 
   // CHECK: "tosa.tanh"(%arg0) : (tensor<4xf32>) -> tensor<4xf32>
   %11 = "tosa.tanh"(%arg0) : (tensor<4xf32>) -> tensor<*xf32>
 
   // 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
 }
 
 // -----
 
 // CHECK-LABEL: @test_unary_i32
 func @test_unary_i32(%arg0 : tensor<4xi32>) -> () {
   // CHECK: "tosa.abs"(%arg0) : (tensor<4xi32>) -> tensor<4xi32>
func @test_unary_i32(%arg0 : tensor<4xi32>) -> () {
   // CHECK: "tosa.negate"(%arg0) : (tensor<4xi32>) -> tensor<4xi32>
   %4 = "tosa.negate"(%arg0) : (tensor<4xi32>) -> tensor<*xi32>
 
   // CHECK: "tosa.reluN"(%arg0) {{.+}} : (tensor<4xi32>) -> tensor<4xi32>
   %5 = "tosa.reluN"(%arg0) { max_int = 10 : i64, min_int = 0 : i64, min_fp = 0.0 : f32, max_fp = 10.0 : f32 } : (tensor<4xi32>) -> tensor<*xi32>
 
   // CHECK: "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<4xi32>) -> tensor<4xi32>
   %6 = "tosa.reverse"(%arg0) { axis = 0 : i64 } : (tensor<4xi32>) -> tensor<?xi32>
+
+  // 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<?xi32>
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_unary_i1
 func @test_unary_i1(%arg0 : tensor<4xi1>) -> () {
   // CHECK: "tosa.logical_not"(%arg0) : (tensor<4xi1>) -> tensor<4xi1>

 // -----
 
 // CHECK-LABEL: @scatter_minimum_static
 func @scatter_minimum_static(%arg0 : tensor<?x4x?xi32>, %arg1 : tensor<3x?xi32>, %arg2 : tensor<?x?x5xi32>) {
   // CHECK: "tosa.scatter"(%arg0, %arg1, %arg2) : (tensor<?x4x?xi32>, tensor<3x?xi32>, tensor<?x?x5xi32>) -> tensor<3x4x5xi32>
   %0 = "tosa.scatter"(%arg0, %arg1, %arg2) : (tensor<?x4x?xi32>, tensor<3x?xi32>, tensor<?x?x5xi32>)  -> (tensor<?x?x?xi32>)
   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<?x?x?x?xi32>
+  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<?x?x?x?xi32>
+  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<?x?x?x?xi32>
+  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<?x?x?x?xi32>
+  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<?x?x?x?xi32>
+  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<?x?x?x?xi32>
+  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<?x?x?x?xi32>
+  return
+}