[Mlir-commits] [mlir] 3aeb6c1 - [mlir][spirv] Re-enable bf16/fp8 for vector composite ops (#204848)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 22 02:02:07 PDT 2026
Author: Dmitriy Smirnov
Date: 2026-06-22T10:02:02+01:00
New Revision: 3aeb6c1f3322b57fe102b7c87d8fd17f87454195
URL: https://github.com/llvm/llvm-project/commit/3aeb6c1f3322b57fe102b7c87d8fd17f87454195
DIFF: https://github.com/llvm/llvm-project/commit/3aeb6c1f3322b57fe102b7c87d8fd17f87454195.diff
LOG: [mlir][spirv] Re-enable bf16/fp8 for vector composite ops (#204848)
Allow bf16 and fp8 vector element types in VectorExtractDynamic,
VectorInsertDynamic, and VectorShuffle.
Added:
Modified:
mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td
mlir/test/Dialect/SPIRV/IR/composite-ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index 316040025810d..68a0ee470709d 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4314,6 +4314,10 @@ def SPIRV_FloatVector : VectorOfRankAndLengthAndType<[1], [2, 3, 4, 8, 16],
def SPIRV_ExtendedFloatVector : VectorOfRankAndLengthAndType<[1], [2, 3, 4, 8, 16],
[SPIRV_BFloat16KHR, SPIRV_Float8E4M3EXT,
SPIRV_Float8E5M2EXT]>;
+def SPIRV_ExtendedVector : VectorOfRankAndLengthAndType<[1], [2, 3, 4, 8, 16],
+ [SPIRV_Bool, SPIRV_Integer, SPIRV_Float,
+ SPIRV_BFloat16KHR, SPIRV_Float8E4M3EXT,
+ SPIRV_Float8E5M2EXT]>;
def SPIRV_AnyFloatVector : AnyTypeOf<[SPIRV_FloatVector, SPIRV_ExtendedFloatVector]>;
// Component type check is done in the type parser for the following SPIR-V
// dialect-specific types so we use "Any" here.
@@ -4344,7 +4348,9 @@ def SPIRV_AnyTensorArm : DialectType<SPIRV_Dialect, SPIRV_IsTensorArmType,
"any SPIR-V tensorArm type">;
def SPIRV_Numerical : AnyTypeOf<[SPIRV_Integer, SPIRV_Float]>;
+def SPIRV_ExtendedNumerical : AnyTypeOf<[SPIRV_Integer, SPIRV_Float, SPIRV_BFloat16KHR, SPIRV_Float8E4M3EXT, SPIRV_Float8E5M2EXT]>;
def SPIRV_Scalar : AnyTypeOf<[SPIRV_Numerical, SPIRV_Bool]>;
+def SPIRV_ExtendedScalar : AnyTypeOf<[SPIRV_ExtendedNumerical, SPIRV_Bool]>;
def SPIRV_Aggregate : AnyTypeOf<[SPIRV_AnyArray, SPIRV_AnyRTArray, SPIRV_AnyStruct]>;
def SPIRV_Composite :
AnyTypeOf<[SPIRV_Vector, SPIRV_ExtendedFloatVector, SPIRV_AnyArray,
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td
index 981131484498d..73f93522d142b 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCompositeOps.td
@@ -207,12 +207,12 @@ def SPIRV_VectorExtractDynamicOp : SPIRV_Op<"VectorExtractDynamic", [
}];
let arguments = (ins
- SPIRV_Vector:$vector,
+ SPIRV_ExtendedVector:$vector,
SPIRV_Integer:$index
);
let results = (outs
- SPIRV_Scalar:$result
+ SPIRV_ExtendedScalar:$result
);
let hasVerifier = 0;
@@ -260,13 +260,13 @@ def SPIRV_VectorInsertDynamicOp : SPIRV_Op<"VectorInsertDynamic", [
}];
let arguments = (ins
- SPIRV_Vector:$vector,
- SPIRV_Scalar:$component,
+ SPIRV_ExtendedVector:$vector,
+ SPIRV_ExtendedScalar:$component,
SPIRV_Integer:$index
);
let results = (outs
- SPIRV_Vector:$result
+ SPIRV_ExtendedVector:$result
);
let hasVerifier = 0;
@@ -320,13 +320,13 @@ def SPIRV_VectorShuffleOp : SPIRV_Op<"VectorShuffle", [
}];
let arguments = (ins
- SPIRV_Vector:$vector1,
- SPIRV_Vector:$vector2,
+ SPIRV_ExtendedVector:$vector1,
+ SPIRV_ExtendedVector:$vector2,
I32ArrayAttr:$components
);
let results = (outs
- SPIRV_Vector:$result
+ SPIRV_ExtendedVector:$result
);
let assemblyFormat = [{
diff --git a/mlir/test/Dialect/SPIRV/IR/composite-ops.mlir b/mlir/test/Dialect/SPIRV/IR/composite-ops.mlir
index 478884f8dfd08..22411fd7294cf 100644
--- a/mlir/test/Dialect/SPIRV/IR/composite-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/composite-ops.mlir
@@ -297,6 +297,24 @@ func.func @vector_dynamic_extract(%vec: vector<4xf32>, %id : i32) -> f32 {
return %0 : f32
}
+func.func @vector_dynamic_extract_bf16(%vec: vector<4xbf16>, %id : i32) -> bf16 {
+ // CHECK: spirv.VectorExtractDynamic %{{.*}}[%{{.*}}] : vector<4xbf16>, i32
+ %0 = spirv.VectorExtractDynamic %vec[%id] : vector<4xbf16>, i32
+ return %0 : bf16
+}
+
+func.func @vector_dynamic_extract_f8E5M2(%vec: vector<4xf8E5M2>, %id : i32) -> f8E5M2 {
+ // CHECK: spirv.VectorExtractDynamic %{{.*}}[%{{.*}}] : vector<4xf8E5M2>, i32
+ %0 = spirv.VectorExtractDynamic %vec[%id] : vector<4xf8E5M2>, i32
+ return %0 : f8E5M2
+}
+
+func.func @vector_dynamic_extract_f8E4M3FN(%vec: vector<4xf8E4M3FN>, %id : i32) -> f8E4M3FN {
+ // CHECK: spirv.VectorExtractDynamic %{{.*}}[%{{.*}}] : vector<4xf8E4M3FN>, i32
+ %0 = spirv.VectorExtractDynamic %vec[%id] : vector<4xf8E4M3FN>, i32
+ return %0 : f8E4M3FN
+}
+
//===----------------------------------------------------------------------===//
// spirv.VectorInsertDynamic
//===----------------------------------------------------------------------===//
@@ -307,6 +325,24 @@ func.func @vector_dynamic_insert(%val: f32, %vec: vector<4xf32>, %id : i32) -> v
return %0 : vector<4xf32>
}
+func.func @vector_dynamic_insert_bf16(%val: bf16, %vec: vector<4xbf16>, %id : i32) -> vector<4xbf16> {
+ // CHECK: spirv.VectorInsertDynamic %{{.*}}, %{{.*}}[%{{.*}}] : vector<4xbf16>, i32
+ %0 = spirv.VectorInsertDynamic %val, %vec[%id] : vector<4xbf16>, i32
+ return %0 : vector<4xbf16>
+}
+
+func.func @vector_dynamic_insert_f8E5M2(%val: f8E5M2, %vec: vector<4xf8E5M2>, %id : i32) -> vector<4xf8E5M2> {
+ // CHECK: spirv.VectorInsertDynamic %{{.*}}, %{{.*}}[%{{.*}}] : vector<4xf8E5M2>, i32
+ %0 = spirv.VectorInsertDynamic %val, %vec[%id] : vector<4xf8E5M2>, i32
+ return %0 : vector<4xf8E5M2>
+}
+
+func.func @vector_dynamic_insert_f8E4M3FN(%val: f8E4M3FN, %vec: vector<4xf8E4M3FN>, %id : i32) -> vector<4xf8E4M3FN> {
+ // CHECK: spirv.VectorInsertDynamic %{{.*}}, %{{.*}}[%{{.*}}] : vector<4xf8E4M3FN>, i32
+ %0 = spirv.VectorInsertDynamic %val, %vec[%id] : vector<4xf8E4M3FN>, i32
+ return %0 : vector<4xf8E4M3FN>
+}
+
// -----
//===----------------------------------------------------------------------===//
@@ -319,6 +355,24 @@ func.func @vector_shuffle(%vector1: vector<4xf32>, %vector2: vector<2xf32>) -> v
return %0: vector<3xf32>
}
+func.func @vector_shuffle_bf16(%vector1: vector<4xbf16>, %vector2: vector<2xbf16>) -> vector<3xbf16> {
+ // CHECK: %{{.+}} = spirv.VectorShuffle [1 : i32, 3 : i32, -1 : i32] %{{.+}}, %arg1 : vector<4xbf16>, vector<2xbf16> -> vector<3xbf16>
+ %0 = spirv.VectorShuffle [1: i32, 3: i32, 0xffffffff: i32] %vector1, %vector2 : vector<4xbf16>, vector<2xbf16> -> vector<3xbf16>
+ return %0: vector<3xbf16>
+}
+
+func.func @vector_shuffle_f8E5M2(%vector1: vector<4xf8E5M2>, %vector2: vector<2xf8E5M2>) -> vector<3xf8E5M2> {
+ // CHECK: %{{.+}} = spirv.VectorShuffle [1 : i32, 3 : i32, -1 : i32] %{{.+}}, %arg1 : vector<4xf8E5M2>, vector<2xf8E5M2> -> vector<3xf8E5M2>
+ %0 = spirv.VectorShuffle [1: i32, 3: i32, 0xffffffff: i32] %vector1, %vector2 : vector<4xf8E5M2>, vector<2xf8E5M2> -> vector<3xf8E5M2>
+ return %0: vector<3xf8E5M2>
+}
+
+func.func @vector_shuffle_f8E4M3FN(%vector1: vector<4xf8E4M3FN>, %vector2: vector<2xf8E4M3FN>) -> vector<3xf8E4M3FN> {
+ // CHECK: %{{.+}} = spirv.VectorShuffle [1 : i32, 3 : i32, -1 : i32] %{{.+}}, %arg1 : vector<4xf8E4M3FN>, vector<2xf8E4M3FN> -> vector<3xf8E4M3FN>
+ %0 = spirv.VectorShuffle [1: i32, 3: i32, 0xffffffff: i32] %vector1, %vector2 : vector<4xf8E4M3FN>, vector<2xf8E4M3FN> -> vector<3xf8E4M3FN>
+ return %0: vector<3xf8E4M3FN>
+}
+
// -----
func.func @vector_shuffle_extra_selector(%vector1: vector<4xf32>, %vector2: vector<2xf32>) -> vector<3xf32> {
More information about the Mlir-commits
mailing list