[mlir][tosa] Add ext-mxfp support for const and cast ops

This commit allows the creation of const/cast operations with
MXFP datatypes.

Note: it doesn't include support for the mxint8 datatype. This will
be added in a separate commit.

Note: this commit adds support as defined in the spec in
https://review.mlplatform.org/c/tosa/specification/+/15362. EXT_MXFP
extension is considered experimental and subject to breaking change.

Change-Id: Idd0477bc947ade524b0fb0213cc7e8d4f892ddab

Author: lhutton1

mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc

@@ -747,7 +747,14 @@ extensionComplianceMap = {
        {{{fp8e5m2T, fp16T}, SpecificationVersion::V_1_0},
         {{fp8e5m2T, fp32T}, SpecificationVersion::V_1_0},
         {{fp16T, fp8e5m2T}, SpecificationVersion::V_1_0},
-        {{fp32T, fp8e5m2T}, SpecificationVersion::V_1_0}}}}},
+        {{fp32T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
+      {{Extension::bf16, Extension::mxfp},
+       {{{fp4e2m1T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp6e3m2T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp6e2m3T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{bf16T, fp4e2m1T}, SpecificationVersion::V_1_1_DRAFT},
+        {{bf16T, fp6e3m2T}, SpecificationVersion::V_1_1_DRAFT},
+        {{bf16T, fp6e2m3T}, SpecificationVersion::V_1_1_DRAFT}}, allOf}}},
     {"tosa.cast_from_block_scaled",
      {{{Extension::bf16, Extension::mxfp},
        {{{fp4e2m1T, fp8ue8m0T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
@@ -784,7 +791,12 @@ extensionComplianceMap = {
       {{Extension::int16}, {{{i48T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e4m3}, {{{fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2}, {{{fp8e5m2T}, SpecificationVersion::V_1_0}}},
-      {{Extension::bf16}, {{{bf16T}, SpecificationVersion::V_1_0}}}}},
+      {{Extension::bf16}, {{{bf16T}, SpecificationVersion::V_1_0}}},
+      {{Extension::mxfp},
+       {{{fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp6e3m2T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp6e2m3T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp4e2m1T}, SpecificationVersion::V_1_1_DRAFT}}}}},
     {"tosa.identity",
      {{{Extension::int4}, {{{i4T, i4T}, SpecificationVersion::V_1_0}}},
       {{Extension::int16}, {{{i48T, i48T}, SpecificationVersion::V_1_0}}},

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

@@ -2464,7 +2464,7 @@ def Tosa_CastOp: Tosa_Op<"cast", [Pure, SameOperandsAndResultShape,
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16, Tosa_EXT_MXFP]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -2643,7 +2643,7 @@ def Tosa_ConstOp : Tosa_Op<"const", [ConstantLike, Pure,
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_INT4, Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_INT4, Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16, Tosa_EXT_MXFP]>,
   ];
 
   let hasFolder = 1;

mlir/test/Dialect/Tosa/availability.mlir

@@ -606,7 +606,7 @@ func.func @test_resize(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x64x64x8xf32> {
 // CHECK-LABEL: cast
 func.func @test_cast1(%arg0: tensor<13x21x3xi32>) -> tensor<13x21x3xf32> {
   // CHECK: profiles: [ [pro_int, pro_fp] ]
-  // CHECK: extensions: [ [fp8e4m3, fp8e5m2, bf16] ]
+  // CHECK: extensions: [ [fp8e4m3, fp8e5m2, bf16, mxfp] ]
   %0 = tosa.cast %arg0 : (tensor<13x21x3xi32>) -> tensor<13x21x3xf32>
   return %0 : tensor<13x21x3xf32>
 }
@@ -626,7 +626,7 @@ func.func @test_rescale(%arg0: tensor<13x21x3x!quant.uniform<u8:f32, 0.015655439
 // CHECK-LABEL: test_const
 func.func @test_const(%arg0 : index) -> tensor<4xi32> {
   // CHECK: profiles: [ [pro_int, pro_fp] ]
-  // CHECK: extensions: [ [int4, int16, fp8e4m3, fp8e5m2, bf16] ]
+  // CHECK: extensions: [ [int4, int16, fp8e4m3, fp8e5m2, bf16, mxfp] ]
     %0 = "tosa.const"() {values = dense<[3, 0, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
     return %0 : tensor<4xi32>
 }

mlir/test/Dialect/Tosa/invalid_extension.mlir

@@ -562,3 +562,17 @@ func.func @test_cast_to_block_scaled_static(%arg0: tensor<4x32xf32>) -> (tensor<
   %0:2 = tosa.cast_to_block_scaled %arg0 {block_size = #tosa.block_size<BLOCK_SIZE_32>} : (tensor<4x32xf32>) -> (tensor<4x32xf6E3M2FN>, tensor<4x1xf8E8M0FNU>)
   return %0#0, %0#1 : tensor<4x32xf6E3M2FN>, tensor<4x1xf8E8M0FNU>
 }
+
+// -----
+func.func @test_const_fp6e3m2(%arg0 : index) -> tensor<4xf6E3M2FN> {
+  // expected-error@+1 {{'tosa.const' op illegal: requires [mxfp] but not enabled in target}}
+    %0 = "tosa.const"() {values = dense<[0.0, 0.0, 0.0, 0.0]> : tensor<4xf6E3M2FN>} : () -> tensor<4xf6E3M2FN>
+    return %0 : tensor<4xf6E3M2FN>
+}
+
+// -----
+func.func @test_cast_f4e2m1(%arg0: tensor<13x21x3xf4E2M1FN>) -> tensor<13x21x3xbf16> {
+  // expected-error@+1 {{'tosa.cast' op illegal: requires all of [bf16, mxfp] but not enabled in target}}
+  %0 = tosa.cast %arg0 : (tensor<13x21x3xf4E2M1FN>) -> tensor<13x21x3xbf16>
+  return %0 : tensor<13x21x3xbf16>
+}

mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir

@@ -50,3 +50,19 @@ func.func @test_cast_to_block_scaled_static(%arg0: tensor<4x32xf32>) -> (tensor<
   %0:2 = tosa.cast_to_block_scaled %arg0 {block_size = #tosa.block_size<BLOCK_SIZE_32>} : (tensor<4x32xf32>) -> (tensor<4x32xf6E3M2FN>, tensor<4x1xf8E8M0FNU>)
   return %0#0, %0#1 : tensor<4x32xf6E3M2FN>, tensor<4x1xf8E8M0FNU>
 }
+
+// -----
+
+// CHECK-LABEL: test_const_fp6e3m2
+func.func @test_const_fp6e3m2(%arg0 : index) -> tensor<4xf6E3M2FN> {
+    %0 = "tosa.const"() {values = dense<[0.0, 0.0, 0.0, 0.0]> : tensor<4xf6E3M2FN>} : () -> tensor<4xf6E3M2FN>
+    return %0 : tensor<4xf6E3M2FN>
+}
+
+// -----
+
+// CHECK-LABEL: test_cast_f4e2m1
+func.func @test_cast_f4e2m1(%arg0: tensor<13x21x3xf4E2M1FN>) -> tensor<13x21x3xbf16> {
+  %0 = tosa.cast %arg0 : (tensor<13x21x3xf4E2M1FN>) -> tensor<13x21x3xbf16>
+  return %0 : tensor<13x21x3xbf16>
+}