[Mlir-commits] [mlir] 956278f - [mlir][tosa] Add MXFP support for GATHER (#227710)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 1 05:34:24 PDT 2026
Author: Maddy Dunn
Date: 2026-10-01T13:34:13+01:00
New Revision: 956278f0f570085eb9a2bc2a1de32d142aeb233f
URL: https://github.com/llvm/llvm-project/commit/956278f0f570085eb9a2bc2a1de32d142aeb233f
DIFF: https://github.com/llvm/llvm-project/commit/956278f0f570085eb9a2bc2a1de32d142aeb233f.diff
LOG: [mlir][tosa] Add MXFP support for GATHER (#227710)
Change-Id: I007889a6d934c8588feae7d2c5a023df60c85ca5
Signed-off-by: Madeleine Dunn <madeleine.dunn at arm.com>
Added:
Modified:
mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
mlir/test/Dialect/Tosa/invalid_extension.mlir
mlir/test/Dialect/Tosa/ops.mlir
mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 216dc65da9610..5dcfe3a737e2d 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -4372,6 +4372,30 @@ extensionComplianceMap = {
{{{fp8e5m2T, i32T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
{{Extension::bf16},
{{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_0}}},
+ {{Extension::mx_common, Extension::mx_fp8e4m3},
+ {{{bs32_fp8ue8m0_fp8e4m3T, i32T, bs32_fp8ue8m0_fp8e4m3T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::mx_common, Extension::mx_fp8e5m2},
+ {{{bs32_fp8ue8m0_fp8e5m2T, i32T, bs32_fp8ue8m0_fp8e5m2T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::mx_common, Extension::mx_fp6e3m2},
+ {{{bs32_fp8ue8m0_fp6e3m2T, i32T, bs32_fp8ue8m0_fp6e3m2T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::mx_common, Extension::mx_fp6e2m3},
+ {{{bs32_fp8ue8m0_fp6e2m3T, i32T, bs32_fp8ue8m0_fp6e2m3T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::mx_common, Extension::mx_fp4e2m1},
+ {{{bs32_fp8ue8m0_fp4e2m1T, i32T, bs32_fp8ue8m0_fp4e2m1T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::mx_common, Extension::mx_int8},
+ {{{bs32_fp8ue8m0_mxint8T, i32T, bs32_fp8ue8m0_mxint8T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
{{Extension::int64},
{{{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
@@ -4388,6 +4412,31 @@ extensionComplianceMap = {
allOf},
{{Extension::bf16, Extension::int64},
{{{bf16T, i64T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::int64, Extension::mx_common, Extension::mx_fp8e4m3},
+ {{{bs32_fp8ue8m0_fp8e4m3T, i64T, bs32_fp8ue8m0_fp8e4m3T,
+ bs32_fp8ue8m0_fp8e4m3T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::int64, Extension::mx_common, Extension::mx_fp8e5m2},
+ {{{bs32_fp8ue8m0_fp8e5m2T, i64T, bs32_fp8ue8m0_fp8e5m2T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::int64, Extension::mx_common, Extension::mx_fp6e3m2},
+ {{{bs32_fp8ue8m0_fp6e3m2T, i64T, bs32_fp8ue8m0_fp6e3m2T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::int64, Extension::mx_common, Extension::mx_fp6e2m3},
+ {{{bs32_fp8ue8m0_fp6e2m3T, i64T, bs32_fp8ue8m0_fp6e2m3T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::int64, Extension::mx_common, Extension::mx_fp4e2m1},
+ {{{bs32_fp8ue8m0_fp4e2m1T, i64T, bs32_fp8ue8m0_fp4e2m1T},
+ SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::int64, Extension::mx_common, Extension::mx_int8},
+ {{{bs32_fp8ue8m0_mxint8T, i64T, bs32_fp8ue8m0_mxint8T},
+ SpecificationVersion::V_1_1_DRAFT}},
allOf}}},
{"tosa.row_gather",
{{{Extension::int16},
diff --git a/mlir/test/Dialect/Tosa/invalid_extension.mlir b/mlir/test/Dialect/Tosa/invalid_extension.mlir
index f4e70710dc320..16462130bb23b 100644
--- a/mlir/test/Dialect/Tosa/invalid_extension.mlir
+++ b/mlir/test/Dialect/Tosa/invalid_extension.mlir
@@ -275,6 +275,13 @@ func.func @test_gather(%arg0: tensor<13x21x3xbf16>, %arg1: tensor<13x26xi32>) ->
return %0 : tensor<13x26x3xbf16>
}
+// -----
+func.func @test_gather_mxfp(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ // expected-error at +1 {{'tosa.gather' op illegal: requires all of [mx_common, mx_fp4e2m1] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
// -----
func.func @test_row_gather(%arg0: tensor<13x21x3xbf16>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xbf16> {
%row_count = "tosa.const"() <{values = dense<2> : tensor<1xi32>}> : () -> tensor<1xi32>
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index 11b6f275bd8b8..eac2f58772d8d 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -1749,6 +1749,13 @@ func.func @test_gather_f8E4M3FN(%arg0: tensor<13x21x3xf8E4M3FN>, %arg1: tensor<1
return %0 : tensor<13x26x3xf8E4M3FN>
}
+// -----
+// CHECK-LABEL: test_block_scaled_gather
+func.func @test_block_scaled_gather(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
// -----
// CHECK-LABEL: scatter_f8E4M3FN
func.func @test_scatter_f8E4M3FN(%arg0: tensor<13x29x3xf8E4M3FN>, %arg1: tensor<13x26xi32>, %arg2: tensor<13x26x3xf8E4M3FN>) -> tensor<13x29x3xf8E4M3FN> {
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 ff9186b19eff6..681025deaebb8 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -108,6 +108,14 @@ func.func @test_gather_bool_i32(%arg0: tensor<13x21x3xi1>, %arg1: tensor<13x26xi
// -----
+func.func @test_gather_mxfp(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ // expected-error at +1 {{'tosa.gather' op illegal: requires specification version compatible with 1.1.draft (got 1.0) and requires all of [mx_common, mx_fp4e2m1] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
+// -----
+
func.func @test_row_gather_i8_i32(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
%row_count = "tosa.const"() <{values = dense<2> : tensor<1xi32>}> : () -> tensor<1xi32>
// expected-error at +1 {{'tosa.row_gather' op illegal: requires specification version compatible with 1.1.draft (got 1.0) OR requires specification version compatible with 1.1.draft (got 1.0) to be specified in the target environment}}
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 c51f7c055a7fa..db25c3b1e685b 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -356,6 +356,22 @@ func.func @test_row_gather_mxfp_i64_indices(%arg0: tensor<13x21x32x!tosa.block_s
// -----
+// CHECK-LABEL: test_gather_mxfp_i32_indices
+func.func @test_gather_mxfp_i32_indices(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
+// -----
+
+// CHECK-LABEL: test_gather_mxfp_i64_indices
+func.func @test_gather_mxfp_i64_indices(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>, %arg1: tensor<13x26xi64>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>> {
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>, tensor<13x26xi64>) -> tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>
+ return %0 : tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>
+}
+
+// -----
+
// CHECK-LABEL: test_scatter_mxfp_i32_indices
func.func @test_scatter_mxfp_i32_indices(%arg0: tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>, %arg2: tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
%0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>, tensor<13x26x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
More information about the Mlir-commits
mailing list