[Mlir-commits] [mlir] 59a1a32 - [mlir][tosa] Add support for MXFP in SLICE (#220557)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 16 03:08:29 PDT 2026
Author: Philip742
Date: 2026-09-16T11:08:24+01:00
New Revision: 59a1a3289d90b76bedb1013c8ac47b0a9ad5c95f
URL: https://github.com/llvm/llvm-project/commit/59a1a3289d90b76bedb1013c8ac47b0a9ad5c95f
DIFF: https://github.com/llvm/llvm-project/commit/59a1a3289d90b76bedb1013c8ac47b0a9ad5c95f.diff
LOG: [mlir][tosa] Add support for MXFP in SLICE (#220557)
Adds support to profile compliance for MXFP variants of SLICE
Change-Id: I184fbd49ef667d8d4e0ee09a2e65a72095d61daa
---------
Signed-off-by: Philip Wilkinson <philip.wilkinson at arm.com>
Added:
Modified:
mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
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
mlir/test/Dialect/Tosa/verifier.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 e371aecf4e931..89f20a2d178d4 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -4254,7 +4254,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.tile",
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
{{Extension::fp8e4m3},
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index e91c5703e347b..1f113bb29db28 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -2764,6 +2764,21 @@ LogicalResult tosa::SliceOp::verify() {
<< startValues << "]";
}
+ // ERROR_IF(is_block_scale<in_out_t>() && start[rank(shape1) - 1] %
+ // get_innermost_block_size<in_out_t>() != 0);
+ const auto elemType = getElementTypeOrSelf(input.getType());
+ if (const auto blockScaledType = llvm::dyn_cast<BlockScaledType>(elemType)) {
+ const auto startBlock = startValues.back();
+ const auto scaleBlock =
+ BlockShapeAttr::getBlockShapeValue(blockScaledType.getBlockShape());
+ if (startBlock % scaleBlock != 0) {
+ return emitOpError(
+ "expected start innermost block size to match data type "
+ "for block scaled input, got start block=")
+ << startBlock << ", scale block=" << scaleBlock;
+ }
+ }
+
SmallVector<int64_t> sizeValues;
if (!tosa::getConstShapeValues(size.getDefiningOp(), sizeValues))
return success();
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index 13201189942c9..0e81a438a3360 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -2139,3 +2139,13 @@ func.func @test_block_scaled_tile(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOC
%0 = tosa.tile %arg0, %multiples : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>, !tosa.shape<3>) -> tensor<39x21x64x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>
return %0 : tensor<39x21x64x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>
}
+
+// -----
+
+// CHECK-LABEL: test_block_scaled_slice
+func.func @test_block_scaled_slice(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>) -> tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>> {
+ %size = tosa.const_shape {values = dense<[4, 11, 32]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %start = tosa.const_shape {values = dense<[6, 8, 0]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %2 = tosa.slice %arg0, %start, %size : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+ return %2 : tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+}
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 e8d7ca5e57338..7a1e789716924 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -363,3 +363,13 @@ func.func @test_transpose_block_scaled_f6E3M2FN(%input: tensor<29x12x13x96x!tosa
%transpose = tosa.transpose %input { perms = array<i32: 2, 0, 1, 3> } : (tensor<29x12x13x96x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>) -> tensor<13x29x12x96x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>
return %transpose : tensor<13x29x12x96x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>
}
+
+// -----
+
+func.func @test_block_scaled_slice(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>) -> tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>> {
+ %size = tosa.const_shape {values = dense<[4, 11, 32]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %start = tosa.const_shape {values = dense<[6, 8, 0]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ // expected-error at +1 {{'tosa.slice' op illegal: requires specification version compatible with 1.1.draft (got 1.0) and requires all of [mx_common, mx_fp8e4m3] profiles/extensions to be specified in the target environment}}
+ %2 = tosa.slice %arg0, %start, %size : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+ return %2 : tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+}
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 d60fee0255ac9..755e4700350fb 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -710,3 +710,13 @@ func.func @test_transpose_block_scaled_f6E3M2FN(%input: tensor<29x12x13x96x!tosa
%transpose = tosa.transpose %input { perms = array<i32: 2, 0, 1, 3> } : (tensor<29x12x13x96x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>) -> tensor<13x29x12x96x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>
return %transpose : tensor<13x29x12x96x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>
}
+
+// -----
+
+// CHECK-LABEL: test_block_scaled_slice
+func.func @test_block_scaled_slice(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>) -> tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>> {
+ %size = tosa.const_shape {values = dense<[4, 11, 32]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %start = tosa.const_shape {values = dense<[6, 8, 0]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %2 = tosa.slice %arg0, %start, %size : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+ return %2 : tensor<4x11x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+}
diff --git a/mlir/test/Dialect/Tosa/verifier.mlir b/mlir/test/Dialect/Tosa/verifier.mlir
index cd71fffbe9962..702a5742f5217 100644
--- a/mlir/test/Dialect/Tosa/verifier.mlir
+++ b/mlir/test/Dialect/Tosa/verifier.mlir
@@ -2437,3 +2437,13 @@ func.func @test_transpose_block_scaled_illegal_perms(%input: tensor<29x12x32x96x
%transpose = tosa.transpose %input { perms = array<i32: 3, 0, 1, 2> } : (tensor<29x12x32x96x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>) -> tensor<96x29x12x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>
return %transpose : tensor<96x29x12x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E3M2FN>>
}
+
+// -----
+
+func.func @test_slice_block_scaled_invalid_block_size(%arg0: tensor<4x4x64x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>) -> tensor<1x1x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+ %start = tosa.const_shape {values = dense<[0, 0, 1]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %size = tosa.const_shape {values = dense<[1, 1, 32]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ // expected-error at +1 {{'tosa.slice' op expected start innermost block size to match data type for block scaled input, got start block=1, scale block=32}}
+ %0 = tosa.slice %arg0, %start, %size {input_unsigned = false} : (tensor<4x4x64x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, !tosa.shape<3>, !tosa.shape<3>) -> tensor<1x1x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+ return %0 : tensor<1x1x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
More information about the Mlir-commits
mailing list