[mlir][tosa] Align tensor rank specifications with current spec

[sjarus]

Summary

Deconstrains several TOSA operators to align with the current TOSA spec, including all the elementwise ops.
Note: some more ops are under consideration for further cleanup; they will follow once the spec has been updated.

[stellaraccident]

Nice improvement - thanks. Do you need me to land this for you?

[mehdi_amini]

Can you upload patches with full context please? Right now I can't see which ops are modified here.

[stellaraccident]

In D102958#2776297, @mehdi_amini wrote:

Can you upload patches with full context please? Right now I can't see which ops are modified here.

I had a similar question but was able to infer based on eliminating the entire classes of restrictions that had been lifted from the spec.

I don't know phab well enough to know why it can't show context (or how to recommend uploading differently) -- Mehdi, do you know?

[sjarus]

In D102958#2776308, @stellaraccident wrote:

In D102958#2776297, @mehdi_amini wrote:

Can you upload patches with full context please? Right now I can't see which ops are modified here.

I had a similar question but was able to infer based on eliminating the entire classes of restrictions that had been lifted from the spec.

I don't know phab well enough to know why it can't show context (or how to recommend uploading differently) -- Mehdi, do you know?

I have the same question. Scrolling down the phab content it does look strange that I can't expand to see the full code. As background, I just did git format-patch -1 HEAD and then uploaded the patch on the phab web interface.

[sjarus]

Trying using arc diff --update from commandline.

[sjarus]

Ok now it seems to look a lot better ? If this is alright, please land it on my behalf - it doesn't seem I can do so myself.

[stellaraccident]

Ah, you were using the web interface. The docs have extra flags that must be included for full context: https://llvm.org/docs/Phabricator.html#id5

I always use arc diff.

[sjarus]

In D102958#2776457, @stellaraccident wrote:

Ah, you were using the web interface. The docs have extra flags that must be included for full context: https://llvm.org/docs/Phabricator.html#id5

I always use arc diff.

Thanks for the information! I'll probably stick to using arc too, from now on. I've updated our internal documentation to recommend its use as well.

[stellaraccident]

Thanks for updating the patch. I re-reviewed this and cross referenced with the spec, and lgtm. Left one comment for the future. Landing as-is.

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

504

Some feedback (not for this change): Reviewing this prompted me to look the operator up in the spec and hunt around for a corresponding floating point division. The verbiage in the spec "Only used for integer operation. Floating point divide should use RECIPROCAL and MUL." is pretty salient and (imo) would be valuable to include here, given that keeping this integer only was an explicit design choice (and naming/convention-wise is not obvious).

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

