[Mlir-commits] [mlir] [mlir][tosa] Add support for mxfp IDENTITY (PR #220624)

[Yuvaraj Venkatesh]

https://github.com/Yuvaraj-Venkatesh updated https://github.com/llvm/llvm-project/pull/220624

>From efedbb27d17ea7f39fe1d8965f1a8c5d209807c8 Mon Sep 17 00:00:00 2001
From: Yuvaraj Venkatesh <yuvaraj.venkatesh at arm.com>
Date: Wed, 2 Sep 2026 11:45:12 +0000
Subject: [PATCH] [mlir][tosa] Add support for mxfp IDENTITY

Tosa specification changes in https://github.com/arm/tosa-specification/pull/61

Signed-off-by: Yuvaraj Venkatesh <yuvaraj.venkatesh at arm.com>
Change-Id: Ib938e4915c14e14aeb465944875b2fed7e5e3698
---
 .../Dialect/Tosa/IR/TosaComplianceData.h.inc  | 26 ++++++++++++++++++-
 mlir/test/Dialect/Tosa/ops.mlir               |  8 ++++++
 .../tosa-validation-version-1p0-invalid.mlir  |  8 ++++++
 .../tosa-validation-version-1p1-valid.mlir    |  8 ++++++
 4 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index c657d67cb5873..79567c65594ab 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -4755,7 +4755,31 @@ extensionComplianceMap = {
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
-      {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+      {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}},
+      {{Extension::mx_common, Extension::mx_fp8e4m3},
+       {{{bs32_fp8ue8m0_fp8e4m3T, bs32_fp8ue8m0_fp8e4m3T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::mx_common, Extension::mx_fp8e5m2},
+       {{{bs32_fp8ue8m0_fp8e5m2T, bs32_fp8ue8m0_fp8e5m2T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::mx_common, Extension::mx_fp6e3m2},
+       {{{bs32_fp8ue8m0_fp6e3m2T, bs32_fp8ue8m0_fp6e3m2T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::mx_common, Extension::mx_fp6e2m3},
+       {{{bs32_fp8ue8m0_fp6e2m3T, bs32_fp8ue8m0_fp6e2m3T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::mx_common, Extension::mx_fp4e2m1},
+       {{{bs32_fp8ue8m0_fp4e2m1T, bs32_fp8ue8m0_fp4e2m1T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::mx_common, Extension::mx_int8},
+       {{{bs32_fp8ue8m0_mxint8T, bs32_fp8ue8m0_mxint8T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf}}},
     {"tosa.variable",
      {{{Extension::variable},
        {{{i8T}, SpecificationVersion::V_1_0},
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index 45fb083e32047..f9d0dcc0e51da 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -1333,6 +1333,14 @@ func.func @test_identity(%arg0: tensor<13x21x3xi32>) -> tensor<13x21x3xi32> {
   return %0 : tensor<13x21x3xi32>
 }
 
+// -----
+
+// CHECK-LABEL: test_block_scaled_identity
+func.func @test_block_scaled_identity(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+  %0 = tosa.identity %arg0 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+  return %0 : tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
 // -----
 // CHECK-LABEL: cond_if
 func.func @test_cond_if(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
index 04f432ef0f04b..36787506a3bc3 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -218,6 +218,14 @@ func.func @test_dyanmic_dims(%arg0: tensor<1x8x16xi8>) {
 
 // -----
 
+func.func @test_block_scaled_identity(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>> {
+  // expected-error at +1 {{'tosa.identity' op illegal: requires specification version compatible with 1.1.draft (got 1.0) and requires all of [mx_common, mx_fp6e2m3] profiles/extensions to be specified in the target environment}}
+  %0 = tosa.identity %arg0 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>
+  return %0 : tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>
+}
+
+// -----
+
 func.func @test_matmul_t_block_scaled(%arg0: tensor<4x8x32xf8E4M3FN>, %arg1: tensor<4x8x1xf8E8M0FNU>, %arg2: tensor<4x16x32xf8E4M3FN>, %arg3: tensor<4x16x1xf8E8M0FNU>) -> tensor<4x8x16xf32> {
   // expected-error at +1 {{'tosa.matmul_t_block_scaled' op illegal: requires specification version compatible with 1.1.draft (got 1.0) and requires any of [mxfp] profiles/extensions to be specified in the target environment}}
   %0 = tosa.matmul_t_block_scaled %arg0, %arg1, %arg2, %arg3 {block_size = #tosa.block_size<BLOCK_SIZE_32>} : (tensor<4x8x32xf8E4M3FN>, tensor<4x8x1xf8E8M0FNU>, tensor<4x16x32xf8E4M3FN>, tensor<4x16x1xf8E8M0FNU>) -> tensor<4x8x16xf32>
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
index 329510db46c68..21b713eab9f90 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -647,6 +647,14 @@ func.func @test_block_scaled_tile(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOC
 
 // -----
 
+// CHECK-LABEL: test_block_scaled_identity
+func.func @test_block_scaled_identity(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>> {
+  %0 = tosa.identity %arg0 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>
+  return %0 : tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>
+}
+
+// -----
+
 // CHECK-LABEL: test_resize_fp8
 func.func @test_resize_fp8(%arg0: tensor<1x32x32x8xf8E4M3FN>, %arg1: tensor<1x32x32x8xf8E5M2>) {
   %scale = tosa.const_shape { values = dense<[4, 2, 4, 2]> : tensor<4xindex> } : () -> !tosa.shape<4>