[mlir][tosa] Add InferReturnTypes for N-ary operations

[rsuderman]

Summary

InferReturnTypes enables return type inference based on input types. This
allows op validation of any tosa op that supports broadcasting along with
non-broadcastable TOSA operations.

[rsuderman]

Going to change from inferring the full return type to only inferring the shape. This encounters issues during dynamic / underspecified cases (verifiers fail) which makes shape inference required more than it should be.

[jpienaar]

You probably want to also set the compatible testing, that allows for the shape to be more or less exact and still not be an error. The default is the strictest (inferred type has to equal inferred - I had considered relaxing that with the tensor type helper class.

If this is unclear comment, let me know, quickly typed on mobile

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 #define MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/Traits.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // TOSA dialect and structs includes.
 //===----------------------------------------------------------------------===//
 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.h.inc"
 #include "mlir/Dialect/Tosa/IR/TosaStructs.h.inc"

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

 //===----------------------------------------------------------------------===//
 
 #ifndef TOSA_OPS
 #define TOSA_OPS
 
 include "mlir/IR/OpBase.td"
 
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Dialect/Tosa/IR/TosaInterfaces.td"
 
 include "mlir/Dialect/Tosa/IR/TosaTypesBase.td"
 include "mlir/Dialect/Tosa/IR/TosaOpBase.td"
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.2

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.3
 // Operator Class: Activation Functions.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: clamp
 //===----------------------------------------------------------------------===//
-def Tosa_ClampOp : Tosa_Op<"clamp", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_ClampOp : Tosa_Op<"clamp", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
+    NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Computes clamp(features, min, max).";
 
   let description = [{
     Clamp to an arbitrary minimum and maximum value. Note that the maximum and
     minimum values are specified as signed quantized values, no scaling happens
     before or after this operation.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
     I64Attr:$min_int,
     I64Attr:$max_int,
     F32Attr:$min_fp,
     F32Attr:$max_fp
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reluN
 //===----------------------------------------------------------------------===//
-def Tosa_ReluNOp : Tosa_Op<"reluN", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_ReluNOp : Tosa_Op<"reluN", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
+    NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Computes rectified linear: `max(features, N)`.";
 
   let description = [{
      ReLU with a scalar maximum value.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
     I64Attr:$max_int,
     F32Attr:$max_fp
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: sigmoid
 //===----------------------------------------------------------------------===//
-def Tosa_SigmoidOp : Tosa_Op<"sigmoid", [NoSideEffect,
-                             SameOperandsAndResultType]> {
+def Tosa_SigmoidOp : Tosa_Op<"sigmoid", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
   let summary = "Computes elementwise sigmoid of input.";
 
   let description = [{
     Sigmoid function: output = 1 / (1 + exp(-input))
     For quantized integer data types, the TABLE operator should be used instead
     with the following definition.  The sigmoid table has 513 entries each of
     16-bit precision and covering the input range -16.0 to +16.0
     in steps of 1/16.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: tanh
 //===----------------------------------------------------------------------===//
-def Tosa_TanhOp : Tosa_Op<"tanh", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_TanhOp : Tosa_Op<"tanh", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
+    NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Computes elementwise hyperbolic tangent of input";
 
   let description = [{
     Parameterized hyperbolic tangent.
     For quantized integer data types, the TABLE operator should be used instead
     with the following definition.  The tanh_table has 513 entries each of
     16-bit precision and covering the input range -8.0 to +8.0 in steps of 1/32.
   }];

 // TOSA Spec Section 2.4
 // Operator Class: Elementwise unary/binary/ternary operators.
 // Operator Subclass: Elementwise binary ops.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: add
 //===----------------------------------------------------------------------===//
-def Tosa_AddOp : Tosa_Op<"add", [ResultsBroadcastableShape, NoSideEffect,
-                         Commutative]> {
+def Tosa_AddOp : Tosa_Op<"add", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
+    ResultsBroadcastableShape, NoSideEffect, Commutative]> {
   let summary = "Elementwise addition operator";
 
   let description = [{
     Elementwise addition of input1 and input2. Axis of size 1 will be broadcast,
     as necessary.  Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: arithmetic_right_shift
 //===----------------------------------------------------------------------===//
-def Tosa_ArithmeticRightShiftOp : Tosa_Op<"arithmetic_right_shift",
-                                          [ResultsBroadcastableShape,
-                                          NoSideEffect]> {
+def Tosa_ArithmeticRightShiftOp : Tosa_Op<"arithmetic_right_shift", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Elementwise Arithmetic Right Shift";
 
   let description = [{
     Elementwise arithmetic right shift of input1 by the amount specified in
     input2. Axis of size 1 will be broadcast, as necessary.  Rank of input
     tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2,
     BoolAttr:$round
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_and
 //===----------------------------------------------------------------------===//
-def Tosa_BitwiseAndOp : Tosa_Op<"bitwise_and", [ResultsBroadcastableShape,
-                                NoSideEffect, Commutative]> {
+def Tosa_BitwiseAndOp : Tosa_Op<"bitwise_and", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Bitwise AND operator";
 
   let description = [{
     Elementwise bitwise AND of input1 and input2. Axis of size 1
     will be broadcast as necessary. Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_or
 //===----------------------------------------------------------------------===//
-def Tosa_BitwiseOrOp : Tosa_Op<"bitwise_or", [ResultsBroadcastableShape,
-                                              NoSideEffect, Commutative]> {
+def Tosa_BitwiseOrOp : Tosa_Op<"bitwise_or", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Bitwise OR operator";
 
   let description = [{
     Elementwise bitwise OR of input1 and input2. Axis of size 1 will be
     broadcast as necessary. Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_xor
 //===----------------------------------------------------------------------===//
-def Tosa_BitwiseXorOp : Tosa_Op<"bitwise_xor", [ResultsBroadcastableShape,
-                                                NoSideEffect, Commutative]> {
+def Tosa_BitwiseXorOp : Tosa_Op<"bitwise_xor", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Bitwise XOR operator";
 
   let description = [{
     Elementwise bitwise XOR of input1 and input2. Axis of size 1 will be
     broadcast as necessary. Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: div
 //===----------------------------------------------------------------------===//
-def Tosa_DivOp : Tosa_Op<"div", [ResultsBroadcastableShape,
-                                 NoSideEffect]> {
+def Tosa_DivOp : Tosa_Op<"div", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Integer divide operator";
 
   let description = [{
     Elementwise integer divide operator of input1 by input2. Axis of size 1
     will be broadcast, as necessary. Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Int32Tensor:$input1,
     Tosa_Int32Tensor:$input2
   );
 
   let results = (outs
     Tosa_Int32Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_and
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalAndOp : Tosa_Op<"logical_and", [ResultsBroadcastableShape,
-                                                Commutative, NoSideEffect]> {
+def Tosa_LogicalAndOp : Tosa_Op<"logical_and", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Returns the truth value of x AND y element-wise.";
 
   let description = [{
     Elementwise logical AND of input1 and input2. Axis of size 1 will be
     broadcast, as necessary.  Rank of input tensors must match.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,
     I1Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_left_shift
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalLeftShiftOp : Tosa_Op<"logical_left_shift",
-                                      [ResultsBroadcastableShape,
-                                       NoSideEffect]> {
+def Tosa_LogicalLeftShiftOp : Tosa_Op<"logical_left_shift", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Elementwise Logical Left Shift";
 
   let description = [{
     Elementwise left shift of input1 and input2. Axis of size 1 will be
     broadcast, as necessary. Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_right_shift
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalRightShiftOp : Tosa_Op<"logical_right_shift",
-                                       [ResultsBroadcastableShape,
-                                        NoSideEffect]> {
+def Tosa_LogicalRightShiftOp : Tosa_Op<"logical_right_shift", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Elementwise Logical Right Shift";
 
   let description = [{
     Elementwise logical right shift of input1 by the amount specified in input2.
     Axis of size 1 will be broadcast, as necessary.
     Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_or
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalOrOp : Tosa_Op<"logical_or", [ResultsBroadcastableShape,
-                                              Commutative, NoSideEffect]> {
+def Tosa_LogicalOrOp : Tosa_Op<"logical_or", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Returns the truth value of x OR y element-wise.";
 
   let description = [{
     Elementwise logical OR of input1 and input2. Axis of size 1 will be
     broadcast as necessary. Rank of input tensors must match.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,
     I1Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_xor
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalXorOp : Tosa_Op<"logical_xor", [ResultsBroadcastableShape,
-                                                Commutative, NoSideEffect]> {
+def Tosa_LogicalXorOp : Tosa_Op<"logical_xor", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Returns the truth value of x XOR y element-wise.";
 
   let description = [{
     Elementwise logical XOR of input1 and input2.  Axis of size 1 will be
     broadcast as necessary. Rank of input tensors must match.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,
     I1Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: maximum
 //===----------------------------------------------------------------------===//
-def Tosa_MaximumOp : Tosa_Op<"maximum", [ResultsBroadcastableShape,
-                                         NoSideEffect, Commutative]> {
+def Tosa_MaximumOp : Tosa_Op<"maximum", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Elementwise Maximum";
 
   let description = [{
     Elementwise max of input1 and input2. Axis of size 1 will be broadcast, as
     necessary.  Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: minimum
 //===----------------------------------------------------------------------===//
-def Tosa_MinimumOp : Tosa_Op<"minimum", [ResultsBroadcastableShape,
-                                         NoSideEffect, Commutative]> {
+def Tosa_MinimumOp : Tosa_Op<"minimum", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Elementwise Minimum";
 
   let description = [{
     Elementwise minimum of input1 and input2. Axis of size 1
     will be broadcast, as necessary.  Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: mul
 //===----------------------------------------------------------------------===//
-def Tosa_MulOp : Tosa_Op<"mul", [ResultsBroadcastableShape, NoSideEffect,
-                                 Commutative]> {
+def Tosa_MulOp : Tosa_Op<"mul", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
+    NoSideEffect, Commutative]> {
   let summary = "Multiplication operator";
 
   let description = [{
     Elementwise multiplication (Hadamard product) of input1 and input2.
     Axis of size 1 will be broadcast, as necessary.
     Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2,
     I32Attr:$shift
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pow
 //===----------------------------------------------------------------------===//
-def Tosa_PowOp : Tosa_Op<"pow", [ResultsBroadcastableShape, NoSideEffect]> {
+def Tosa_PowOp : Tosa_Op<"pow", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
+    NoSideEffect]> {
   let summary = "Computes the power of one value to another.";
 
   let description = [{
     Elementwise input1 raised to the power of input2.
     Axis of size 1 will be broadcast, as necessary.
     Rank of input tensors must match.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: sub
 //===----------------------------------------------------------------------===//
-def Tosa_SubOp : Tosa_Op<"sub", [ResultsBroadcastableShape, NoSideEffect]> {
+def Tosa_SubOp : Tosa_Op<"sub", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
+    NoSideEffect]> {
   let summary = "Elementwise subtraction operator";
 
   let description = [{
     Elementwise subtraction of input1 and input2. Axis of size 1 will be
     broadcast as necessary.  Rank of input tensors must match.
   }];
 
   let arguments = (ins

 // TOSA Spec Section 2.5
 // Operator Class: Elementwise unary/binary/ternary operators.
 // Operator Subclass: Elementwise unary ops.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: abs
 //===----------------------------------------------------------------------===//
-def Tosa_AbsOp : Tosa_Op<"abs", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_AbsOp : Tosa_Op<"abs", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
+    SameOperandsAndResultType]> {
   let summary = "Elementwise abs op";
 
   let description = [{
     Elementwise absolute value operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_not
 //===----------------------------------------------------------------------===//
-def Tosa_BitwiseNotOp : Tosa_Op<"bitwise_not", [ResultsBroadcastableShape,
-                                                NoSideEffect]> {
+def Tosa_BitwiseNotOp : Tosa_Op<"bitwise_not", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Bitwise NOT operator";
 
   let description = [{
     Elementwise bitwise NOT of input tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: ceil
 //===----------------------------------------------------------------------===//
-def Tosa_CeilOp : Tosa_Op<"ceil", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_CeilOp : Tosa_Op<"ceil", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
+    SameOperandsAndResultType]> {
   let summary = "Elementwise ceil op";
 
   let description = [{
     Elementwise ceiling operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: clz
 //===----------------------------------------------------------------------===//
-def Tosa_ClzOp : Tosa_Op<"clz", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_ClzOp : Tosa_Op<"clz", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
+    SameOperandsAndResultType]> {
   let summary = "Elementwise count leading zero op";
 
   let description = [{
     Elementwise count leading zeros operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: exp
 //===----------------------------------------------------------------------===//
-def Tosa_ExpOp : Tosa_Op<"exp", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_ExpOp : Tosa_Op<"exp", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
+    SameOperandsAndResultType]> {
   let summary = "Elementwise exp op";
 
   let description = [{
     Elementwise e to the x operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: floor
 //===----------------------------------------------------------------------===//
-def Tosa_FloorOp : Tosa_Op<"floor", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_FloorOp : Tosa_Op<"floor", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
+    SameOperandsAndResultType]> {
   let summary = "Elementwise floor op";
 
   let description = [{
     Elementwise floor operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: log
 //===----------------------------------------------------------------------===//
-def Tosa_LogOp : Tosa_Op<"log", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_LogOp : Tosa_Op<"log", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
+    SameOperandsAndResultType]> {
   let summary = "Elementwise log op";
 
   let description = [{
     Elementwise natural logarithm operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_not
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalNotOp : Tosa_Op<"logical_not", [NoSideEffect,
-                                                SameOperandsAndResultType]> {
+def Tosa_LogicalNotOp : Tosa_Op<"logical_not", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
   let summary = "Returns the truth value of NOT x element-wise.";
 
   let description = [{
     Elementwise logical NOT of input.
   }];
 
   let arguments = (ins
     I1Tensor:$input1
   );
 
   let results = (outs
     I1Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: negate
 //===----------------------------------------------------------------------===//
-def Tosa_NegateOp : Tosa_Op<"negate", [NoSideEffect,
-                                       SameOperandsAndResultType]> {
+def Tosa_NegateOp : Tosa_Op<"negate", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
   let summary = "Elementwise negate op";
 
   let description = [{
     Elementwise negation operation
   }];
 
   let arguments = (ins
       Tosa_Tensor:$input1,
       OptionalAttr<Tosa_UnaryOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
   let builders = [Tosa_UnaryOpQuantInfoBuilder];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reciprocal
 //===----------------------------------------------------------------------===//
-def Tosa_ReciprocalOp : Tosa_Op<"reciprocal", [NoSideEffect,
-                                               SameOperandsAndResultType]> {
+def Tosa_ReciprocalOp : Tosa_Op<"reciprocal", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
   let summary = "Elementwise reciprocal op";
 
   let description = [{
     Elementwise reciprocal operation. For integer operation, a TABLE should be
     used with the appropriate ranges.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: rsqrt
 //===----------------------------------------------------------------------===//
-def Tosa_RsqrtOp : Tosa_Op<"rsqrt", [NoSideEffect, SameOperandsAndResultType]> {
+def Tosa_RsqrtOp : Tosa_Op<"rsqrt", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect,
+    SameOperandsAndResultType]> {
   let summary = "Elementwise 1/sqrt op";
 
   let description = [{
     Elementwise reciprocal square root operation. For integer operation, a TABLE
     should be used with the appropriate ranges.
   }];
 
   let arguments = (ins

 // TOSA Spec Section 2.6
 // Operator Class: Elementwise unary/binary/ternary operators.
 // Operator Subclass: Elementwise ternary ops.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: select
 //===----------------------------------------------------------------------===//
-def Tosa_SelectOp : Tosa_Op<"select", [NoSideEffect]> {
+def Tosa_SelectOp : Tosa_Op<"select", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect]> {
   let summary = "Elementwise select operator";
 
   let description = [{
     Elementwise select of the output based on a condition.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.7
 // Operator Class: Logical Operations.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: equal
 //===----------------------------------------------------------------------===//
-def Tosa_EqualOp : Tosa_Op<"equal", [ResultsBroadcastableShape, Commutative,
-                                     NoSideEffect]> {
+def Tosa_EqualOp : Tosa_Op<"equal", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
+    Commutative, NoSideEffect]> {
   let summary = "Returns the truth value of (x == y) element-wise.";
 
   let description = [{
      Elementwise comparison operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: greater
 //===----------------------------------------------------------------------===//
-def Tosa_GreaterOp : Tosa_Op<"greater", [ResultsBroadcastableShape,
-                                         NoSideEffect]> {
+def Tosa_GreaterOp : Tosa_Op<"greater", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Returns the truth value of (x > y) element-wise.";
 
   let description = [{
     Elementwise greater than comparison operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: greater_equal
 //===----------------------------------------------------------------------===//
-def Tosa_GreaterEqualOp : Tosa_Op<"greater_equal", [ResultsBroadcastableShape,
-                                                    NoSideEffect]> {
+def Tosa_GreaterEqualOp : Tosa_Op<"greater_equal", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>, ResultsBroadcastableShape,
   let summary = "Returns the truth value of (x >= y) element-wise.";
 
   let description = [{
     Elementwise comparison operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
def Tosa_PadOp : Tosa_Op<"pad", [NoSideEffect]> {
 
   let builders = [Tosa_PadOpQuantInfoBuilder];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reshape
 //===----------------------------------------------------------------------===//
 def Tosa_ReshapeOp: Tosa_Op<"reshape", [
-   NoSideEffect]> {
+  DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect]> {
   let summary = "Reshape operator";
 
   let description = [{
     Returns a tensor with the same type/values as the input, with a new shape
     specified by the shape argument. Reshape may operate on tensors of any rank.
     No data conversion happens during a reshape operation.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     I64ArrayAttr:$new_shape
   );
 
   let results = (outs
     Tosa_RankedTensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reverse
 //===----------------------------------------------------------------------===//
-def Tosa_ReverseOp: Tosa_Op<"reverse", [NoSideEffect]> {
+def Tosa_ReverseOp: Tosa_Op<"reverse", [
+  DeclareOpInterfaceMethods<InferTypeOpInterface>, NoSideEffect]> {
   let summary = "Reverse operator";
 
   let description = [{
     Returns a tensor with the same type/values as the input, with the data
     reversed along the given axis. No data conversion happens during a reverse
     operation.
   }];
 

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

static void buildPadOpWithQuantInfo(OpBuilder &builder, OperationState &result,
   result.addOperands({input, paddings});
   auto quantAttr = buildPadOpQuantizationAttr(builder, input);
   if (quantAttr)
     result.addAttribute("quantization_info", quantAttr);
   result.types.push_back(outputType);
 }
 
 //===----------------------------------------------------------------------===//
+// TOSA Operator Return Type Inference.
+//===----------------------------------------------------------------------===//
+
+static void getI64Values(ArrayAttr arrayAttr, SmallVector<int64_t> &values) {
+  for (auto it : arrayAttr) {
+    values.push_back(it.cast<IntegerAttr>().getValue().getSExtValue());
+  }
+}
+
+LogicalResult tosa::ReshapeOp::inferReturnTypes(
+    MLIRContext *context, ::llvm::Optional<Location> location,
+    ValueRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  auto newShape = attributes.get("new_shape").cast<ArrayAttr>();
+  auto eTy = operands.front().getType().cast<ShapedType>().getElementType();
+  llvm::SmallVector<int64_t> newShapeValue;
+  getI64Values(newShape, newShapeValue);
+
+  inferredReturnTypes.push_back(RankedTensorType::get(newShapeValue, eTy));
+  return success();
+}
+
+static LogicalResult ResolveBroadcastShape(ValueRange operands,
+                                           SmallVector<int64_t> &outShape) {
+  int64_t outRank = 0;
+  for (auto operand : operands) {
+    auto type = operand.getType().cast<ShapedType>();
+    outRank = std::max<int64_t>(outRank, type.getRank());
+  }
+
+  outShape.resize(outRank, 1);
+
+  for (auto operand : operands) {
+    auto type = operand.getType().cast<ShapedType>();
+    auto shape = type.getShape();
+    auto rankDiff = outShape.size() - shape.size();
+
+    for (size_t i = 0; i < shape.size(); i++) {
+      auto dim1 = outShape[i + rankDiff];
+      auto dim2 = shape[i];
+      auto resolvedDim = dim1;
+
+      if (dim1 == 1) {
+        resolvedDim = dim2;
+      } else if (dim2 == 1) {
+        resolvedDim = dim1;
+      } else if (dim1 != dim2) {
+        return failure();
+      }
+      outShape[i + rankDiff] = resolvedDim;
+    }
+  }
+
+  return success();
+}
+
+static LogicalResult
+NAryInferReturnTypes(ValueRange operands,
+                     SmallVectorImpl<Type> &inferredReturnTypes) {
+  llvm::SmallVector<int64_t> outShape;
+  if (ResolveBroadcastShape(operands, outShape).failed())
+    return failure();
+  inferredReturnTypes.push_back(RankedTensorType::get(
+      outShape,
+      operands.front().getType().cast<ShapedType>().getElementType()));
+  return success();
+}
+
+#define NARY_SHAPE_INFER(OP)                                                   \
+  LogicalResult OP::inferReturnTypes(                                          \
+      MLIRContext *context, ::llvm::Optional<Location> location,               \
+      ValueRange operands, DictionaryAttr attributes, RegionRange regions,     \
+      SmallVectorImpl<Type> &inferredReturnTypes) {                            \
+    return NAryInferReturnTypes(operands, inferredReturnTypes);                \
+  }
+
+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::CeilOp)
+NARY_SHAPE_INFER(tosa::ClampOp)
+NARY_SHAPE_INFER(tosa::ClzOp)
+NARY_SHAPE_INFER(tosa::DivOp)
+NARY_SHAPE_INFER(tosa::ExpOp)
+NARY_SHAPE_INFER(tosa::FloorOp)
+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::NegateOp)
+NARY_SHAPE_INFER(tosa::PowOp)
+NARY_SHAPE_INFER(tosa::ReciprocalOp)
+NARY_SHAPE_INFER(tosa::ReluNOp)
+NARY_SHAPE_INFER(tosa::ReverseOp)
+NARY_SHAPE_INFER(tosa::RsqrtOp)
+NARY_SHAPE_INFER(tosa::SubOp)
+NARY_SHAPE_INFER(tosa::TanhOp)
+NARY_SHAPE_INFER(tosa::SigmoidOp)
+#undef UNARY_SHAPE_INFER
+
+#define PRED_SHAPE_INFER(OP)                                                   \
+  LogicalResult OP::inferReturnTypes(                                          \
+      MLIRContext *context, ::llvm::Optional<Location> location,               \
+      ValueRange operands, DictionaryAttr attributes, RegionRange regions,     \
+      SmallVectorImpl<Type> &inferredReturnTypes) {                            \
+    llvm::SmallVector<int64_t> outShape;                                       \
+    if (ResolveBroadcastShape(operands, outShape).failed())                    \
+      return failure();                                                        \
+    inferredReturnTypes.push_back(                                             \
+        RankedTensorType::get(outShape, IntegerType::get(context, 1)));        \
+    return success();                                                          \
+  }
+PRED_SHAPE_INFER(tosa::EqualOp)
+PRED_SHAPE_INFER(tosa::GreaterOp)
+PRED_SHAPE_INFER(tosa::GreaterEqualOp)
+#undef PRED_SHAPE_INFER
+
+LogicalResult tosa::MulOp::inferReturnTypes(
+    MLIRContext *context, ::llvm::Optional<Location> location,
+    ValueRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  llvm::SmallVector<int64_t> outShape;
+  if (ResolveBroadcastShape(operands, outShape).failed())
+    return failure();
+  auto eTy = operands[1].getType().cast<ShapedType>().getElementType();
+  if (eTy.isa<IntegerType>())
+    eTy = IntegerType::get(context, 32);
+  inferredReturnTypes.push_back(RankedTensorType::get(outShape, eTy));
+  return success();
+}
+
+LogicalResult tosa::SelectOp::inferReturnTypes(
+    MLIRContext *context, ::llvm::Optional<Location> location,
+    ValueRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  llvm::SmallVector<int64_t> outShape;
+  if (ResolveBroadcastShape(operands, outShape).failed())
+    return failure();
+  inferredReturnTypes.push_back(RankedTensorType::get(
+      outShape, operands[1].getType().cast<ShapedType>().getElementType()));
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func @test_reshape_samerank(%arg0: tensor<3x2xf32>) -> tensor<2x3xf32> {
   return %0 : tensor<2x3xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_downrank_6D
 func @test_reshape_downrank_6D(%arg0: tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32> {
   // CHECK: linalg.tensor_collapse_shape %arg0 {{\[}}[0, 1, 2], [3], [4, 5]]
-  %0 = "tosa.reshape"(%arg0) {new_shape = [2, 3]} : (tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32>
+  %0 = "tosa.reshape"(%arg0) {new_shape = [6, 5, 77]} : (tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32>
   return %0 : tensor<6x5x77xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_identity
 func @test_identity(%arg0: tensor<1xf32>, %arg1: tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>) {
   %0 = "tosa.identity"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>