[mlir][tosa] Disable tosa shape verification between operands/results

[rsuderman]

Summary

Tosa shape verification prevent shape propagation when coming from a dialect
of known shape. Relax this constraint to allow ingestion / shape propagation
from these other dialects.

[rsuderman]

Relaxed tosa tensor types to allow unranked in all cases as they could be
refined to valid types.

[eric-k256]

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

122

I think it may confuse people if the Tosa_TensorND now includes unranked along with the given rank. Most of the ops which are flexible in their rank already take Tosa_Tensor, which is supposed to take ranked and unranked.
If we make this change, do we have the right validation that tensors are the right rank, as that was previously enforced by the type?
A comment noting why Tosa_UnrankedTensor is included in the defined tensor types would be helpful to add here .

[rsuderman]

Updated with comment.

The short is we wanted unranked to be supported so that shape inference can
propagate to unranked tensors. If unranked tensors are illegal it requires all
constructed TOSA IR to be always valid.

By enabling unranked we can allow the shape inference to propagate known shapes
then fail if unranked tensors remain in the IR.

[rsuderman]

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

122

Added a comment to clarify. Short story is we need flexibility on allowing unranked tensors so that shape inference on unknown shape lowerings to still lower to TOSA. Then Tosa's shape inference can handle the remaining shape inference.

Essentially we can't say an UnrankedTensor is invalid because its rank is unknown. Rather we should depend on shape inference to propagate the shape and then (only) when that shape cannot be propagate we can label the IR as invalid.

[eric-k256]

Thanks for the clarification and comment.

[stellaraccident]

Makes sense. Thanks!

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

 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: clamp
 //===----------------------------------------------------------------------===//
 def Tosa_ClampOp : Tosa_Op<"clamp", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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.
   }];
 

 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reluN
 //===----------------------------------------------------------------------===//
 def Tosa_ReluNOp : Tosa_Op<"reluN", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    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.

 }
 
 //===----------------------------------------------------------------------===//
 // Operator: tanh
 //===----------------------------------------------------------------------===//
 def Tosa_TanhOp : Tosa_Op<"tanh", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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.
   }];

 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: abs
 //===----------------------------------------------------------------------===//
 def Tosa_AbsOp : Tosa_Op<"abs", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   let summary = "Elementwise abs op";
 
   let description = [{
     Elementwise absolute value operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1

 }
 
 //===----------------------------------------------------------------------===//
 // Operator: ceil
 //===----------------------------------------------------------------------===//
 def Tosa_CeilOp : Tosa_Op<"ceil", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    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", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    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", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    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", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
-    NoSideEffect, SameOperandsAndResultType]> {
+    NoSideEffect]> {
   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.12
 // Operator Class: Type Conversion.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: cast
 //===----------------------------------------------------------------------===//
-def Tosa_CastOp: Tosa_Op<"cast", [NoSideEffect, SameOperandsAndResultShape]> {
+def Tosa_CastOp: Tosa_Op<"cast", [NoSideEffect]> {
 
   let summary = "Cast operation";
 
   let description = [{
     Performs a set of permissible cast operations
         Mode                    Input   Output
         ---------------------------------------
         signed 8 to bool        int8    Boolean

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

 
 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]>;
+// We include unranked tensors as a supported type for all possible tosa
+// Tensors as unranked does not guarantee invalid. If unranked tensors exist
+// they should be shape propagate used Tosa's shape inference pass and verified
+// to not include any remaining unranked tensors.
+def Tosa_UnrankedTensor : UnrankedTensorOf<[Tosa_AnyNumber]>;
+
+def Tosa_Tensor1D : AnyTypeOf<[Tosa_UnrankedTensor, 1DTensorOf<[Tosa_AnyNumber]>]>;
+def Tosa_Tensor2D : AnyTypeOf<[Tosa_UnrankedTensor, 2DTensorOf<[Tosa_AnyNumber]>]>;
+def Tosa_Tensor3D : AnyTypeOf<[Tosa_UnrankedTensor, 3DTensorOf<[Tosa_AnyNumber]>]>;
+def Tosa_Tensor4D : AnyTypeOf<[Tosa_UnrankedTensor, 4DTensorOf<[Tosa_AnyNumber]>]>;
+def Tosa_Tensor5D : AnyTypeOf<[Tosa_UnrankedTensor, TensorRankOf<[Tosa_AnyNumber], [5]>]>;
 
 // Ranked tensors up to given rank.
-def Tosa_Tensor1Dto4D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4]>;
-def Tosa_Tensor1Dto6D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4,5,6]>;
+def Tosa_Tensor1Dto4D : AnyTypeOf<[
+  Tosa_UnrankedTensor, TensorRankOf<[Tosa_AnyNumber], [1,2,3,4]>]>;
+def Tosa_Tensor1Dto6D : AnyTypeOf<[
+  Tosa_UnrankedTensor, TensorRankOf<[Tosa_AnyNumber], [1,2,3,4,5,6]>]>;
 
-def Tosa_TensorUpto4D : TensorRankOf<[Tosa_AnyNumber], [0,1,2,3,4]>;
+def Tosa_TensorUpto4D : AnyTypeOf<[
+  Tosa_UnrankedTensor, TensorRankOf<[Tosa_AnyNumber], [0,1,2,3,4]>]>;
 
-def Tosa_Int32TensorUpto4D : TensorRankOf<[Tosa_Int32], [0,1,2,3,4]>;
+def Tosa_Int32TensorUpto4D : AnyTypeOf<[
+  Tosa_UnrankedTensor, 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<[
      AnyTypeOf<types>.predicate,
      VectorOf<types>.predicate,

mlir/include/mlir/IR/OpBase.td

 
 // Whether a type is a MemRefType.
 def IsMemRefTypePred : CPred<"$_self.isa<::mlir::MemRefType>()">;
 
 // Whether a type is an UnrankedMemRefType
 def IsUnrankedMemRefTypePred
         : CPred<"$_self.isa<::mlir::UnrankedMemRefType>()">;
 
+// Whether a type is an UnrankedTensorType
+def IsUnrankedTensorTypePred
+        : CPred<"$_self.isa<::mlir::UnrankedTensorType>()">;
+
 // Whether a type is a BaseMemRefType
 def IsBaseMemRefTypePred
         : CPred<"$_self.isa<::mlir::BaseMemRefType>()">;
 
 // Whether a type is a ShapedType.
 def IsShapedTypePred : CPred<"$_self.isa<::mlir::ShapedType>()">;
 
 // For a ShapedType, verify that it has a static shape.
ShapedContainerType<allowedTypes, IsTensorTypePred, "tensor",
                       "::mlir::TensorType">;
 
 class RankedTensorOf<list<Type> allowedTypes> :
   ShapedContainerType<allowedTypes, And<[IsTensorTypePred, HasRankPred]>,
   "ranked tensor", "::mlir::TensorType">;
 
 def AnyTensor : TensorOf<[AnyType]>;
 
+// Unranked Memref type
+class UnrankedTensorOf<list<Type> allowedTypes> :
+    ShapedContainerType<allowedTypes,
+                        IsUnrankedTensorTypePred,
+                        "unranked.tensor", "::mlir::UnrankedTensorType">;
+
 def AnyRankedTensor : RankedTensorOf<[AnyType]>;
 
 // TODO: Have an easy way to add another constraint to a type.
 class StaticShapeTensorOf<list<Type> allowedTypes>
     : Type<And<[TensorOf<allowedTypes>.predicate, HasStaticShapePred]>,
            "statically shaped " # TensorOf<allowedTypes>.summary,
            "::mlir::TensorType">;