[Mlir-commits] [mlir] [mlir][spirv] Add SPV_EXT_FP8 type support to SPIR-V TOSA ops (PR #193199)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Apr 21 04:36:49 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Davide Grohmann (davidegrohmann)
<details>
<summary>Changes</summary>
Add SPV_EXT_FP8 support for SPIR-V TOSA ops by updating the shared type definitions and extending op constraints for the newly supported element types.
Also update verifier coverage to reflect the new constraints:
- refresh existing negative tests whose diagnostics now list FP8 types
- add negative tests for SPV_EXT_FP8-specific output, weight, accumulator, and cast restrictions
---
Patch is 59.11 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/193199.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td (+76-52)
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaTypes.td (+20-13)
- (modified) mlir/test/Dialect/SPIRV/IR/tosa-ops-verification.mlir (+279-5)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td
index c873e3069733c..4e3689e6001f6 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td
@@ -156,14 +156,20 @@ class SPIRV_TosaConvolutionOp<string mnemonic, int opcode, list<Trait> traits =
TypeConstraintImplicationOn<"input", BF16, "output", [BF16]>,
TypeConstraintImplicationOn<"input", F16, "output", [F16]>,
TypeConstraintImplicationOn<"input", F32, "output", [F32]>,
+ TypeConstraintImplicationOn<"input", F8E4M3FN, "output", [F16]>,
+ TypeConstraintImplicationOn<"input", F8E5M2, "output", [F16]>,
TypeConstraintImplicationOn<"input", BF16, "weight", [BF16]>,
TypeConstraintImplicationOn<"input", F16, "weight", [F16]>,
TypeConstraintImplicationOn<"input", F32, "weight", [F32]>,
+ TypeConstraintImplicationOn<"input", F8E4M3FN, "weight", [F8E4M3FN]>,
+ TypeConstraintImplicationOn<"input", F8E5M2, "weight", [F8E5M2]>,
TypeImpliesAccType<"input", I8, ["INT32"]>,
TypeImpliesAccType<"input", I16, ["INT48"]>,
TypeImpliesAccType<"input", F16, ["FP16", "FP32"]>,
TypeImpliesAccType<"input", BF16, ["FP32"]>,
TypeImpliesAccType<"input", F32, ["FP32"]>,
+ TypeImpliesAccType<"input", F8E4M3FN, ["FP16"]>,
+ TypeImpliesAccType<"input", F8E5M2, ["FP16"]>,
AllElementTypesMatch<["bias", "output"]>,
AllElementTypesMatch<["input", "input_zp"]>,
AllElementTypesMatch<["weight", "weight_zp"]>])> {
@@ -249,7 +255,7 @@ def SPIRV_TosaArgMaxOp : SPIRV_TosaOpWithResult<"ArgMax", 0, [Pure,
let arguments = (ins
SPIRV_TensorArmAxisAttr: $axis,
SPIRV_TosaExtNaNPropagationModeAttr: $nan_mode,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm: $input
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $input
);
let results = (outs
@@ -277,6 +283,8 @@ def SPIRV_TosaAvgPool2DOp : SPIRV_TosaOpWithResult<"AvgPool2D", 1, [NoMemoryEffe
TypeImpliesAccType<"input", F16, ["FP16", "FP32"]>,
TypeImpliesAccType<"input", BF16, ["FP32"]>,
TypeImpliesAccType<"input", F32, ["FP32"]>,
+ TypeImpliesAccType<"input", F8E4M3FN, ["FP16"]>,
+ TypeImpliesAccType<"input", F8E5M2, ["FP16"]>,
AllElementTypesMatch<["input", "input_zp", "output", "output_zp"]>]> {
let summary = "Performs average pooling on the input.";
@@ -304,13 +312,13 @@ def SPIRV_TosaAvgPool2DOp : SPIRV_TosaOpWithResult<"AvgPool2D", 1, [NoMemoryEffe
SPIRV_I32_1DTensorArmOfLength2Attr: $stride,
SPIRV_I32_1DTensorArmOfLength4Attr: $pad,
SPIRV_TosaExtAccTypeAttr: $acc_type,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D: $input,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $input_zp,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $output_zp
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $input,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $input_zp,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $output_zp
);
let results = (outs
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D: $output
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $output
);
let assemblyFormat = [{
@@ -361,11 +369,11 @@ def SPIRV_TosaConv2DOp : SPIRV_TosaConvolutionOp<"Conv2D", 2> {
SPIRV_I32_1DTensorArmOfLength2Attr: $dilation,
SPIRV_TosaExtAccTypeAttr: $acc_type,
SPIRV_BoolConstAttr: $local_bound,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D: $input,
- SPIRV_I8OrF16OrF32OrBF16_TensorArm4D: $weight,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $input,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $weight,
SPIRV_I32OrI64OrF16OrF32OrBF16_TensorArm1D: $bias,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $input_zp,
- SPIRV_I8OrF16OrF32OrBF16_1DTensorArmOfLength1: $weight_zp
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $input_zp,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $weight_zp
);
let results = (outs
@@ -416,11 +424,11 @@ def SPIRV_TosaConv3DOp : SPIRV_TosaConvolutionOp<"Conv3D", 3> {
SPIRV_I32_1DTensorArmOfLength3Attr: $dilation,
SPIRV_TosaExtAccTypeAttr: $acc_type,
SPIRV_BoolConstAttr: $local_bound,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm5D: $input,
- SPIRV_I8OrF16OrF32OrBF16_TensorArm5D: $weight,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm5D: $input,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm5D: $weight,
SPIRV_I32OrI64OrF16OrF32OrBF16_TensorArm1D: $bias,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $input_zp,
- SPIRV_I8OrF16OrF32OrBF16_1DTensorArmOfLength1: $weight_zp
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $input_zp,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $weight_zp
);
let results = (outs
@@ -472,11 +480,11 @@ def SPIRV_TosaDepthwiseConv2DOp : SPIRV_TosaConvolutionOp<"DepthwiseConv2D", 4>
SPIRV_I32_1DTensorArmOfLength2Attr: $dilation,
SPIRV_TosaExtAccTypeAttr: $acc_type,
SPIRV_BoolConstAttr: $local_bound,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D: $input,
- SPIRV_I8OrF16OrF32OrBF16_TensorArm4D: $weight,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $input,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $weight,
SPIRV_I32OrI64OrF16OrF32OrBF16_TensorArm1D: $bias,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $input_zp,
- SPIRV_I8OrF16OrF32OrBF16_1DTensorArmOfLength1: $weight_zp
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $input_zp,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $weight_zp
);
let results = (outs
@@ -557,6 +565,8 @@ def SPIRV_TosaMatMulOp : SPIRV_TosaOpWithResult<"MatMul", 6, [NoMemoryEffect,
TypeConstraintImplicationOn<"A", BF16, "output", [F32]>,
TypeConstraintImplicationOn<"A", F16, "output", [F16, F32]>,
TypeConstraintImplicationOn<"A", F32, "output", [F32]>,
+ TypeConstraintImplicationOn<"A", F8E4M3FN, "output", [F16]>,
+ TypeConstraintImplicationOn<"A", F8E5M2, "output", [F16]>,
AllElementTypesMatch<["A", "A_zp", "B", "B_zp"]>]> {
let summary = "Matrix Multiplication operator.";
@@ -579,10 +589,10 @@ def SPIRV_TosaMatMulOp : SPIRV_TosaOpWithResult<"MatMul", 6, [NoMemoryEffect,
}];
let arguments = (ins
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm3D: $A,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm3D: $B,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $A_zp,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $B_zp
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D: $A,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D: $B,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $A_zp,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $B_zp
);
let results = (outs
@@ -634,11 +644,11 @@ def SPIRV_TosaMaxPool2DOp : SPIRV_TosaOpWithResult<"MaxPool2D", 7, [Pure,
SPIRV_I32_1DTensorArmOfLength2Attr: $stride,
SPIRV_I32_1DTensorArmOfLength4Attr: $pad,
SPIRV_TosaExtNaNPropagationModeAttr: $nan_mode,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D: $input
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $input
);
let results = (outs
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D: $output
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $output
);
let assemblyFormat = [{
@@ -734,11 +744,11 @@ def SPIRV_TosaTransposeConv2DOp : SPIRV_TosaConvolutionOp<"TransposeConv2D", 9>
SPIRV_I32_1DTensorArmOfLength2Attr: $stride,
SPIRV_TosaExtAccTypeAttr: $acc_type,
SPIRV_BoolConstAttr: $local_bound,
- SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D: $input,
- SPIRV_I8OrF16OrF32OrBF16_TensorArm4D: $weight,
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $input,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D: $weight,
SPIRV_I32OrI64OrF16OrF32OrBF16_TensorArm1D: $bias,
- SPIRV_I8OrI16OrF16OrF32OrBF16_1DTensorArmOfLength1: $input_zp,
- SPIRV_I8OrF16OrF32OrBF16_1DTensorArmOfLength1: $weight_zp
+ SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $input_zp,
+ SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $weight_zp
);
let results = (outs
@@ -2167,11 +2177,11 @@ def SPIRV_TosaConcatOp : SPIRV_TosaOpWithResult<"Concat", 54, [Pure,
let arguments = (ins
SPIRV_TensorArmAxisAttr: $axis,
- Variadic<SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm>: $input1
+ Variadic<SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm>: $input1
);
let results = (outs
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
@@ -2214,13 +2224,13 @@ def SPIRV_TosaPadOp : SPIRV_TosaOpWithResult<"Pad", 55, [Pure,
}];
let arguments = (ins
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $input1,
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $input1,
SPIRV_I32_1DTensorArmOfEvenLength2To12: $padding,
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_1DTensorArmOfLength1: $pad_const
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_1DTensorArmOfLength1: $pad_const
);
let results = (outs
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
@@ -2267,12 +2277,12 @@ def SPIRV_TosaReshapeOp : SPIRV_TosaOpWithResult<"Reshape", 56, [Pure,
}];
let arguments = (ins
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $input1,
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $input1,
SPIRV_I32_1DTensorArmOfLength1To6: $shape
);
let results = (outs
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
@@ -2315,11 +2325,11 @@ def SPIRV_TosaReverseOp : SPIRV_TosaOpWithResult<"Reverse", 57, [Pure,
let arguments = (ins
SPIRV_TensorArmAxisAttr: $axis,
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $input1
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $input1
);
let results = (outs
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
@@ -2362,13 +2372,13 @@ def SPIRV_TosaSliceOp : SPIRV_TosaOpWithResult<"Slice", 58, [Pure,
}];
let arguments = (ins
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $input1,
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $input1,
SPIRV_I32_1DTensorArmOfLength1To6: $start,
SPIRV_I32_1DTensorArmOfLength1To6: $size
);
let results = (outs
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
@@ -2416,12 +2426,12 @@ def SPIRV_TosaTileOp : SPIRV_TosaOpWithResult<"Tile", 59, [Pure,
}];
let arguments = (ins
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $input1,
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $input1,
SPIRV_I32_1DTensorArmOfLength1To6: $multiples
);
let results = (outs
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
@@ -2466,11 +2476,11 @@ def SPIRV_TosaTransposeOp : SPIRV_TosaOpWithResult<"Transpose", 60, [Pure,
let arguments = (ins
SPIRV_I32_1DTensorArmOfLength1To6Attr: $perms,
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $input1
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $input1
);
let results = (outs
- SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
@@ -2512,12 +2522,12 @@ def SPIRV_TosaGatherOp : SPIRV_TosaOpWithResult<"Gather", 61, [NoMemoryEffect,
}];
let arguments = (ins
- SPIRV_I8OrI16OrI32OrF16OrF32OrBF16_TensorArm3D: $values,
+ SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D: $values,
SPIRV_I32_TensorArm2D: $indices
);
let results = (outs
- SPIRV_I8OrI16OrI32OrF16OrF32OrBF16_TensorArm3D: $output
+ SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D: $output
);
let assemblyFormat = [{
@@ -2566,13 +2576,13 @@ def SPIRV_TosaScatterOp : SPIRV_TosaOpWithResult<"Scatter", 62, [NoMemoryEffect,
}];
let arguments = (ins
- SPIRV_I8OrI16OrI32OrF16OrF32OrBF16_TensorArm3D: $values_in,
+ SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D: $values_in,
SPIRV_I32_TensorArm2D: $indices,
- SPIRV_I8OrI16OrI32OrF16OrF32OrBF16_TensorArm3D: $input
+ SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D: $input
);
let results = (outs
- SPIRV_I8OrI16OrI32OrF16OrF32OrBF16_TensorArm3D: $values_out
+ SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D: $values_out
);
let assemblyFormat = [{
@@ -2687,13 +2697,15 @@ def SPIRV_TosaResizeOp : SPIRV_TosaOpWithResult<"Resize", 63, [Pure,
def SPIRV_TosaCastOp : SPIRV_TosaOpWithResult<"Cast", 64, [Pure,
AllShapesMatch<["input", "output"]>,
- TypeConstraintImplicationOn<"input", F16, "output", [F32, I16, I32, I8]>,
- TypeConstraintImplicationOn<"input", F32, "output", [F16, I16, I32, I8, BF16]>,
+ TypeConstraintImplicationOn<"input", F16, "output", [F32, I16, I32, I8, F8E4M3FN, F8E5M2]>,
+ TypeConstraintImplicationOn<"input", F32, "output", [F16, I16, I32, I8, BF16, F8E4M3FN, F8E5M2]>,
TypeConstraintImplicationOn<"input", I16, "output", [F16, F32, I32, I8, SPIRV_Bool, BF16]>,
TypeConstraintImplicationOn<"input", I32, "output", [F16, F32, I16, I8, SPIRV_Bool, BF16]>,
TypeConstraintImplicationOn<"input", I8, "output", [F16, F32, I16, I32, SPIRV_Bool, BF16]>,
TypeConstraintImplicationOn<"input", SPIRV_Bool, "output", [I16, I32, I8]>,
- TypeConstraintImplicationOn<"input", BF16, "output", [F32, I16, I32, I8]>]> {
+ TypeConstraintImplicationOn<"input", BF16, "output", [F32, I16, I32, I8, F8E4M3FN, F8E5M2]>,
+ TypeConstraintImplicationOn<"input", F8E4M3FN, "output", [F16, F32, BF16]>,
+ TypeConstraintImplicationOn<"input", F8E5M2, "output", [F16, F32, BF16]>]> {
let summary = "Cast operation.";
let description = [{
@@ -2737,6 +2749,18 @@ def SPIRV_TosaCastOp : SPIRV_TosaOpWithResult<"Cast", 64, [Pure,
| int16 | bf16 |
| int32 | bf16 |
| int8 | bf16 |
+ | bf16 | fp8e4m3 |
+ | fp8e4m3 | bf16 |
+ | bf16 | fp8e5m2 |
+ | fp8e5m2 | bf16 |
+ | float16 | fp8e4m3 |
+ | float32 | fp8e4m3 |
+ | fp8e4m3 | float16 |
+ | fp8e4m3 | float32 |
+ | float16 | fp8e5m2 |
+ | float32 | fp8e5m2 |
+ | fp8e5m2 | float16 |
+ | fp8e5m2 | float32 |
References:
* https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_cast
@@ -2750,11 +2774,11 @@ def SPIRV_TosaCastOp : SPIRV_TosaOpWithResult<"Cast", 64, [Pure,
}];
let arguments = (ins
- SPIRV_BoolOrI8OrI16OrI32OrBF16OrF16OrF32_TensorArm: $input
+ SPIRV_BoolOrI8OrI16OrI32OrBF16OrF16OrF32OrF8E4M3FNOrF8E5M2_TensorArm: $input
);
let results = (outs
- SPIRV_I8OrI16OrI32OrBoolOrF16OrF32OrBF16_TensorArm: $output
+ SPIRV_I8OrI16OrI32OrBoolOrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm: $output
);
let assemblyFormat = [{
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaTypes.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaTypes.td
index 6c918aec28845..316dc02da5b1e 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaTypes.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaTypes.td
@@ -22,16 +22,20 @@ def SPIRV_I8OrI16OrI32OrI64 : AnyIntOfWidths<[8, 16, 32, 64]>;
def SPIRV_I16OrI32 : AnyIntOfWidths<[16, 32]>;
def SPIRV_I32OrI64 : AnyIntOfWidths<[32, 64]>;
def SPIRV_F16OrF32OrBF16 : AnyTypeOf<[SPIRV_Float16, SPIRV_Float32, SPIRV_BFloat16KHR]>;
-def SPIRV_I8OrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_Int8, SPIRV_F16OrF32OrBF16]>;
+def SPIRV_F16OrF32OrBF16OrF8E4M3FNOrF8E5M2 : AnyTypeOf<[SPIRV_Float16, SPIRV_Float32, SPIRV_BFloat16KHR, SPIRV_Float8E4M3EXT, SPIRV_Float8E5M2EXT]>;
def SPIRV_I8OrI16OrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_I8OrI16, SPIRV_F16OrF32OrBF16]>;
+def SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2 : AnyTypeOf<[SPIRV_Int8, SPIRV_F16OrF32OrBF16OrF8E4M3FNOrF8E5M2]>;
+def SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2 : AnyTypeOf<[SPIRV_I8OrI16, SPIRV_F16OrF32OrBF16OrF8E4M3FNOrF8E5M2]>;
def SPIRV_I32OrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_Int32, SPIRV_F16OrF32OrBF16]>;
def SPIRV_I8OrI16OrI32OrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_I8OrI16OrI32, SPIRV_F16OrF32OrBF16]>;
+def SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2 : AnyTypeOf<[SPIRV_I8OrI16OrI32, SPIRV_F16OrF32OrBF16OrF8E4M3FNOrF8E5M2]>;
def SPIRV_I32OrI64OrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_I32OrI64, SPIRV_F16OrF32OrBF16]>;
def SPIRV_I32OrI64OrF16OrF32 : AnyTypeOf<[SPIRV_I32OrI64, SPIRV_Float16, SPIRV_Float32]>;
def SPIRV_I32OrI8OrI64OrI16OrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_I8OrI16OrI32OrI64, SPIRV_F16OrF32OrBF16]>;
-def SPIRV_BoolOrI8OrI16OrI32OrBF16OrF16OrF32 : AnyTypeOf<[SPIRV_Bool, SPIRV_I8OrI16OrI32, SPIRV_F16OrF32OrBF16]>;
def SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_Bool, SPIRV_I8OrI16OrI32, SPIRV_F16OrF32OrBF16]>;
-def SPIRV_I8OrI16OrI32OrBoolOrF16OrF32OrBF16 : AnyTypeOf<[SPIRV_I8OrI16OrI32, SPIRV_Bool, SPIRV_F16OrF32OrBF16]>;
+def SPIRV_BoolOrI8OrI16OrI32OrBF16OrF16OrF32OrF8E4M3FNOrF8E5M2 : AnyTypeOf<[SPIRV_Bool, SPIRV_I8OrI16OrI32, SPIRV_F16OrF32OrBF16OrF8E4M3FNOrF8E5M2]>;
+def SPIRV_BoolOrI8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2 : AnyTypeOf<[SPIRV_Bool, SPIRV_I8OrI16OrI32, SPIRV_F16OrF32OrBF16OrF8E4M3FNOrF8E5M2]>;
+def SPIRV_I8OrI16OrI32OrBoolOrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2 : AnyTypeOf<[SPIRV_I8OrI16OrI32, SPIRV_Bool, SPIRV_F16OrF32OrBF16OrF8E4M3FNOrF8E5M2]>;
def SPIRV_I8OrI32 : AnyTypeOf<[SPIRV_Int8, SPIRV_Int32]>;
def SPIRV_TensorArmAxisAttr : ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<5>]>;
@@ -57,23 +61,26 @@ def SPIRV_I32_TensorArm2D : TensorArmRankOf<[SPIRV_Int32], [2]>;
def SPIRV_F32_TensorArm3D: TensorArmRankOf<[SPIRV_Float32], [3]>;
def SPIRV_I32OrI64OrF16OrF32OrBF16_TensorArm1D : TensorArmRankOf<[SPIRV_I32OrI64OrF16OrF32OrBF16], [1]>;
def SPIRV_I8OrI16_TensorArm1D : TensorArmRankOf<[SPIRV_I8OrI16], [1]>;
-def SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm3D : TensorArmRankOf<[SPIRV_I8OrI16OrF16OrF32OrBF16], [3]>;
-def SPIRV_I8OrI16OrI32OrF16OrF32OrBF16_TensorArm3D : TensorArmRankOf<[SPIRV_I8OrI16OrI32OrF16OrF32OrBF16], [3]>;
+def SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D : TensorArmRankOf<[SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2], [3]>;
+def SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm3D : TensorArmRankOf<[SPIRV_I8OrI16OrI32OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2], [3]>;
def SPIRV_I32OrI64OrF16OrF32_TensorArm3D : TensorArmRankOf<[SPIRV_I32OrI64OrF16OrF32], [3]>;
def SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm4D : TensorArmRankOf<[SPIRV_I8OrI16OrF16OrF32OrBF16], [4]>;
-def SPIRV_I8OrF16OrF32OrBF16_TensorArm4D : TensorArmRankOf<[SPIRV_I8OrF16OrF32OrBF16], [4]>;
+def SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D : TensorArmRankOf<[SPIRV_I8OrI16OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2], [4]>;
+def SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2_TensorArm4D : TensorArmRankOf<[SPIRV_I8OrF16OrF32OrBF16OrF8E4M3FNOrF8E5M2], [4]>;
def SPIRV_I32OrI64OrF16OrF32OrBF16_TensorArm4D : TensorArmRankOf<[SPIRV_I32OrI64OrF16OrF32OrBF16], [4]>;
def SPIRV_I32OrI8OrI64OrI16OrF16OrF32OrBF16_TensorArm4D : TensorArmRankOf<[SPIRV_I32OrI8OrI64OrI16OrF16OrF32OrBF16], [4]>;
-def SPIRV_I8OrI16OrF16OrF32OrBF16_TensorArm5D : TensorArmRankOf<[SPIRV_I8OrI16OrF16OrF32OrBF16], [5]>;
-...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/193199
More information about the Mlir-commits
mailing list