[mlir][tosa] Add type checking traits to the appropriate ops

[tatwaichong]

Summary

Add the trait SameOperandsAndResultElementType and
SameOperandsElementType to verify ops that are known
to have the same input and output type rather than generate
an invalid tosa IR with mixed data types like:

"tosa.add"(%0, %1) : (tensor<nxbf16>, tensor<nxf32>) -> tensor<nxf32>

Thus apply tosa.cast prior if needed.

Change-Id: Ie866b84e371e3b571ec04f7abb090c216dd39c33

[mgehre-amd]

Thanks for fixing this!

I think clamp, sigmoid and tanh are also Tosa_ElemWiseUnaryOp.
Other ops like identity, reshape, concat, pad, reverse, slice, tile and transpose should also have same element types on input and output.

[tatwaichong]

Declaring clamp, sigmoid and tanh as Tosa_ElemWiseUnaryOp.

[tatwaichong]

In D150472#4341193, @mgehre-amd wrote:

Thanks for fixing this!

I think clamp, sigmoid and tanh are also Tosa_ElemWiseUnaryOp.
Other ops like identity, reshape, concat, pad, reverse, slice, tile and transpose should also have same element types on input and output.

Good spot! Yes, these activation functions are element-wise operations essentially. Add them into the group as suggested.
For other non element-wise operations ops like pad, transpose, etc, there is a representation difference between tosa definition in mlir dialect and tosa spec, e.g. padding is operand type in mlir, whereas it is an attribute type in tosa spec. That way we have more chance let tosa operation to support dynamic shape. However, as a result, mlir built-in type checking trait treat those arguments as operand but not attribute so end up happening element types mismatch as the attribute type can differ from the input tensor's.

Can we cover element-wise op in this PR for now? Let's think about how to implement the checking for other in a sperate patch?

[mgehre-amd]

In D150472#4343973, @tatwaichong wrote:

In D150472#4341193, @mgehre-amd wrote:

Thanks for fixing this!

I think clamp, sigmoid and tanh are also Tosa_ElemWiseUnaryOp.
Other ops like identity, reshape, concat, pad, reverse, slice, tile and transpose should also have same element types on input and output.

Good spot! Yes, these activation functions are element-wise operations essentially. Add them into the group as suggested.
For other non element-wise operations ops like pad, transpose, etc, there is a representation difference between tosa definition in mlir dialect and tosa spec, e.g. padding is operand type in mlir, whereas it is an attribute type in tosa spec. That way we have more chance let tosa operation to support dynamic shape. However, as a result, mlir built-in type checking trait treat those arguments as operand but not attribute so end up happening element types mismatch as the attribute type can differ from the input tensor's.

Can we cover element-wise op in this PR for now? Let's think about how to implement the checking for other in a sperate patch?

Sounds good!

[tatwaichong]

Fix the typo of Binary to Unary for clamp, sigmoid and tanh.

[tatwaichong]

put tosa.mul into element-wise binary group.

[mgehre-amd]

With this PR, "tosa.mul"(%0, %1) {shift = 0 : i32} : (tensor<4xi8>, tensor<4xi8>) -> tensor<4xi32> is marked as illegal because result and operand types don't match, even though i8xi8 = i32 is according to the spec: https://www.mlplatform.org/tosa/tosa_spec.html#_mul

[tatwaichong]

@mgehre-amd, oh, you're right. I'll send a PR to fix this soon.

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

 //===----------------------------------------------------------------------===//
 // TOSA Operator.
 //===----------------------------------------------------------------------===//
 
 class Tosa_Op<string mnemonic, list<Trait> traits = []> :
     Op<Tosa_Dialect, mnemonic, !listconcat(traits, [TosaOpInterface])> {
 }
 
+class Tosa_ElemWiseUnaryOp<string mnemonic, list<Trait> traits = []> :
+    Tosa_Op<mnemonic, !listconcat(traits, [
+              DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                                        ["inferReturnTypeComponents"]>,
+              Pure, SameOperandsAndResultElementType])> {
+}
+
+class Tosa_ElemWiseBinaryOp<string mnemonic, list<Trait> traits = []> :
+    Tosa_Op<mnemonic, !listconcat(traits, [
+              DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                                        ["inferReturnTypeComponents"]>,
+              ResultsBroadcastableShape, Pure, SameOperandsAndResultElementType])> {
+}
+
 #endif // TOSA_OP_BASE

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

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.3
 // Operator Class: Activation Functions.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: clamp
 //===----------------------------------------------------------------------===//
