Allow WorkgroupSize on variables for Kernels (#4627)

* Update tests

3 files changed, 36 insertions, 9 deletions

diff --git a/source/val/validate_annotation.cpp b/source/val/validate_annotation.cpp
index 16d44906d..3a77552e1 100644
--- a/source/val/validate_annotation.cpp
+++ b/source/val/validate_annotation.cpp
@@ -266,7 +266,8 @@ spv_result_t ValidateDecorationTarget(ValidationState_t& _, SpvDecoration dec,
                << "BuiltIns can only target variables, structure members or "
                   "constants";
       }
-      if (inst->GetOperandAs<SpvBuiltIn>(2) == SpvBuiltInWorkgroupSize) {
+      if (_.HasCapability(SpvCapabilityShader) &&
+          inst->GetOperandAs<SpvBuiltIn>(2) == SpvBuiltInWorkgroupSize) {
         if (!spvOpcodeIsConstant(target->opcode())) {
           return fail() << "must be a constant for WorkgroupSize";
         }
diff --git a/test/val/val_annotation_test.cpp b/test/val/val_annotation_test.cpp
index cd560057f..889c76cae 100644
--- a/test/val/val_annotation_test.cpp
+++ b/test/val/val_annotation_test.cpp
@@ -32,6 +32,40 @@ using ::testing::Eq;
 using ::testing::HasSubstr;
 using ::testing::Values;
 
+using DecorationTest = spvtest::ValidateBase<bool>;
+
+TEST_F(DecorationTest, WorkgroupSizeShader) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpDecorate %ones BuiltIn WorkgroupSize
+%int = OpTypeInt 32 0
+%int3 = OpTypeVector %int 3
+%int_1 = OpConstant %int 1
+%ones = OpConstantComposite %int3 %int_1 %int_1 %int_1
+)";
+
+  CompileSuccessfully(text);
+  EXPECT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
+TEST_F(DecorationTest, WorkgroupSizeKernel) {
+  const std::string text = R"(
+OpCapability Kernel
+OpCapability Linkage
+OpMemoryModel Logical OpenCL
+OpDecorate %var BuiltIn WorkgroupSize
+%int = OpTypeInt 32 0
+%int3 = OpTypeVector %int 3
+%ptr = OpTypePointer Input %int3
+%var = OpVariable %ptr Input
+)";
+
+  CompileSuccessfully(text);
+  EXPECT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
 using MemberOnlyDecorations = spvtest::ValidateBase<std::string>;
 
 TEST_P(MemberOnlyDecorations, MemberDecoration) {
diff --git a/test/val/val_capability_test.cpp b/test/val/val_capability_test.cpp
index ffd9e0ae2..c432c3cfa 100644
--- a/test/val/val_capability_test.cpp
+++ b/test/val/val_capability_test.cpp
@@ -1702,14 +1702,6 @@ std::make_pair(std::string(kOpenCLMemoryModel) +
           AllCapabilities()),
 std::make_pair(std::string(kOpenCLMemoryModel) +
           "OpEntryPoint Kernel %func \"compute\" \n" +
-          "OpDecorate %ones BuiltIn WorkgroupSize\n"
-          "%intt = OpTypeInt 32 0\n"
-          "%int3 = OpTypeVector %intt 3\n"
-          "%int_1 = OpConstant %intt 1\n"
-          "%ones = OpConstantComposite %int3 %int_1 %int_1 %int_1\n" + std::string(kVoidFVoid),
-          AllCapabilities()),
-std::make_pair(std::string(kOpenCLMemoryModel) +
-          "OpEntryPoint Kernel %func \"compute\" \n" +
           "OpDecorate %var BuiltIn WorkgroupId\n"
           "%intt = OpTypeInt 32 0\n"
           "%ptr = OpTypePointer Input %intt\n"