[mlir][tosa] Add TOSA f64 type support for cast op

[AmosLewis]

Summary

Add TOSA f64 type support for cast op

[AmosLewis]

Find this f64 support issue in https://github.com/llvm/torch-mlir/pull/1802

[eric-k256]

The TOSA specification doesn't support F64, so I think we should be careful in how we expand the data types in the dialect. Adding it has implications for TOSA consumers.

At a minimum, code should be added to the tosa-validate pass that warns backends that attempt to align with the TOSA specification that non-compliant data types are present. The pass exists, it should be simple to report errors if FP64 operands are detected for any profile, as fp64 doesn't exist in a profile.

I'd also recommend that we proceed more incrementally with fp64. Rather than changing Tosa_Float, which changes almost all of the operators, create a new type, which is Tosa_Float + fp64, and use that only where needed. So if we only need fp64 for the CAST operator, the new type only gets used there. This would allow consumers of TOSA dialect to understand where f64 is needed, and would help reason about how fp64 would get added to the specification.

If any consumers of TOSA dialect have thoughts on F64 support, whether they want to avoid it, or are already capable of consuming it, it would be good to hear, even outside of this specific review.

[AmosLewis]

[mlir][tosa] Add TOSA f64 type support for cast op

[eric-k256]

Can you add a simple check to TosaValidation::runOnOperation that goes into the existing for loop and checks for float64 type, and does a signalPassFailure there? You don't need to check profileType, as f64 isn't in the spec. This pass is for implementations to use to check against the spec, so dialect specific changes that aren't in the spec should be flagged when it is run.

[AmosLewis]

Add TosaValidation float64 type check and signalPassFailure

[AmosLewis]

In D142599#4101032, @eric-k256 wrote:

Can you add a simple check to TosaValidation::runOnOperation that goes into the existing for loop and checks for float64 type, and does a signalPassFailure there? You don't need to check profileType, as f64 isn't in the spec. This pass is for implementations to use to check against the spec, so dialect specific changes that aren't in the spec should be flagged when it is run.

I am not familiar with this part of the code, could you check if adding is correct? Or any way to check if I am adding it correctly?

if (getElementTypeOrSelf(operand).isF64()) {
  return signalPassFailure();
}

[eric-k256]

Thanks. Looks good to me now.

[AmosLewis]

In D142599#4101676, @eric-k256 wrote:

Thanks. Looks good to me now.

How long will it take for this patch to land?

[eric-k256]

I don't have commit privileges. @rsuderman or @NatashaKnk, can you take a look?

[jpienaar]

I'll help in a bit

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

Performs a set of permissible cast operations
         signed 16 to signed 8   int16   int8
         signed 16 to signed 32  int16   int32
         signed 32 to signed 8   int32   int8
         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
+        float 32 to float 64    float32 float64
+        float 64 to float 32    float64 float32
   }];
 
   let arguments = (ins
-    Tosa_Tensor:$input
+    Tosa_Tensor_Cast:$input
   );
 
   let results = (outs
-    Tosa_Tensor:$output
+    Tosa_Tensor_Cast:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: rescale
 //===----------------------------------------------------------------------===//

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

def Tosa_Float : AnyTypeOf<[
           F16,
           BF16]>;
 
 //===----------------------------------------------------------------------===//
 // Multi-category types.
 //===----------------------------------------------------------------------===//
 def Tosa_AnyNumber : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float],
                                "number">;
+// Add F64 type support for just for tosa::CastOp
+def Tosa_AnyNumber_Cast : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float, F64],
+                               "number_cast">;
 
 //===----------------------------------------------------------------------===//
 // 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]>;
+def Tosa_Tensor_Cast : TensorOf<[Tosa_AnyNumber_Cast]>;
 // 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 = ""> :

mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp

void TosaValidation::runOnOperation() {
   profileType = symbolizeEnum<TosaProfileEnum>(profileName);
 
   getOperation().walk([&](Operation *op) {
     for (Value operand : op->getOperands()) {
       if ((profileType == TosaProfileEnum::BaseInference) &&
           getElementTypeOrSelf(operand).isa<FloatType>()) {
         return signalPassFailure();
       }
+      if (getElementTypeOrSelf(operand).isF64()) {
+        return signalPassFailure();
+      }
     }
   });
 }
 } // namespace
 
 std::unique_ptr<Pass> mlir::tosa::createTosaValidationPass() {
   return std::make_unique<TosaValidation>();
 }