Skip to content

Commit 21c81e0

Browse files
authored
val: Detect additional unsupported arithmetic ops in invalid_type_pass (KhronosGroup#6707)
@0cc4m reported OpMatrixTimesScalar with bf16 coopmat slipping through the cracks. I asked codex to fill any remaining gaps in the arithmetic instructions.
1 parent 5418da0 commit 21c81e0

2 files changed

Lines changed: 191 additions & 26 deletions

File tree

source/val/validate_invalid_type.cpp

Lines changed: 81 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,57 @@
2424
namespace spvtools {
2525
namespace val {
2626

27+
namespace {
28+
29+
bool IsBfloat16ScalarOrCompositeType(ValidationState_t& _, uint32_t type_id) {
30+
const Instruction* inst = _.FindDef(type_id);
31+
if (!inst) return false;
32+
33+
if (_.IsBfloat16Type(type_id)) return true;
34+
35+
switch (inst->opcode()) {
36+
case spv::Op::OpTypeMatrix:
37+
case spv::Op::OpTypeCooperativeMatrixNV:
38+
return _.IsBfloat16ScalarType(_.GetComponentType(type_id));
39+
default:
40+
return false;
41+
}
42+
}
43+
44+
bool IsFP8ScalarOrCompositeType(ValidationState_t& _, uint32_t type_id) {
45+
const Instruction* inst = _.FindDef(type_id);
46+
if (!inst) return false;
47+
48+
if (_.IsFP8Type(type_id)) return true;
49+
50+
switch (inst->opcode()) {
51+
case spv::Op::OpTypeMatrix:
52+
case spv::Op::OpTypeCooperativeMatrixNV:
53+
return _.IsFP8ScalarType(_.GetComponentType(type_id));
54+
default:
55+
return false;
56+
}
57+
}
58+
59+
spv_result_t CheckInvalidScalarOrCompositeType(ValidationState_t& _,
60+
const Instruction* inst,
61+
uint32_t type_id) {
62+
const spv::Op opcode = inst->opcode();
63+
if (IsBfloat16ScalarOrCompositeType(_, type_id)) {
64+
return _.diag(SPV_ERROR_INVALID_DATA, inst)
65+
<< spvOpcodeString(opcode) << " doesn't support BFloat16 type.";
66+
}
67+
if (IsFP8ScalarOrCompositeType(_, type_id)) {
68+
return _.diag(SPV_ERROR_INVALID_DATA, inst)
69+
<< spvOpcodeString(opcode)
70+
<< " doesn't support FP8 E4M3/E5M2 types.";
71+
}
72+
73+
return SPV_SUCCESS;
74+
}
75+
76+
} // namespace
77+
2778
// Validates correctness of certain special type instructions.
2879
spv_result_t InvalidTypePass(ValidationState_t& _, const Instruction* inst) {
2980
const spv::Op opcode = inst->opcode();
@@ -39,6 +90,13 @@ spv_result_t InvalidTypePass(ValidationState_t& _, const Instruction* inst) {
3990
case spv::Op::OpFRem:
4091
case spv::Op::OpFMod:
4192
case spv::Op::OpFNegate:
93+
case spv::Op::OpFmaKHR:
94+
case spv::Op::OpVectorTimesScalar:
95+
case spv::Op::OpMatrixTimesScalar:
96+
case spv::Op::OpVectorTimesMatrix:
97+
case spv::Op::OpMatrixTimesVector:
98+
case spv::Op::OpMatrixTimesMatrix:
99+
case spv::Op::OpOuterProduct:
42100
// Derivative Instructions
43101
case spv::Op::OpDPdx:
44102
case spv::Op::OpDPdy:
@@ -69,10 +127,30 @@ spv_result_t InvalidTypePass(ValidationState_t& _, const Instruction* inst) {
69127
case spv::Op::OpGroupNonUniformFMul:
70128
case spv::Op::OpGroupNonUniformFMin: {
71129
const uint32_t result_type = inst->type_id();
72-
if (_.IsBfloat16Type(result_type)) {
73-
return _.diag(SPV_ERROR_INVALID_DATA, inst)
74-
<< spvOpcodeString(opcode) << " doesn't support BFloat16 type.";
130+
if (spv_result_t result =
131+
CheckInvalidScalarOrCompositeType(_, inst, result_type)) {
132+
return result;
75133
}
134+
135+
break;
136+
}
137+
138+
case spv::Op::OpCooperativeMatrixMulAddNV: {
139+
if (spv_result_t result =
140+
CheckInvalidScalarOrCompositeType(_, inst, inst->type_id())) {
141+
return result;
142+
}
143+
for (uint32_t operand_index = 2; operand_index <= 4; ++operand_index) {
144+
if (spv_result_t result = CheckInvalidScalarOrCompositeType(
145+
_, inst, _.GetOperandTypeId(inst, operand_index))) {
146+
return result;
147+
}
148+
}
149+
break;
150+
}
151+
152+
case spv::Op::OpDot: {
153+
const uint32_t result_type = inst->type_id();
76154
if (_.IsFP8Type(result_type)) {
77155
return _.diag(SPV_ERROR_INVALID_DATA, inst)
78156
<< spvOpcodeString(opcode)
@@ -143,29 +221,6 @@ spv_result_t InvalidTypePass(ValidationState_t& _, const Instruction* inst) {
143221

144222
break;
145223
}
146-
147-
case spv::Op::OpMatrixTimesMatrix: {
148-
const uint32_t result_type = inst->type_id();
149-
uint32_t res_num_rows = 0;
150-
uint32_t res_num_cols = 0;
151-
uint32_t res_col_type = 0;
152-
uint32_t res_component_type = 0;
153-
if (_.GetMatrixTypeInfo(result_type, &res_num_rows, &res_num_cols,
154-
&res_col_type, &res_component_type)) {
155-
if (_.IsBfloat16Type(res_component_type)) {
156-
return _.diag(SPV_ERROR_INVALID_DATA, inst)
157-
<< spvOpcodeString(opcode)
158-
<< " doesn't support BFloat16 type.";
159-
}
160-
if (_.IsFP8Type(res_component_type)) {
161-
return _.diag(SPV_ERROR_INVALID_DATA, inst)
162-
<< spvOpcodeString(opcode)
163-
<< " doesn't support FP8 E4M3/E5M2 types.";
164-
}
165-
}
166-
break;
167-
}
168-
169224
default:
170225
break;
171226
}

test/val/val_invalid_type_test.cpp

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,10 @@ std::string GenerateBFloatCode(const std::string& main_body) {
3939
OpCapability Shader
4040
OpCapability BFloat16TypeKHR
4141
OpCapability AtomicFloat16AddEXT
42+
OpCapability FMAKHR
4243
OpCapability GroupNonUniformShuffle
4344
OpExtension "SPV_EXT_shader_atomic_float16_add"
45+
OpExtension "SPV_KHR_fma"
4446
OpExtension "SPV_KHR_bfloat16"
4547
%1 = OpExtInstImport "GLSL.std.450"
4648
OpMemoryModel Logical GLSL450
@@ -73,6 +75,43 @@ OpFunctionEnd)";
7375
return prefix + main_body + suffix;
7476
}
7577

78+
std::string GenerateBFloatCoopMatCode(const std::string& main_body) {
79+
const std::string prefix =
80+
R"(
81+
OpCapability Shader
82+
OpCapability VulkanMemoryModel
83+
OpCapability BFloat16TypeKHR
84+
OpCapability BFloat16CooperativeMatrixKHR
85+
OpCapability CooperativeMatrixKHR
86+
OpExtension "SPV_KHR_bfloat16"
87+
OpExtension "SPV_KHR_cooperative_matrix"
88+
OpExtension "SPV_KHR_vulkan_memory_model"
89+
OpMemoryModel Logical Vulkan
90+
OpEntryPoint GLCompute %main "main"
91+
OpExecutionMode %main LocalSize 32 1 1
92+
OpSource GLSL 450
93+
OpName %main "main"
94+
%void = OpTypeVoid
95+
%bfloat16 = OpTypeFloat 16 BFloat16KHR
96+
%func = OpTypeFunction %void
97+
%u32 = OpTypeInt 32 0
98+
%u32_8 = OpConstant %u32 8
99+
%subgroup = OpConstant %u32 3
100+
%useA = OpConstant %u32 0
101+
%bf16_1 = OpConstant %bfloat16 1
102+
%bf16matA = OpTypeCooperativeMatrixKHR %bfloat16 %subgroup %u32_8 %u32_8 %useA
103+
%bf16mat_A_1 = OpConstantComposite %bf16matA %bf16_1
104+
%main = OpFunction %void None %func
105+
%main_entry = OpLabel)";
106+
107+
const std::string suffix =
108+
R"(
109+
OpReturn
110+
OpFunctionEnd)";
111+
112+
return prefix + main_body + suffix;
113+
}
114+
76115
TEST_F(ValidateInvalidType, Bfloat16InvalidArithmeticInstruction) {
77116
const std::string body = R"(
78117
%v1 = OpVariable %_ptr_Function_bfloat16 Function
@@ -89,6 +128,45 @@ TEST_F(ValidateInvalidType, Bfloat16InvalidArithmeticInstruction) {
89128
HasSubstr("FMul doesn't support BFloat16 type."));
90129
}
91130

131+
TEST_F(ValidateInvalidType, Bfloat16InvalidFmaInstruction) {
132+
const std::string body = R"(
133+
%15 = OpFmaKHR %bfloat16 %bf16_1 %bf16_1 %bf16_1
134+
)";
135+
136+
CompileSuccessfully(GenerateBFloatCode(body).c_str(), SPV_ENV_VULKAN_1_3);
137+
ASSERT_EQ(SPV_ERROR_INVALID_DATA,
138+
ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
139+
EXPECT_THAT(getDiagnosticString(),
140+
HasSubstr("FmaKHR doesn't support BFloat16 type."));
141+
}
142+
143+
TEST_F(ValidateInvalidType, Bfloat16InvalidVectorTimesScalarInstruction) {
144+
const std::string body = R"(
145+
%v1 = OpVariable %_ptr_Function_v2bfloat16 Function
146+
%12 = OpLoad %v2bfloat16 %v1
147+
%15 = OpVectorTimesScalar %v2bfloat16 %12 %bf16_1
148+
)";
149+
150+
CompileSuccessfully(GenerateBFloatCode(body).c_str(), SPV_ENV_VULKAN_1_3);
151+
ASSERT_EQ(SPV_ERROR_INVALID_DATA,
152+
ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
153+
EXPECT_THAT(getDiagnosticString(),
154+
HasSubstr("VectorTimesScalar doesn't support BFloat16 type."));
155+
}
156+
157+
TEST_F(ValidateInvalidType, Bfloat16InvalidMatrixTimesScalarInstruction) {
158+
const std::string body = R"(
159+
%15 = OpMatrixTimesScalar %bf16matA %bf16mat_A_1 %bf16_1
160+
)";
161+
162+
CompileSuccessfully(GenerateBFloatCoopMatCode(body).c_str(),
163+
SPV_ENV_UNIVERSAL_1_6);
164+
ASSERT_EQ(SPV_ERROR_INVALID_DATA,
165+
ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
166+
EXPECT_THAT(getDiagnosticString(),
167+
HasSubstr("MatrixTimesScalar doesn't support BFloat16 type."));
168+
}
169+
92170
TEST_F(ValidateInvalidType, Bfloat16InvalidRelationalInstruction) {
93171
const std::string body = R"(
94172
%v1 = OpVariable %_ptr_Function_bfloat16 Function
@@ -209,6 +287,38 @@ TEST_F(ValidateInvalidType, FP8E5M2InvalidArithmeticInstruction) {
209287
HasSubstr("FMul doesn't support FP8 E4M3/E5M2 types."));
210288
}
211289

290+
TEST_F(ValidateInvalidType, FP8E4M3InvalidDotInstruction) {
291+
const std::string body = R"(
292+
%v1 = OpVariable %_ptr_Function_v2fp8e4m3 Function
293+
%v2 = OpVariable %_ptr_Function_v2fp8e4m3 Function
294+
%12 = OpLoad %v2fp8e4m3 %v1
295+
%14 = OpLoad %v2fp8e4m3 %v2
296+
%15 = OpDot %fp8e4m3 %12 %14
297+
)";
298+
299+
CompileSuccessfully(GenerateFP8Code(body).c_str(), SPV_ENV_VULKAN_1_3);
300+
ASSERT_EQ(SPV_ERROR_INVALID_DATA,
301+
ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
302+
EXPECT_THAT(getDiagnosticString(),
303+
HasSubstr("Dot doesn't support FP8 E4M3/E5M2 types."));
304+
}
305+
306+
TEST_F(ValidateInvalidType, FP8E5M2InvalidDotInstruction) {
307+
const std::string body = R"(
308+
%v1 = OpVariable %_ptr_Function_v2fp8e5m2 Function
309+
%v2 = OpVariable %_ptr_Function_v2fp8e5m2 Function
310+
%12 = OpLoad %v2fp8e5m2 %v1
311+
%14 = OpLoad %v2fp8e5m2 %v2
312+
%15 = OpDot %fp8e5m2 %12 %14
313+
)";
314+
315+
CompileSuccessfully(GenerateFP8Code(body).c_str(), SPV_ENV_VULKAN_1_3);
316+
ASSERT_EQ(SPV_ERROR_INVALID_DATA,
317+
ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
318+
EXPECT_THAT(getDiagnosticString(),
319+
HasSubstr("Dot doesn't support FP8 E4M3/E5M2 types."));
320+
}
321+
212322
TEST_F(ValidateInvalidType, FP8E4M3InvalidRelationalInstruction) {
213323
const std::string body = R"(
214324
%v1 = OpVariable %_ptr_Function_fp8e4m3 Function

0 commit comments

Comments
 (0)