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

[AmosLewis]

Summary

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

[AmosLewis]

This patch fix torch.prim.NumToTensor.Scalar float64 support https://github.com/llvm/torch-mlir/pull/1802
This fix should have been added in
https://reviews.llvm.org/rGa2dcd994a7f8cc33640f58105276b78acf3483e5

[jpienaar]

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

1824

This naming and comment around the def is confusing now. Could you update?

[AmosLewis]

Chang the tensor with f64 support name from Tosa_Tensor_Cast to Tosa_Tensor_F64

[AmosLewis]

Chang the tensor with f64 support name from Tosa_Tensor_Cast to Tosa_Tensor_64

[AmosLewis]

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

1824

How about now? I name all the tensor support f64 with Tosa_Tensor_64.

[eric-k256]

This is definitely an improvement, I'm okay with this type, but it reads as only a 64-bit tensor. Perhaps Jacques is better at naming than I am, but I'd go something closer to Tosa_Tensor_Plus_F64, indicating that this is an extension of Tosa_Tensor with F64 (as opposed to I64).

[AmosLewis]

Chang the tensor with f64 support name from Tosa_Tensor_Cast to Tosa_Tensor_Plus_F64

[AmosLewis]

In D145336#4173669, @eric-k256 wrote:

This is definitely an improvement, I'm okay with this type, but it reads as only a 64-bit tensor. Perhaps Jacques is better at naming than I am, but I'd go something closer to Tosa_Tensor_Plus_F64, indicating that this is an extension of Tosa_Tensor with F64 (as opposed to I64).

Done

[eric-k256]

Looks reasonable to me, although I don't have merge priveleges.

[manupak]

Quick question @eric-k256 @AmosLewis,

Will it also fold it to be non-F64 ? i.e. tosa consumers are able to not to witness f64 when tosa is in the canonical form.

[jpienaar]

In D145336#4177567, @manupak wrote:

Quick question @eric-k256 @AmosLewis,

Will it also fold it to be non-F64 ? i.e. tosa consumers are able to not to witness f64 when tosa is in the canonical form.

I think this should be explicit pass as it may be interprocedural. Consumers post the "verify for profile" (or runtime, forgot what it is called) should not be getting these. Either by way of it errors out during that verify pass or converted to f32 (or some such pass, which would be opt in and may change precision).

For me canonical would be a function of profile.

[jpienaar]

(which effectively means no not canonical, but one has sections with clear expectations and post "deployment" profile verification there should be none of these for consumers)

[manupak]

Thanks, sounds like a no for now and agree with interprocedural implications.

I was curious because https://reviews.llvm.org/D142599 sounded like F64 wasn't in any profile
and also there are folders for splat based consts (which might help a bit) that also have the same said implications.

[eric-k256]

You're right that F64 isn't in any of the TOSA profiles today, which is why I was discouraging it's use in the overall definition of Tosa_Tensor. It's something we look at adding, but it adds significant requirements to any profiles it goes into. It would be good to get a sense of what networks need F64 vs ending up with that type as default. Many systems that support f64 do it at a performance cost against f32, and some systems don't implement it at all. Ideally the tooling would have a way to guide developers to minimize their use of f64 to where it has a significant improvement on results justifying the extra computation cost.

[manupak]

I concur.

Hence, the question: should we also look to fold it as much as possible ?

I think the cases we have seen are just F64 constants being created (maybe due to default python behaviour?) and then immediately casted to lower bitwidths.
Exactly what this patch is doing in other words -- so its a step in the right direction.

I think splats are folded; so we are good there - it might cover most bases.

I was just curious whether extending the folding of tosa.cast to tosa.const operands to limit the appearance of f64 where possible -- was a good idea or not, while we are here.
Especially if we can catch them before outlining -- i.e. scenarios where we wouldn't hit the inter-procedural case.

[AmosLewis]

In D145336#4178091, @jpienaar wrote:

In D145336#4177567, @manupak wrote:

Quick question @eric-k256 @AmosLewis,

Will it also fold it to be non-F64 ? i.e. tosa consumers are able to not to witness f64 when tosa is in the canonical form.

I think this should be explicit pass as it may be interprocedural. Consumers post the "verify for profile" (or runtime, forgot what it is called) should not be getting these. Either by way of it errors out during that verify pass or converted to f32 (or some such pass, which would be opt in and may change precision).

For me canonical would be a function of profile.

@jpienaar could you land this patch?

[mehdi_amini]

Can you add tests please?

[AmosLewis]

[mlir][tosa] Add a test for tosa::const 64

[AmosLewis]

delete extra line

[AmosLewis]

In D145336#4191332, @mehdi_amini wrote:

Can you add tests please?

Done

[AmosLewis]

@jpienaar could you help land this patch when you are available?

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

Performs a set of permissible cast operations
         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_Cast:$input
+    Tosa_Tensor_Plus_F64:$input
   );
 
   let results = (outs
-    Tosa_Tensor_Cast:$output
+    Tosa_Tensor_Plus_F64:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: rescale
 //===----------------------------------------------------------------------===//
let description = [{
     hold data in any of the supported data formats.
   }];
 
   let arguments = (ins
     ElementsAttr:$value
   );
 
   let results = (outs
-    Tosa_Tensor:$output
+    Tosa_Tensor_Plus_F64:$output
   );
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: identity
 //===----------------------------------------------------------------------===//
 def Tosa_IdentityOp: Tosa_Op<"identity", [Pure,

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">;
+// Add F64 type support just for tosa::CastOp and tosa::ConstOp
+def Tosa_AnyNumber_Plus_F64 : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float, F64],
+                               "number_plus_f64">;
 
 //===----------------------------------------------------------------------===//
 // 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]>;
+def Tosa_Tensor_Plus_F64 : TensorOf<[Tosa_AnyNumber_Plus_F64]>;
 // 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/test/Dialect/Tosa/constant_folding.mlir

 // RUN: mlir-opt --test-constant-fold %s | FileCheck %s
 
 // CHECK-LABEL: func @test_const
 func.func @test_const(%arg0 : index) -> tensor<4xi32> {
   // CHECK: "tosa.const"
   %0 = "tosa.const"() {value = dense<[3, 0, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
   return %0 : tensor<4xi32>
 }
 
+// CHECK-LABEL: func @test_const_i64
+func.func @test_const_i64(%arg0 : index) -> tensor<4xi64> {
+  // CHECK: "tosa.const"
+  %0 = "tosa.const"() {value = dense<[3, 0, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64>
+  return %0 : tensor<4xi64>
+}
+
 // CHECK-LABEL: func @try_fold_equal_with_unranked_tensor
 func.func @try_fold_equal_with_unranked_tensor(%arg0: tensor<4xi32>, %arg1: tensor<i32>) {
   // CHECK: "tosa.equal"
   // CHECK-NEXT: return
   %0 = "tosa.equal"(%arg0, %arg1) : (tensor<4xi32>, tensor<i32>) -> tensor<*xi1>
   return
 }