def Tosa_ClampOp : Tosa_Op<"clamp", [NoSideEffect, SameOperandsAndResultType]> {
 
   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_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$min_int,
     I64Attr:$max_int,
     F32Attr:$min_fp,
     F32Attr:$max_fp
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reluN
 //===----------------------------------------------------------------------===//
 def Tosa_ReluNOp : Tosa_Op<"reluN", [NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Computes rectified linear: `max(features, N)`.";
 
   let description = [{
      ReLU with a scalar maximum value.
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$max_int,
     F32Attr:$max_fp
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: sigmoid
 //===----------------------------------------------------------------------===//
 def Tosa_SigmoidOp : Tosa_Op<"sigmoid", [NoSideEffect,
                              SameOperandsAndResultType]> {
   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_Tensor1Dto4D:$input
+    Tosa_Tensor:$input
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: tanh
 //===----------------------------------------------------------------------===//
 def Tosa_TanhOp : Tosa_Op<"tanh", [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.
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input
+    Tosa_Tensor:$input
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // 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]> {
   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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: arithmetic_right_shift
 //===----------------------------------------------------------------------===//
 def Tosa_ArithmeticRightShiftOp : Tosa_Op<"arithmetic_right_shift",
                                           [ResultsBroadcastableShape,
                                           NoSideEffect]> {
   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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2,
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2,
     BoolAttr:$round
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_and
 //===----------------------------------------------------------------------===//
 def Tosa_BitwiseAndOp : Tosa_Op<"bitwise_and", [ResultsBroadcastableShape,
                                 NoSideEffect, Commutative]> {
   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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_or
 //===----------------------------------------------------------------------===//
 def Tosa_BitwiseOrOp : Tosa_Op<"bitwise_or", [ResultsBroadcastableShape,
                                               NoSideEffect, Commutative]> {
   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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_xor
 //===----------------------------------------------------------------------===//
 def Tosa_BitwiseXorOp : Tosa_Op<"bitwise_xor", [ResultsBroadcastableShape,
                                                 NoSideEffect, Commutative]> {
   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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: div
 //===----------------------------------------------------------------------===//
 def Tosa_DivOp : Tosa_Op<"div", [ResultsBroadcastableShape,
                                  NoSideEffect]> {
   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_Int32TensorUpto4D:$input1,
-    Tosa_Int32TensorUpto4D:$input2
+    Tosa_Int32Tensor:$input1,
+    Tosa_Int32Tensor:$input2
   );
 
   let results = (outs
-    Tosa_Int32TensorUpto4D:$output
+    Tosa_Int32Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_and
 //===----------------------------------------------------------------------===//
 def Tosa_LogicalAndOp : Tosa_Op<"logical_and", [ResultsBroadcastableShape,
                                                 Commutative, NoSideEffect]> {
def Tosa_LogicalLeftShiftOp : Tosa_Op<"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. Rank of input tensors must match.
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_right_shift
 //===----------------------------------------------------------------------===//
 def Tosa_LogicalRightShiftOp : Tosa_Op<"logical_right_shift",
                                        [ResultsBroadcastableShape,
                                         NoSideEffect]> {
   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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_or
 //===----------------------------------------------------------------------===//
 def Tosa_LogicalOrOp : Tosa_Op<"logical_or", [ResultsBroadcastableShape,
                                               Commutative, NoSideEffect]> {
def Tosa_MaximumOp : Tosa_Op<"maximum", [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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: minimum
 //===----------------------------------------------------------------------===//
 def Tosa_MinimumOp : Tosa_Op<"minimum", [ResultsBroadcastableShape,
                                          NoSideEffect, Commutative]> {
   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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: mul
 //===----------------------------------------------------------------------===//
 def Tosa_MulOp : Tosa_Op<"mul", [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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2,
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2,
     I32Attr:$shift
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pow
 //===----------------------------------------------------------------------===//
 def Tosa_PowOp : Tosa_Op<"pow", [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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$z
+    Tosa_Tensor:$z
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: sub
 //===----------------------------------------------------------------------===//
 def Tosa_SubOp : Tosa_Op<"sub", [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_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: table
 //===----------------------------------------------------------------------===//
 def Tosa_TableOp : Tosa_Op<"table", [NoSideEffect]> {
   let summary = "Table lookup op";
* A RESCALE node is expected before the TABLE operator to scale the input
       to a full int16_t range for the table lookup
     * If an int16_t result is required then follow the TABLE operator with a
       RESCALE with a right shift of 7
     * If an int8_t result is required then follow the TABLE operator with a
       RESCALE with a right shift of 15
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D: $input,
+    Tosa_Tensor: $input,
     Tosa_Tensor1D: $table
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 
 }
 
 //===----------------------------------------------------------------------===//
 // 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]> {
   let summary = "Elementwise abs op";
 
   let description = [{
     Elementwise absolute value operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: bitwise_not
 //===----------------------------------------------------------------------===//
 def Tosa_BitwiseNotOp : Tosa_Op<"bitwise_not", [ResultsBroadcastableShape,
                                                 NoSideEffect]> {
   let summary = "Bitwise NOT operator";
 
   let description = [{
     Elementwise bitwise NOT of input tensor.
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: ceil
 //===----------------------------------------------------------------------===//
 def Tosa_CeilOp : Tosa_Op<"ceil", [NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Elementwise ceil op";
 
   let description = [{
     Elementwise ceiling operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: clz
 //===----------------------------------------------------------------------===//
 def Tosa_ClzOp : Tosa_Op<"clz", [NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Elementwise count leading zero op";
 
   let description = [{
     Elementwise count leading zeros operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: exp
 //===----------------------------------------------------------------------===//
 def Tosa_ExpOp : Tosa_Op<"exp", [NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Elementwise exp op";
 
   let description = [{
     Elementwise e to the x operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: floor
 //===----------------------------------------------------------------------===//
 def Tosa_FloorOp : Tosa_Op<"floor", [NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Elementwise floor op";
 
   let description = [{
     Elementwise floor operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: log
 //===----------------------------------------------------------------------===//
 def Tosa_LogOp : Tosa_Op<"log", [NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Elementwise log op";
 
   let description = [{
     Elementwise natural logarithm operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_not
 //===----------------------------------------------------------------------===//
 def Tosa_LogicalNotOp : Tosa_Op<"logical_not", [NoSideEffect,
                                                 SameOperandsAndResultType]> {
def Tosa_NegateOp : Tosa_Op<"negate", [NoSideEffect,
                                        SameOperandsAndResultType]> {
   let summary = "Elementwise negate op";
 
   let description = [{
     Elementwise negation operation
   }];
 
   let arguments = (ins
-      Tosa_TensorUpto4D:$input1,
+      Tosa_Tensor:$input1,
       OptionalAttr<Tosa_UnaryOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 
   let builders = [Tosa_UnaryOpQuantInfoBuilder];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reciprocal
 //===----------------------------------------------------------------------===//
 def Tosa_ReciprocalOp : Tosa_Op<"reciprocal", [NoSideEffect,
                                                SameOperandsAndResultType]> {
   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_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: rsqrt
 //===----------------------------------------------------------------------===//
 def Tosa_RsqrtOp : Tosa_Op<"rsqrt", [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_TensorUpto4D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // 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]> {
   let summary = "Elementwise select operator";
 
   let description = [{
     Elementwise select of the output based on a condition.
   }];
 
   let arguments = (ins
     I1Tensor:$input1,
-    Tosa_TensorUpto4D:$input2,
-    Tosa_TensorUpto4D:$input3
+    Tosa_Tensor:$input2,
+    Tosa_Tensor:$input3
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.7
 // Operator Class: Logical Operations.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: equal
 //===----------------------------------------------------------------------===//
 def Tosa_EqualOp : Tosa_Op<"equal", [ResultsBroadcastableShape, Commutative,
                                      NoSideEffect]> {
   let summary = "Returns the truth value of (x == y) element-wise.";
 
   let description = [{
      Elementwise comparison operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: greater
 //===----------------------------------------------------------------------===//
 def Tosa_GreaterOp : Tosa_Op<"greater", [ResultsBroadcastableShape,
                                          NoSideEffect]> {
   let summary = "Returns the truth value of (x > y) element-wise.";
 
   let description = [{
     Elementwise greater than comparison operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: greater_equal
 //===----------------------------------------------------------------------===//
 def Tosa_GreaterEqualOp : Tosa_Op<"greater_equal", [ResultsBroadcastableShape,
                                                     NoSideEffect]> {
   let summary = "Returns the truth value of (x >= y) element-wise.";
 
   let description = [{
     Elementwise comparison operation
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input1,
-    Tosa_TensorUpto4D:$input2
+    Tosa_Tensor:$input1,
+    Tosa_Tensor:$input2
   );
 
   let results = (outs
     I1Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
def Tosa_ConcatOp : Tosa_Op<"concat", [NoSideEffect]> {
   let summary = "Concatenates tensors along one dimension.";
 
   let description = [{
     Concatenate a variadic amount of tensors along a given axis. No data 
     conversion happens during a concat operation.
   }];
 
   let arguments = (ins
-    Variadic<Tosa_Tensor1Dto4D>:$input1,
+    Variadic<Tosa_RankedTensor>:$input1,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_RankedTensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pad
 //===----------------------------------------------------------------------===//
 def Tosa_PadOp : Tosa_Op<"pad", [NoSideEffect]> {
   let summary = "Pads a tensor with zeros.";
 
   let description = [{
     Zero-pads a tensor along borders of each dimension.
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input1,
+    Tosa_RankedTensor:$input1,
     Tosa_Int32Or64Tensor:$padding,
     OptionalAttr<Tosa_PadOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_RankedTensor:$output
   );
 
   let builders = [Tosa_PadOpQuantInfoBuilder];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reshape
 //===----------------------------------------------------------------------===//
 def Tosa_ReshapeOp: Tosa_Op<"reshape", [
    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_TensorUpto6D:$input1,
+    Tosa_Tensor:$input1,
     I64ArrayAttr:$new_shape
   );
 
   let results = (outs
-    Tosa_TensorUpto6D:$output
+    Tosa_RankedTensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reverse
 //===----------------------------------------------------------------------===//
 def Tosa_ReverseOp: Tosa_Op<"reverse", [NoSideEffect]> {
   let summary = "Reverse operator";
Performs a set of permissible cast operations
         signed 32 to signed 16  int32   int16
         float to signed 8       float   int8
         float to signed 16      float   int16
         signed 8 to float       int8    float
         signed 16 to float      int16   float
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input
+    Tosa_Tensor:$input
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: rescale
 //===----------------------------------------------------------------------===//
 def Tosa_RescaleOp: Tosa_Op<"rescale", [NoSideEffect]> {
   let summary = "Tosa rescale operator";
let description = [{
     signed 48 to 8          int48   int8
     signed 48 to 16         int48   int16
     signed 48 to 32         int48   int32
     unsigned 8 to signed 8  uint8   int8
     signed 8 to unsigned 8  int8    uint8
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto4D:$input,
+    Tosa_Tensor:$input,
     I32Attr:$input_zp,
     I32Attr:$output_zp,
     I32ArrayAttr:$multiplier,
     I32ArrayAttr:$shift,
     BoolAttr:$scale32,
     BoolAttr:$double_round,
     BoolAttr:$per_channel
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.13
 // Operator Class: Data Node Ops.
 //===----------------------------------------------------------------------===//
 
let description = [{
     hold data in any of the supported data formats.
   }];
 
   let arguments = (ins
     ElementsAttr:$value
   );
 
   let results = (outs
-    Tosa_TensorUpto4D:$output
+    Tosa_Tensor:$output
   );
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: identity
 //===----------------------------------------------------------------------===//
 def Tosa_IdentityOp: Tosa_Op<"identity", [NoSideEffect]> {
   let summary = "Identity operator";
   let description = [{
     Returns a tensor with the same shape, size, type
     and content as the input.
   }];
 
   let arguments = (ins
-    Tosa_TensorUpto6D:$input1
+    Tosa_Tensor:$input1
   );
 
   let results = (outs
-    Tosa_TensorUpto6D:$output
+    Tosa_Tensor:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.14
 // Operator Class: Custom Operators.
 //===----------------------------------------------------------------------===//
 

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

 //===----------------------------------------------------------------------===//
 def Tosa_AnyNumber : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float],
                                "number">;
 
 //===----------------------------------------------------------------------===//
 // Tensor types
 //===----------------------------------------------------------------------===//
 
+def Tosa_Int32Tensor : TensorOf<[Tosa_Int32]>;
 def Tosa_Int32Or64Tensor : TensorOf<[Tosa_Int32Or64]>;
 
+// Either ranked or unranked tensor of TOSA supported element types. 
 def Tosa_Tensor : TensorOf<[Tosa_AnyNumber]>;
+// Must be ranked but no further constraints
+def Tosa_RankedTensor : RankedTensorOf<[Tosa_AnyNumber]>;
 
 // Any tensor element type allowed in Tosa ops.
 def Tosa_ElementType : Type<Or<[Tosa_Int.predicate, Tosa_QuantizedInt.predicate,
                                 Tosa_Float.predicate]>, "tosa.dtype">;
 
 class Tosa_TensorOfOrNone<list<Type> allowedTypes, string description = ""> :
   AnyTypeOf<[TensorOf<allowedTypes>, NoneType], description>;
 
 //===----------------------------------------------------------------------===//
 // Tensor types with constrained ranks.
 //===----------------------------------------------------------------------===//
 
 // Must be listed rank.
 def Tosa_Tensor1D : 1DTensorOf<[Tosa_AnyNumber]>;
 def Tosa_Tensor2D : 2DTensorOf<[Tosa_AnyNumber]>;
 def Tosa_Tensor3D : 3DTensorOf<[Tosa_AnyNumber]>;
 def Tosa_Tensor4D : 4DTensorOf<[Tosa_AnyNumber]>;
 def Tosa_Tensor5D : TensorRankOf<[Tosa_AnyNumber], [5]>;
-def Tosa_Tensor6D : TensorRankOf<[Tosa_AnyNumber], [6]>;
 
 // Ranked tensors up to given rank.
-def Tosa_Tensor1Dto2D : TensorRankOf<[Tosa_AnyNumber], [1,2]>;
 def Tosa_Tensor1Dto4D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4]>;
-def Tosa_Tensor1Dto5D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4,5]>;
 def Tosa_Tensor1Dto6D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4,5,6]>;
 
 def Tosa_TensorUpto4D : TensorRankOf<[Tosa_AnyNumber], [0,1,2,3,4]>;
-def Tosa_TensorUpto6D : TensorRankOf<[Tosa_AnyNumber], [0,1,2,3,4,5,6]>;
 
 def Tosa_Int32TensorUpto4D : TensorRankOf<[Tosa_Int32], [0,1,2,3,4]>;
 
 //===----------------------------------------------------------------------===//
 // Generic scalar, vector, or tensor of a particular type.
 //===----------------------------------------------------------------------===//
 
 class Tosa_TypeLike<list<Type> types, string description = ""> : TypeConstraint<Or<[