[val] Do not reject OpExtInst for BFloat16/FP8 result types (KhronosGroup#6734)

Needed for llvm/llvm-project#202859

Author: aobolensk

source/val/validate_extensions.cpp

@@ -26,6 +26,7 @@
 #include "source/latest_version_glsl_std_450_header.h"
 #include "source/latest_version_opencl_std_header.h"
 #include "source/opcode.h"
+#include "source/operand.h"
 #include "source/spirv_constant.h"
 #include "source/table2.h"
 #include "source/val/instruction.h"
@@ -163,6 +164,38 @@ std::string GetExtInstName(const ValidationState_t& _,
   return ss.str();
 }
 
+// Rejects an instruction whose result or any operand uses a BFloat16 or FP8
+// (E4M3/E5M2) type, i.e. an OpTypeFloat that is not IEEE 754 encoded.
+spv_result_t ValidateExtInstFloatEncoding(ValidationState_t& _,
+                                          const Instruction* inst) {
+  auto check = [&](uint32_t type_id) -> spv_result_t {
+    if (_.IsBfloat16Type(type_id)) {
+      return _.diag(SPV_ERROR_INVALID_DATA, inst)
+             << GetExtInstName(_, inst) << ": doesn't support BFloat16 type.";
+    }
+    if (_.IsFP8Type(type_id)) {
+      return _.diag(SPV_ERROR_INVALID_DATA, inst)
+             << GetExtInstName(_, inst)
+             << ": doesn't support FP8 E4M3/E5M2 types.";
+    }
+    return SPV_SUCCESS;
+  };
+
+  if (spv_result_t result = check(inst->type_id())) return result;
+
+  const uint32_t num_operands = static_cast<uint32_t>(inst->operands().size());
+  for (uint32_t operand_index = 4; operand_index < num_operands;
+       ++operand_index) {
+    // Some ext instructions take literal operands (e.g. OpenCL.std vloadn's
+    // component count); only <id> operands have a meaningful result type.
+    if (!spvIsIdType(inst->operand(operand_index).type)) continue;
+    if (spv_result_t result = check(_.GetOperandTypeId(inst, operand_index)))
+      return result;
+  }
+
+  return SPV_SUCCESS;
+}
+
 // Returns the declared NSDI version from the OpExtInstImport referenced by
 // |inst|.  Returns 0 if not a NonSemantic.Shader.DebugInfo import.
 uint32_t GetNSDIVersion(const ValidationState_t& _, const Instruction* inst) {
@@ -4420,6 +4453,14 @@ spv_result_t ValidateExtInst(ValidationState_t& _, const Instruction* inst) {
   const spv_ext_inst_type_t ext_inst_type =
       spv_ext_inst_type_t(inst->ext_inst_type());
 
+  // GLSL.std.450 and OpenCL.std define a floating-point type as an OpTypeFloat
+  // using the IEEE 754 encoding, so they don't support BFloat16 or FP8 types.
+  if (ext_inst_type == SPV_EXT_INST_TYPE_GLSL_STD_450 ||
+      ext_inst_type == SPV_EXT_INST_TYPE_OPENCL_STD) {
+    if (spv_result_t result = ValidateExtInstFloatEncoding(_, inst))
+      return result;
+  }
+
   if (ext_inst_type == SPV_EXT_INST_TYPE_GLSL_STD_450) {
     return ValidateExtInstGlslStd450(_, inst);
   } else if (ext_inst_type == SPV_EXT_INST_TYPE_OPENCL_STD) {

source/val/validate_invalid_type.cpp

@@ -90,8 +90,6 @@ spv_result_t InvalidTypePass(ValidationState_t& _, const Instruction* inst) {
   const spv::Op opcode = inst->opcode();
 
   switch (opcode) {
-    // OpExtInst
-    case spv::Op::OpExtInst:
     // Arithmetic Instructions
     case spv::Op::OpFAdd:
     case spv::Op::OpFSub:

test/val/val_invalid_type_test.cpp

@@ -208,6 +208,44 @@ TEST_F(ValidateInvalidType, Bfloat16InvalidGroupNonUniformShuffle) {
       HasSubstr("GroupNonUniformShuffle doesn't support BFloat16 type."));
 }
 
+TEST_F(ValidateInvalidType, Bfloat16ExtInstruction) {
+  const std::string body = R"(
+%15 = OpExtInst %bfloat16 %1 FClamp %bf16_1 %bf16_1 %bf16_1
+)";
+
+  CompileSuccessfully(GenerateBFloatCode(body).c_str(), SPV_ENV_VULKAN_1_3);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("FClamp: doesn't support BFloat16 type."));
+}
+
+TEST_F(ValidateInvalidType, Bfloat16ExtInstructionRequiresExtension) {
+  const std::string spirv = R"(
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %main "main"
+OpExecutionMode %main LocalSize 1 1 1
+%void = OpTypeVoid
+%func = OpTypeFunction %void
+%bfloat16 = OpTypeFloat 16 BFloat16KHR
+%bf16_1 = OpConstant %bfloat16 1
+%main = OpFunction %void None %func
+%entry = OpLabel
+%15 = OpExtInst %bfloat16 %1 FClamp %bf16_1 %bf16_1 %bf16_1
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(spirv.c_str(), SPV_ENV_VULKAN_1_3);
+  EXPECT_EQ(SPV_ERROR_INVALID_CAPABILITY,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("Operand 3 of TypeFloat requires one of these "
+                        "capabilities: BFloat16TypeKHR"));
+}
+
 std::string GenerateFP8Code(const std::string& main_body) {
   const std::string prefix =
       R"(
@@ -401,6 +439,56 @@ TEST_F(ValidateInvalidType, FP8E5M2InvalidGroupNonUniformShuffle) {
       HasSubstr("GroupNonUniformShuffle doesn't support FP8 E4M3/E5M2 types."));
 }
 
+TEST_F(ValidateInvalidType, FP8E4M3ExtInstruction) {
+  const std::string body = R"(
+%15 = OpExtInst %fp8e4m3 %1 FClamp %fp8e4m3_1 %fp8e4m3_1 %fp8e4m3_1
+)";
+
+  CompileSuccessfully(GenerateFP8Code(body).c_str(), SPV_ENV_VULKAN_1_3);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("FClamp: doesn't support FP8 E4M3/E5M2 types."));
+}
+
+TEST_F(ValidateInvalidType, FP8E5M2ExtInstruction) {
+  const std::string body = R"(
+%15 = OpExtInst %fp8e5m2 %1 FClamp %fp8e5m2_1 %fp8e5m2_1 %fp8e5m2_1
+)";
+
+  CompileSuccessfully(GenerateFP8Code(body).c_str(), SPV_ENV_VULKAN_1_3);
+  EXPECT_EQ(SPV_ERROR_INVALID_DATA,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("FClamp: doesn't support FP8 E4M3/E5M2 types."));
+}
+
+TEST_F(ValidateInvalidType, FP8ExtInstructionRequiresExtension) {
+  const std::string spirv = R"(
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %main "main"
+OpExecutionMode %main LocalSize 1 1 1
+%void = OpTypeVoid
+%func = OpTypeFunction %void
+%fp8e4m3 = OpTypeFloat 8 Float8E4M3EXT
+%fp8e4m3_1 = OpConstant %fp8e4m3 1
+%main = OpFunction %void None %func
+%entry = OpLabel
+%15 = OpExtInst %fp8e4m3 %1 FClamp %fp8e4m3_1 %fp8e4m3_1 %fp8e4m3_1
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(spirv.c_str(), SPV_ENV_VULKAN_1_3);
+  EXPECT_EQ(SPV_ERROR_INVALID_CAPABILITY,
+            ValidateInstructions(SPV_ENV_UNIVERSAL_1_6));
+  EXPECT_THAT(getDiagnosticString(),
+              HasSubstr("Operand 3 of TypeFloat requires one of these "
+                        "capabilities: Float8EXT"));
+}
+
 struct OCPMicroscalingCase {
   const char* capability;
   const char* type_decl;