[Mlir-commits] [mlir] [mlir][tosa] Add support for MXFP row_gather (PR #207933)
Luke Hutton
llvmlistbot at llvm.org
Mon Jul 13 07:01:50 PDT 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/207933
>From 26ed8a15ec67e0efd28a71b04b42c6941c50c4f2 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Mon, 29 Jun 2026 16:00:18 +0100
Subject: [PATCH] [mlir][tosa] Add support for MXFP row_gather
Adds profile compliance support for MXFP row_gather.
Change-Id: I3a256bd60e857955c5a849b951143fdb2df48fab
---
.../Dialect/Tosa/IR/TosaComplianceData.h.inc | 112 +++++++++++++-----
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 17 ++-
mlir/test/Dialect/Tosa/availability.mlir | 2 +-
mlir/test/Dialect/Tosa/invalid_extension.mlir | 8 ++
mlir/test/Dialect/Tosa/ops.mlir | 8 ++
.../tosa-validation-version-1p0-invalid.mlir | 9 ++
.../tosa-validation-version-1p1-valid.mlir | 18 +++
7 files changed, 138 insertions(+), 36 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index ef644845ed937..3d9f40efac459 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -992,37 +992,87 @@ extensionComplianceMap = {
{{{bf16T, i64T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
allOf}}},
{"tosa.row_gather",
- {
- {{Extension::int16},
- {{{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT}}},
- {{Extension::int64},
- {{{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
- {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
- {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
- {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
- {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
- {{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
- {{fp16T, i64T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
- {{fp32T, i64T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
- {{Extension::int16, Extension::int64},
- {{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT}},
- allOf},
- {{Extension::fp8e4m3},
- {{{fp8e4m3T, i32T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}}},
- {{Extension::fp8e5m2},
- {{{fp8e5m2T, i32T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}}},
- {{Extension::bf16},
- {{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_1_DRAFT}}},
- {{Extension::fp8e4m3, Extension::int64},
- {{{fp8e4m3T, i64T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}},
- allOf},
- {{Extension::fp8e5m2, Extension::int64},
- {{{fp8e5m2T, i64T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}},
- allOf},
- {{Extension::bf16, Extension::int64},
- {{{bf16T, i64T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
- allOf},
- }},
+ {{{Extension::int16},
+ {{{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::int64},
+ {{{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, i64T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, i64T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::int16, Extension::int64},
+ {{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::fp8e4m3},
+ {{{fp8e4m3T, i32T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e5m2},
+ {{{fp8e5m2T, i32T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::bf16},
+ {{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{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::fp8e4m3, Extension::int64},
+ {{{fp8e4m3T, i64T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::fp8e5m2, Extension::int64},
+ {{{fp8e5m2T, i64T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}},
+ 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},
+ 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_block_scaled",
{{{Extension::fp8e4m3},
{{{fp8e4m3T, i32T, i32T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}}},
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 3518a62800963..ffa80a24b231e 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -2596,7 +2596,15 @@ def Tosa_RowGatherOp : Tosa_InferShapedTypeOp<"row_gather", [NoMemoryEffect]> {
list<Availability> availability = [Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
Extension<[Tosa_EXT_FP8E4M3,
Tosa_EXT_FP8E5M2, Tosa_EXT_BF16,
- Tosa_EXT_INT16, Tosa_EXT_INT64]>,
+ Tosa_EXT_INT16,
+ Tosa_EXT_INT64,
+ Tosa_EXT_MX_COMMON,
+ Tosa_EXT_MX_FP4E2M1,
+ Tosa_EXT_MX_FP6E2M3,
+ Tosa_EXT_MX_FP6E3M2,
+ Tosa_EXT_MX_FP8E4M3,
+ Tosa_EXT_MX_FP8E5M2,
+ Tosa_EXT_MX_INT8]>,
];
let hasCanonicalizer = 1;
@@ -2608,6 +2616,8 @@ def Tosa_RowGatherOp : Tosa_InferShapedTypeOp<"row_gather", [NoMemoryEffect]> {
//===----------------------------------------------------------------------===//
// Operator: row_gather_block_scaled
+//
+// Note: This operation is deprecated. It will be removed in the future.
//===----------------------------------------------------------------------===//
def Tosa_RowGatherBlockScaledOp
: Tosa_InferShapedTypeOp<"row_gather_block_scaled", [NoMemoryEffect]> {
@@ -2624,11 +2634,10 @@ def Tosa_RowGatherBlockScaledOp
- non-block-scaled: 1 input value tensor and 1 result tensor
- block-scaled: data + scale input tensors and data + scale result tensors
- This operation follows the TOSA 1.1 draft specification and may evolve as
- the specification is updated.
-
This operation is not pure. Undefined behaviour may occur if the specified
indices are out of range.
+
+ Note: This operation is deprecated. It will be removed in the future.
}];
let arguments = (ins Variadic<Tosa_Tensor3D>:$values,
diff --git a/mlir/test/Dialect/Tosa/availability.mlir b/mlir/test/Dialect/Tosa/availability.mlir
index c358460bbc662..151a9a27302ae 100644
--- a/mlir/test/Dialect/Tosa/availability.mlir
+++ b/mlir/test/Dialect/Tosa/availability.mlir
@@ -600,7 +600,7 @@ func.func @test_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) ->
func.func @test_row_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf32> {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
// CHECK: profiles: [ [pro_int, pro_fp] ]
- // CHECK: extensions: [ [fp8e4m3, fp8e5m2, bf16, int16, int64] ]
+ // CHECK: extensions: [ [fp8e4m3, fp8e5m2, bf16, int16, int64, mx_common, mx_fp4e2m1, mx_fp6e2m3, mx_fp6e3m2, mx_fp8e4m3, mx_fp8e5m2, mx_int8] ]
%0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf32>
return %0 : tensor<13x52x3xf32>
}
diff --git a/mlir/test/Dialect/Tosa/invalid_extension.mlir b/mlir/test/Dialect/Tosa/invalid_extension.mlir
index 5e4dd17b24090..e42650dc88cc1 100644
--- a/mlir/test/Dialect/Tosa/invalid_extension.mlir
+++ b/mlir/test/Dialect/Tosa/invalid_extension.mlir
@@ -257,6 +257,14 @@ func.func @test_row_gather(%arg0: tensor<13x21x3xbf16>, %arg1: tensor<13x26xi32>
return %0 : tensor<13x52x3xbf16>
}
+// -----
+func.func @test_row_gather_mxfp(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op illegal: requires all of [mx_common, mx_fp4e2m1] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
// -----
func.func @test_row_gather_block_scaled(%arg0: tensor<13x21x32xf4E2M1FN>, %arg1: tensor<13x21x1xf8E8M0FNU>, %arg2: tensor<13x26xi32>) -> (tensor<13x52x32xf4E2M1FN>, tensor<13x52x1xf8E8M0FNU>) {
%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 ec3d0c881d857..81137921be19f 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -1438,6 +1438,14 @@ func.func @test_row_gather_f8E5M2(%arg0: tensor<13x21x3xf8E5M2>, %arg1: tensor<1
return %0 : tensor<13x52x3xf8E5M2>
}
+// -----
+// CHECK-LABEL: row_gather_mxfp
+func.func @test_row_gather_mxfp(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
// -----
// CHECK-LABEL: pad_f8E5M2
func.func @test_pad_f8E5M2(%arg0: tensor<13x21x3xf8E5M2>) -> tensor<13x21x3xf8E5M2> {
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 a5e4d4397f290..87d407c11e386 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -97,6 +97,15 @@ func.func @test_row_gather_i8_i32(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26
// -----
+func.func @test_row_gather_mxfp(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ %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) and requires all of [mx_common, mx_fp4e2m1] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
+// -----
+
func.func @test_row_gather_block_scaled_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_block_scaled' 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 e7a4fe6735457..9aee861c6a8d3 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -260,6 +260,24 @@ func.func @test_row_gather_f8e5m2_i32_indices(%arg0: tensor<13x21x3xf8E5M2>, %ar
// -----
+// CHECK-LABEL: test_row_gather_mxfp_i32_indices
+func.func @test_row_gather_mxfp_i32_indices(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<13x26xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
+// -----
+
+// CHECK-LABEL: test_row_gather_mxfp_i64_indices
+func.func @test_row_gather_mxfp_i64_indices(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>, %arg1: tensor<13x26xi64>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>, tensor<13x26xi64>, tensor<1xi32>) -> tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>
+ return %0 : tensor<13x52x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:!tosa.mxint8>>
+}
+
+// -----
+
// CHECK-LABEL: test_row_gather_block_scaled_i8_i32_indices
func.func @test_row_gather_block_scaled_i8_i32_indices(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
More information about the Mlir-commits
mailing list