Add validation of storage classes for WebGPU (#2446)

Fixes #2445

source/val/validate_instruction.cpp

@@ -55,6 +55,18 @@ std::string ToString(const CapabilitySet& capabilities,
   return ss.str();
 }
 
+bool IsValidWebGPUStorageClass(SpvStorageClass storage_class) {
+  return storage_class == SpvStorageClassUniformConstant ||
+         storage_class == SpvStorageClassUniform ||
+         storage_class == SpvStorageClassStorageBuffer ||
+         storage_class == SpvStorageClassInput ||
+         storage_class == SpvStorageClassOutput ||
+         storage_class == SpvStorageClassImage ||
+         storage_class == SpvStorageClassWorkgroup ||
+         storage_class == SpvStorageClassPrivate ||
+         storage_class == SpvStorageClassFunction;
+}
+
 // Returns capabilities that enable an opcode.  An empty result is interpreted
 // as no prohibition of use of the opcode.  If the result is non-empty, then
 // the opcode may only be used if at least one of the capabilities is specified
@@ -483,6 +495,15 @@ spv_result_t InstructionPass(ValidationState_t& _, const Instruction* inst) {
     if (auto error = LimitCheckNumVars(_, inst->id(), storage_class)) {
       return error;
     }
+
+    if (spvIsWebGPUEnv(_.context()->target_env) &&
+        !IsValidWebGPUStorageClass(storage_class)) {
+      return _.diag(SPV_ERROR_INVALID_BINARY, inst)
+             << "For WebGPU, OpVariable storage class must be one of "
+                "UniformConstant, Uniform, StorageBuffer, Input, Output, "
+                "Image, Workgroup, Private, Function for WebGPU";
+    }
+
     if (storage_class == SpvStorageClassGeneric)
       return _.diag(SPV_ERROR_INVALID_BINARY, inst)
              << "OpVariable storage class cannot be Generic";
@@ -506,6 +527,15 @@ spv_result_t InstructionPass(ValidationState_t& _, const Instruction* inst) {
                   "outside of a function";
       }
     }
+  } else if (opcode == SpvOpTypePointer) {
+    const auto storage_class = inst->GetOperandAs<SpvStorageClass>(1);
+    if (spvIsWebGPUEnv(_.context()->target_env) &&
+        !IsValidWebGPUStorageClass(storage_class)) {
+      return _.diag(SPV_ERROR_INVALID_BINARY, inst)
+             << "For WebGPU, OpTypePointer storage class must be one of "
+                "UniformConstant, Uniform, StorageBuffer, Input, Output, "
+                "Image, Workgroup, Private, Function";
+    }
   }
 
   // SPIR-V Spec 2.16.3: Validation Rules for Kernel Capabilities: The

test/val/val_storage_test.cpp

@@ -26,7 +26,10 @@ namespace val {
 namespace {
 
 using ::testing::HasSubstr;
+using ::testing::Values;
 using ValidateStorage = spvtest::ValidateBase<std::string>;
+using ValidateStorageClass =
+    spvtest::ValidateBase<std::tuple<std::string, bool, bool, std::string>>;
 
 TEST_F(ValidateStorage, FunctionStorageInsideFunction) {
   char str[] = R"(
@@ -245,6 +248,72 @@ TEST_F(ValidateStorage, RelaxedLogicalPointerFunctionParamBad) {
               HasSubstr("OpFunctionCall Argument <id> '"));
 }
 
+std::string GetVarDeclStr(const std::string& storage_class) {
+  if (storage_class != "Output" && storage_class != "Private" &&
+      storage_class != "Function") {
+    return "%var    = OpVariable %ptrt " + storage_class + "\n";
+  } else {
+    return "%var    = OpVariable %ptrt " + storage_class + " %null\n";
+  }
+}
+
+TEST_P(ValidateStorageClass, WebGPU) {
+  std::string storage_class = std::get<0>(GetParam());
+  bool is_local = std::get<1>(GetParam());
+  bool is_valid = std::get<2>(GetParam());
+  std::string error = std::get<3>(GetParam());
+
+  std::string str = R"(
+          OpCapability Shader
+          OpCapability VulkanMemoryModelKHR
+          OpExtension "SPV_KHR_vulkan_memory_model"
+          OpMemoryModel Logical VulkanKHR
+          OpEntryPoint Fragment %func "func"
+          OpExecutionMode %func OriginUpperLeft
+%intt   = OpTypeInt 32 1
+%voidt  = OpTypeVoid
+%vfunct = OpTypeFunction %voidt
+%null   = OpConstantNull %intt
+)";
+  str += "%ptrt   = OpTypePointer " + storage_class + " %intt\n";
+  if (!is_local) str += GetVarDeclStr(storage_class);
+  str += R"(
+%func   = OpFunction %voidt None %vfunct
+%funcl  = OpLabel
+)";
+  if (is_local) str += GetVarDeclStr(storage_class);
+  str += R"(
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(str, SPV_ENV_WEBGPU_0);
+  if (is_valid) {
+    ASSERT_EQ(SPV_SUCCESS, ValidateInstructions(SPV_ENV_WEBGPU_0));
+  } else {
+    ASSERT_EQ(SPV_ERROR_INVALID_BINARY, ValidateInstructions(SPV_ENV_WEBGPU_0));
+    EXPECT_THAT(getDiagnosticString(), HasSubstr(error));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    StorageClass, ValidateStorageClass,
+    Values(std::make_tuple("UniformConstant", false, true, ""),
+           std::make_tuple("Uniform", false, true, ""),
+           std::make_tuple("StorageBuffer", false, true, ""),
+           std::make_tuple("Input", false, true, ""),
+           std::make_tuple("Output", false, true, ""),
+           std::make_tuple("Image", false, true, ""),
+           std::make_tuple("Workgroup", false, true, ""),
+           std::make_tuple("Private", false, true, ""),
+           std::make_tuple("Function", true, true, ""),
+           std::make_tuple(
+               "CrossWorkgroup", false, false,
+               "For WebGPU, OpTypePointer storage class must be one of"),
+           std::make_tuple(
+               "PushConstant", false, false,
+               "For WebGPU, OpTypePointer storage class must be one of")));
+
 }  // namespace
 }  // namespace val
 }  // namespace spvtools