[Mlir-commits] [mlir] 706db1e - [mlir][tosa] Add support for MXFP row_gather (#207933)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 15 06:11:53 PDT 2026
Author: Luke Hutton
Date: 2026-07-15T14:11:47+01:00
New Revision: 706db1e0959f45d6392edd2ad196a4ea38bb99d3
URL: https://github.com/llvm/llvm-project/commit/706db1e0959f45d6392edd2ad196a4ea38bb99d3
DIFF: https://github.com/llvm/llvm-project/commit/706db1e0959f45d6392edd2ad196a4ea38bb99d3.diff
LOG: [mlir][tosa] Add support for MXFP row_gather (#207933)
Adds profile compliance support for MXFP row_gather.
Added:
Modified:
mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
mlir/test/Dialect/Tosa/availability.mlir
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 4e23d8dd2fbbb..9bbbed98402f0 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -1475,37 +1475,82 @@ 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},
+ {{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 24de011e90c92..46ddffd245a15 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -2603,7 +2603,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;
@@ -2615,6 +2623,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]> {
@@ -2631,11 +2641,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 bc014fb9b928e..9d982395b9407 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 c68487169da63..ffd60ba5e62ad 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 74ba92bdb25ab..643c04588b6ef 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -1465,6 +1465,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 bea262a11a391..a46da2ab9abd3 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 facb49379248a..70a65c3448f39 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -303,6 +303,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