-def Tosa_ClampOp : Tosa_Op<"clamp", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_ClampOp : Tosa_ElemWiseUnaryOp<"clamp"> {
   let summary = "Computes clamp(features, min, max).";
 
   let description = [{
     Clamp to an arbitrary minimum and maximum value.
     Maximum and minimum values are specified as values in the range of the
     input type.
     No zero point subtraction is done to the values, thus to clamp to the zero
     point value, the zero point itself should be supplied as the minimum value.
def Tosa_ClampOp : Tosa_ElemWiseUnaryOp<"clamp"> {
   );
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: sigmoid
 //===----------------------------------------------------------------------===//
-def Tosa_SigmoidOp : Tosa_Op<"sigmoid", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_SigmoidOp : Tosa_ElemWiseUnaryOp<"sigmoid"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_TanhOp : Tosa_ElemWiseUnaryOp<"tanh"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure, Commutative]> {
+def Tosa_AddOp : Tosa_ElemWiseBinaryOp<"add", [Commutative]> {
   let summary = "Elementwise addition operator";
 
   let description = [{
     Elementwise addition of input1 and input2. Axis of size 1 will be broadcast,
     as necessary.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: arithmetic_right_shift
 //===----------------------------------------------------------------------===//
-def Tosa_ArithmeticRightShiftOp : Tosa_Op<"arithmetic_right_shift", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+def Tosa_ArithmeticRightShiftOp : Tosa_ElemWiseBinaryOp<"arithmetic_right_shift"> {
   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.
   }];
 
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure, Commutative]> {
+def Tosa_BitwiseAndOp : Tosa_ElemWiseBinaryOp<"bitwise_and", [Commutative]> {
   let summary = "Bitwise AND operator";
 
   let description = [{
     Elementwise bitwise AND of input1 and input2. Axis of size 1
     will be broadcast as necessary.
   }];
 
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure, Commutative]> {
+def Tosa_BitwiseOrOp : Tosa_ElemWiseBinaryOp<"bitwise_or", [Commutative]> {
   let summary = "Bitwise OR operator";
 
   let description = [{
     Elementwise bitwise OR of input1 and input2. Axis of size 1 will be
     broadcast as necessary.
   }];
 
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure, Commutative]> {
+def Tosa_BitwiseXorOp : Tosa_ElemWiseBinaryOp<"bitwise_xor", [Commutative]> {
   let summary = "Bitwise XOR operator";
 
   let description = [{
     Elementwise bitwise XOR of input1 and input2. Axis of size 1 will be
     broadcast as necessary.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: div
 //===----------------------------------------------------------------------===//
-def Tosa_DivOp : Tosa_Op<"div", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+def Tosa_DivOp : Tosa_ElemWiseBinaryOp<"div"> {
   let summary = "Integer divide operator";
 
   let description = [{
     Elementwise integer divide operator of input1 by input2. Axis of size 1
     will be broadcast, as necessary.
   }];
 
   let arguments = (ins
     Tosa_Int32Tensor:$input1,
     Tosa_Int32Tensor:$input2
   );
 
   let results = (outs
     Tosa_Int32Tensor:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_and
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalAndOp : Tosa_Op<"logical_and", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Commutative, Pure]> {
+def Tosa_LogicalAndOp : Tosa_ElemWiseBinaryOp<"logical_and", [Commutative]> {
   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.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,
     I1Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_left_shift
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalLeftShiftOp : Tosa_Op<"logical_left_shift", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+def Tosa_LogicalLeftShiftOp : Tosa_ElemWiseBinaryOp<"logical_left_shift"> {
   let summary = "Elementwise Logical Left Shift";
 
   let description = [{
     Elementwise left shift of input1 and input2. Axis of size 1 will be
     broadcast, as necessary.
   }];
 
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+def Tosa_LogicalRightShiftOp : Tosa_ElemWiseBinaryOp<"logical_right_shift"> {
   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.
   }];
 
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Commutative, Pure]> {
+def Tosa_LogicalOrOp : Tosa_ElemWiseBinaryOp<"logical_or", [Commutative]> {
   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.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,
     I1Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_xor
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalXorOp : Tosa_Op<"logical_xor", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Commutative, Pure]> {
+def Tosa_LogicalXorOp : Tosa_ElemWiseBinaryOp<"logical_xor", [Commutative]> {
   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.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,
     I1Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: maximum
 //===----------------------------------------------------------------------===//
-def Tosa_MaximumOp : Tosa_Op<"maximum", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure, Commutative]> {
+def Tosa_MaximumOp : Tosa_ElemWiseBinaryOp<"maximum", [Commutative]> {
   let summary = "Elementwise Maximum";
 
   let description = [{
     Elementwise max of input1 and input2. Axis of size 1 will be broadcast, as
     necessary.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: minimum
 //===----------------------------------------------------------------------===//
-def Tosa_MinimumOp : Tosa_Op<"minimum", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure, Commutative]> {
+def Tosa_MinimumOp : Tosa_ElemWiseBinaryOp<"minimum", [Commutative]> {
   let summary = "Elementwise Minimum";
 
   let description = [{
     Elementwise minimum of input1 and input2. Axis of size 1
     will be broadcast, as necessary.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: mul
 //===----------------------------------------------------------------------===//
-def Tosa_MulOp : Tosa_Op<"mul", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure, Commutative]> {
+def Tosa_MulOp : Tosa_ElemWiseBinaryOp<"mul", [Commutative]> {
   let summary = "Multiplication operator";
 
   let description = [{
     Elementwise multiplication (Hadamard product) of input1 and input2.
     Axis of size 1 will be broadcast, as necessary.
+    i8/i16 input type can be promoted to i32 result type.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2,
     I32Attr:$shift
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pow
 //===----------------------------------------------------------------------===//
-def Tosa_PowOp : Tosa_Op<"pow", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+def Tosa_PowOp : Tosa_ElemWiseBinaryOp<"pow"> {
   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.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     Tosa_Tensor:$input2
   );
 
   let results = (outs
     Tosa_Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: sub
 //===----------------------------------------------------------------------===//
-def Tosa_SubOp : Tosa_Op<"sub", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+def Tosa_SubOp : Tosa_ElemWiseBinaryOp<"sub"> {
   let summary = "Elementwise subtraction operator";
 
   let description = [{
     Elementwise subtraction of input1 and input2. Axis of size 1 will be
     broadcast as necessary.
   }];
 
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_AbsOp : Tosa_ElemWiseUnaryOp<"abs"> {
   let summary = "Elementwise abs op";
 
   let description = [{
     Elementwise absolute value operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_not
 //===----------------------------------------------------------------------===//
-def Tosa_BitwiseNotOp : Tosa_Op<"bitwise_not", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+def Tosa_BitwiseNotOp : Tosa_ElemWiseUnaryOp<"bitwise_not"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_CeilOp : Tosa_ElemWiseUnaryOp<"ceil"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_ClzOp : Tosa_ElemWiseUnaryOp<"clz"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_ExpOp : Tosa_ElemWiseUnaryOp<"exp"> {
   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
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: floor
 //===----------------------------------------------------------------------===//
-def Tosa_FloorOp : Tosa_Op<"floor", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_FloorOp : Tosa_ElemWiseUnaryOp<"floor"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_LogOp : Tosa_ElemWiseUnaryOp<"log"> {
   let summary = "Elementwise log op";
 
   let description = [{
     Elementwise natural logarithm operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_not
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalNotOp : Tosa_Op<"logical_not", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure, SameOperandsAndResultType]> {
+def Tosa_LogicalNotOp : Tosa_ElemWiseUnaryOp<"logical_not"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_NegateOp : Tosa_ElemWiseUnaryOp<"negate"> {
   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];
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reciprocal
 //===----------------------------------------------------------------------===//
-def Tosa_ReciprocalOp : Tosa_Op<"reciprocal", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_ReciprocalOp : Tosa_ElemWiseUnaryOp<"reciprocal"> {
   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", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_RsqrtOp : Tosa_ElemWiseUnaryOp<"rsqrt"> {
   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.7
 // Operator Class: Logical Operations.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: equal
 //===----------------------------------------------------------------------===//
 def Tosa_EqualOp : Tosa_Op<"equal", [InferTensorType, ResultsBroadcastableShape,
-    Commutative, Pure]> {
+    Commutative, Pure, SameOperandsElementType]> {
   let summary = "Returns the truth value of (x == y) element-wise.";
 
   let description = [{
      Elementwise comparison operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,

 }
 
 //===----------------------------------------------------------------------===//
 // Operator: greater
 //===----------------------------------------------------------------------===//
 def Tosa_GreaterOp : Tosa_Op<"greater", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+    ResultsBroadcastableShape, Pure, SameOperandsElementType]> {
   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
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: greater_equal
 //===----------------------------------------------------------------------===//
 def Tosa_GreaterEqualOp : Tosa_Op<"greater_equal", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    ResultsBroadcastableShape, Pure]> {
+    ResultsBroadcastableShape, Pure, SameOperandsElementType]> {
   let summary = "Returns the truth value of (x >= y) element-wise.";
 
   let description = [{
     Elementwise comparison operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,

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

func.func @test_simple_f16(%arg0: tensor<1xf16>) -> () {
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_simple_i16
 func.func @test_simple_i16(%arg0: tensor<1xi16>) -> () {
   // CHECK: linalg.generic
-  // CHECK: arith.extsi
-  // CHECK: arith.extsi
   // CHECK: arith.muli
-  %0 = "tosa.mul"(%arg0, %arg0) {shift = 0 : i32} : (tensor<1xi16>, tensor<1xi16>) -> tensor<1xi32>
+  %0 = "tosa.mul"(%arg0, %arg0) {shift = 0 : i32} : (tensor<1xi16>, tensor<1xi16>) -> tensor<1xi16>
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_simple_ui8
 func.func @test_simple_ui8(%arg0: tensor<1xui8>) -> () {