commit dad0f9a717193da223ebf3c443b0fe347096caa9
author Natalie Chouinard <chouinard@google.com> Tue Jul 16 16:31:41 2024 +0000
committer Dawn LUCI CQ <dawn-scoped@luci-project-accounts.iam.gserviceaccount.com> Tue Jul 16 16:31:41 2024 +0000
tree f6ea5250141f5253d397a0bcd66bb142e5b800a8
parent 188dd4295d02f99e2bcfc44047788903b1a4121e

[Tint] Validate subgroup builtins as input only

Validate that subgroup_size and subgroup_invocation_id are input only.

Bug: 352540632
Change-Id: Iaeab5781f9022cf085177664711ab0a2a93fd5ac
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/198574
Reviewed-by: dan sinclair <dsinclair@chromium.org>
Commit-Queue: Natalie Chouinard <chouinard@google.com>

src/tint/lang/wgsl/resolver/subgroups_extension_test.cc

diff --git a/src/tint/lang/wgsl/resolver/subgroups_extension_test.cc b/src/tint/lang/wgsl/resolver/subgroups_extension_test.cc
index 5e6a8b1..c93edfc 100644
--- a/src/tint/lang/wgsl/resolver/subgroups_extension_test.cc
+++ b/src/tint/lang/wgsl/resolver/subgroups_extension_test.cc
@@ -30,6 +30,7 @@
 
 #include "gmock/gmock.h"
 
+using namespace tint::core::fluent_types;     // NOLINT
 using namespace tint::core::number_suffixes;  // NOLINT
 
 namespace tint::resolver {
@@ -165,6 +166,47 @@
               "error: '@builtin(subgroup_size)' is only valid as a compute shader input");
 }
 
+TEST_F(ResolverSubgroupsExtensionTest, SubgroupSizeComputeShaderOutput) {
+    Enable(wgsl::Extension::kSubgroups);
+
+    Func("main", tint::Empty, ty.u32(),
+         Vector{
+             Return(Call<u32>()),
+         },
+         Vector{
+             Stage(ast::PipelineStage::kCompute),
+             WorkgroupSize(1_i),
+         },
+         Vector{Builtin(Source{{1, 2}}, core::BuiltinValue::kSubgroupSize)});
+
+    EXPECT_FALSE(r()->Resolve());
+    EXPECT_EQ(r()->error(),
+              "1:2 error: '@builtin(subgroup_size)' is only valid as a compute shader input");
+}
+
+TEST_F(ResolverSubgroupsExtensionTest, SubgroupSizeComputeShaderStructOutput) {
+    Enable(wgsl::Extension::kSubgroups);
+
+    auto* s = Structure(
+        "Output", Vector{
+                      Member("a", ty.u32(), Vector{Builtin(core::BuiltinValue::kSubgroupSize)}),
+                  });
+
+    Func("main", tint::Empty, ty.Of(s),
+         Vector{
+             Return(Call(ty.Of(s), Call<u32>())),
+         },
+         Vector{
+             Stage(ast::PipelineStage::kCompute),
+             WorkgroupSize(1_i),
+         });
+
+    EXPECT_FALSE(r()->Resolve());
+    EXPECT_EQ(r()->error(),
+              R"(error: '@builtin(subgroup_size)' is only valid as a compute shader input
+note: while analyzing entry point 'main')");
+}
+
 // Using builtin(subgroup_invocation_id) for anything other than a compute shader input should fail.
 TEST_F(ResolverSubgroupsExtensionTest, SubgroupInvocationIdFragmentShader) {
     Enable(wgsl::Extension::kSubgroups);
@@ -177,5 +219,24 @@
               "error: '@builtin(subgroup_invocation_id)' is only valid as a compute shader input");
 }
 
+TEST_F(ResolverSubgroupsExtensionTest, SubgroupInvocationIdComputeShaderOutput) {
+    Enable(wgsl::Extension::kSubgroups);
+
+    Func("main", tint::Empty, ty.u32(),
+         Vector{
+             Return(Call<u32>()),
+         },
+         Vector{
+             Stage(ast::PipelineStage::kCompute),
+             WorkgroupSize(1_i),
+         },
+         Vector{Builtin(Source{{1, 2}}, core::BuiltinValue::kSubgroupInvocationId)});
+
+    EXPECT_FALSE(r()->Resolve());
+    EXPECT_EQ(
+        r()->error(),
+        "1:2 error: '@builtin(subgroup_invocation_id)' is only valid as a compute shader input");
+}
+
 }  // namespace
 }  // namespace tint::resolver

src/tint/lang/wgsl/resolver/validator.cc

diff --git a/src/tint/lang/wgsl/resolver/validator.cc b/src/tint/lang/wgsl/resolver/validator.cc
index 964ebea..32d158d 100644
--- a/src/tint/lang/wgsl/resolver/validator.cc
+++ b/src/tint/lang/wgsl/resolver/validator.cc
@@ -1109,7 +1109,8 @@
                 err_builtin_type("u32");
                 return false;
             }
-            if (stage != ast::PipelineStage::kNone && stage != ast::PipelineStage::kCompute) {
+            if (stage != ast::PipelineStage::kNone &&
+                !(stage == ast::PipelineStage::kCompute && is_input)) {
                 AddError(attr->source)
                     << style::Attribute("@builtin") << style::Code("(", style::Enum(builtin), ")")
                     << " is only valid as a compute shader input";