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

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

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D150472

Author: tatwaichong

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

@@ -215,4 +215,18 @@ 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