[TOSA] Introduce Tosa_ElementwiseUnaryOp with Type and Shape Enforcement (#115784)

* Enforce that Tosa_ElementwiseUnaryOp requires output tensors to match
the input tensor's type and shape.
* Update the following ops to conform to Tosa_ElementwiseUnaryOp: clamp,
erf, sigmoid, tanh, cos, sin, abs, bitwise_not, ceil, clz, exp, floor,
log, logical_not, negate, reciprocal, rsqrt.
* Add invalid tests for each operator to ensure compliance with TOSA
v1.0 Specification.

Signed-off-by: Peng Sun <[email protected]>

Author: psunn

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

@@ -234,6 +234,11 @@ class Tosa_ElementwiseOp<string mnemonic, list<Trait> traits = []> :
       "operands attr-dict `:` functional-type(operands, results)";
 }
 
+class Tosa_ElementwiseUnaryOp<string mnemonic, list<Trait> traits = []> :
+    Tosa_ElementwiseOp<mnemonic, !listconcat(traits, [
+              SameOperandsAndResultShape,
+              SameOperandsAndResultElementType])> {}
+
 class Tosa_InferTensorTypeOp<string mnemonic, list<Trait> traits = []>
     : Tosa_Op<mnemonic, !listconcat(traits, [InferTensorTypeAdaptor, Pure])> {
   let assemblyFormat =

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

@@ -367,7 +367,7 @@ def Tosa_TransposeConv2DOp : Tosa_InferShapedTypeOp<"transpose_conv2d"> {
 //===----------------------------------------------------------------------===//
 // Operator: clamp
 //===----------------------------------------------------------------------===//
-def Tosa_ClampOp : Tosa_ElementwiseOp<"clamp"> {
+def Tosa_ClampOp : Tosa_ElementwiseUnaryOp<"clamp"> {
   let summary = "Computes clamp(features, min, max).";
 
   let description = [{
@@ -397,7 +397,7 @@ def Tosa_ClampOp : Tosa_ElementwiseOp<"clamp"> {
 //===----------------------------------------------------------------------===//
 // Operator: sigmoid
 //===----------------------------------------------------------------------===//
-def Tosa_SigmoidOp : Tosa_ElementwiseOp<"sigmoid"> {
+def Tosa_SigmoidOp : Tosa_ElementwiseUnaryOp<"sigmoid"> {
   let summary = "Computes elementwise sigmoid of input.";
 
   let description = [{
@@ -420,7 +420,7 @@ def Tosa_SigmoidOp : Tosa_ElementwiseOp<"sigmoid"> {
 //===----------------------------------------------------------------------===//
 // Operator: tanh
 //===----------------------------------------------------------------------===//
-def Tosa_TanhOp : Tosa_ElementwiseOp<"tanh", [SameOperandsAndResultElementType]> {
+def Tosa_TanhOp : Tosa_ElementwiseUnaryOp<"tanh"> {
   let summary = "Computes elementwise hyperbolic tangent of input";
 
   let description = [{
@@ -442,10 +442,7 @@ def Tosa_TanhOp : Tosa_ElementwiseOp<"tanh", [SameOperandsAndResultElementType]>
 //===----------------------------------------------------------------------===//
 // Operator: erf
 //===----------------------------------------------------------------------===//
-def Tosa_ErfOp : Tosa_Op<"erf", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_ErfOp : Tosa_ElementwiseUnaryOp<"erf"> {
   let summary = "Computes gauss error function of input";
 
   let description = [{
@@ -906,7 +903,7 @@ def Tosa_TableOp : Tosa_InferShapedTypeOp<"table"> {
 //===----------------------------------------------------------------------===//
 // Operator: abs
 //===----------------------------------------------------------------------===//
-def Tosa_AbsOp : Tosa_ElementwiseOp<"abs", [SameOperandsAndResultElementType]> {
+def Tosa_AbsOp : Tosa_ElementwiseUnaryOp<"abs"> {
   let summary = "Elementwise abs op";
 
   let description = [{
@@ -933,8 +930,7 @@ def Tosa_AbsOp : Tosa_ElementwiseOp<"abs", [SameOperandsAndResultElementType]> {
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_not
 //===----------------------------------------------------------------------===//
-def Tosa_BitwiseNotOp : Tosa_ElementwiseOp<"bitwise_not",
-    [SameOperandsAndResultElementType]> {
+def Tosa_BitwiseNotOp : Tosa_ElementwiseUnaryOp<"bitwise_not"> {
   let summary = "Bitwise NOT operator";
 
   let description = [{
@@ -953,7 +949,7 @@ def Tosa_BitwiseNotOp : Tosa_ElementwiseOp<"bitwise_not",
 //===----------------------------------------------------------------------===//
 // Operator: ceil
 //===----------------------------------------------------------------------===//
-def Tosa_CeilOp : Tosa_ElementwiseOp<"ceil", [SameOperandsAndResultElementType]> {
+def Tosa_CeilOp : Tosa_ElementwiseUnaryOp<"ceil"> {
   let summary = "Elementwise ceil op";
 
   let description = [{
@@ -972,7 +968,7 @@ def Tosa_CeilOp : Tosa_ElementwiseOp<"ceil", [SameOperandsAndResultElementType]>
 //===----------------------------------------------------------------------===//
 // Operator: clz
 //===----------------------------------------------------------------------===//
-def Tosa_ClzOp : Tosa_ElementwiseOp<"clz", [SameOperandsAndResultElementType]> {
+def Tosa_ClzOp : Tosa_ElementwiseUnaryOp<"clz"> {
   let summary = "Elementwise count leading zero op";
 
   let description = [{
@@ -991,8 +987,7 @@ def Tosa_ClzOp : Tosa_ElementwiseOp<"clz", [SameOperandsAndResultElementType]> {
 //===----------------------------------------------------------------------===//
 // Operator: cos
 //===----------------------------------------------------------------------===//
-def Tosa_CosOp : Tosa_ElementwiseOp<"cos",
-    [SameOperandsAndResultElementType]> {
+def Tosa_CosOp : Tosa_ElementwiseUnaryOp<"cos"> {
   let summary = "Elementwise cos op";
 
   let description = [{
@@ -1011,7 +1006,7 @@ def Tosa_CosOp : Tosa_ElementwiseOp<"cos",
 //===----------------------------------------------------------------------===//
 // Operator: exp
 //===----------------------------------------------------------------------===//
-def Tosa_ExpOp : Tosa_ElementwiseOp<"exp", [SameOperandsAndResultElementType]> {
+def Tosa_ExpOp : Tosa_ElementwiseUnaryOp<"exp"> {
   let summary = "Elementwise exp op";
 
   let description = [{
@@ -1032,7 +1027,7 @@ def Tosa_ExpOp : Tosa_ElementwiseOp<"exp", [SameOperandsAndResultElementType]> {
 //===----------------------------------------------------------------------===//
 // Operator: floor
 //===----------------------------------------------------------------------===//
-def Tosa_FloorOp : Tosa_ElementwiseOp<"floor", [SameOperandsAndResultElementType]> {
+def Tosa_FloorOp : Tosa_ElementwiseUnaryOp<"floor"> {
   let summary = "Elementwise floor op";
 
   let description = [{
@@ -1051,7 +1046,7 @@ def Tosa_FloorOp : Tosa_ElementwiseOp<"floor", [SameOperandsAndResultElementType
 //===----------------------------------------------------------------------===//
 // Operator: log
 //===----------------------------------------------------------------------===//
-def Tosa_LogOp : Tosa_ElementwiseOp<"log", [SameOperandsAndResultElementType]> {
+def Tosa_LogOp : Tosa_ElementwiseUnaryOp<"log"> {
   let summary = "Elementwise log op";
 
   let description = [{
@@ -1072,8 +1067,7 @@ def Tosa_LogOp : Tosa_ElementwiseOp<"log", [SameOperandsAndResultElementType]> {
 //===----------------------------------------------------------------------===//
 // Operator: logical_not
 //===----------------------------------------------------------------------===//
-def Tosa_LogicalNotOp : Tosa_ElementwiseOp<"logical_not",
-    [SameOperandsAndResultElementType]> {
+def Tosa_LogicalNotOp : Tosa_ElementwiseUnaryOp<"logical_not"> {
   let summary = "Returns the truth value of NOT x element-wise.";
 
   let description = [{
@@ -1092,8 +1086,7 @@ def Tosa_LogicalNotOp : Tosa_ElementwiseOp<"logical_not",
 //===----------------------------------------------------------------------===//
 // Operator: negate
 //===----------------------------------------------------------------------===//
-def Tosa_NegateOp : Tosa_ElementwiseOp<"negate",
-    [SameOperandsAndResultElementType]> {
+def Tosa_NegateOp : Tosa_ElementwiseUnaryOp<"negate"> {
   let summary = "Elementwise negate op";
 
   let description = [{
@@ -1117,8 +1110,7 @@ def Tosa_NegateOp : Tosa_ElementwiseOp<"negate",
 //===----------------------------------------------------------------------===//
 // Operator: reciprocal
 //===----------------------------------------------------------------------===//
-def Tosa_ReciprocalOp : Tosa_ElementwiseOp<"reciprocal",
-    [SameOperandsAndResultElementType]> {
+def Tosa_ReciprocalOp : Tosa_ElementwiseUnaryOp<"reciprocal"> {
   let summary = "Elementwise reciprocal op";
 
   let description = [{
@@ -1149,8 +1141,7 @@ def Tosa_ReciprocalOp : Tosa_ElementwiseOp<"reciprocal",
 //===----------------------------------------------------------------------===//
 // Operator: rsqrt
 //===----------------------------------------------------------------------===//
-def Tosa_RsqrtOp : Tosa_ElementwiseOp<"rsqrt",
-    [SameOperandsAndResultElementType]> {
+def Tosa_RsqrtOp : Tosa_ElementwiseUnaryOp<"rsqrt"> {
   let summary = "Elementwise 1/sqrt op";
 
   let description = [{
@@ -1170,8 +1161,7 @@ def Tosa_RsqrtOp : Tosa_ElementwiseOp<"rsqrt",
 //===----------------------------------------------------------------------===//
 // Operator: sin
 //===----------------------------------------------------------------------===//
-def Tosa_SinOp : Tosa_ElementwiseOp<"sin",
-    [SameOperandsAndResultElementType]> {
+def Tosa_SinOp : Tosa_ElementwiseUnaryOp<"sin"> {
   let summary = "Elementwise sin op";
 
   let description = [{