From 1bb637195b530750e119eb9ad223236ba2c43be2 Mon Sep 17 00:00:00 2001
From: James Price <jrprice@google.com>
Date: Mon, 13 Sep 2021 17:11:58 +0000
Subject: [PATCH] resolver: Enable the num_workgroups builtin

The subsequent commits will implement this for all remaining backends.

Bug: tint:752
Change-Id: Id26f920fea93b5678466705612dcbfc229dc878a
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/63921
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: James Price <jrprice@google.com>
---
 src/resolver/builtins_validation_test.cc | 31 +++++++++++++++++++++++-
 src/resolver/resolver.cc                 |  5 +---
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/resolver/builtins_validation_test.cc b/src/resolver/builtins_validation_test.cc
index 38d04e04c0..e29b6f224e 100644
--- a/src/resolver/builtins_validation_test.cc
+++ b/src/resolver/builtins_validation_test.cc
@@ -128,6 +128,16 @@ static constexpr Params cases[] = {
                          ast::PipelineStage::kCompute,
                          true),
 
+    ParamsFor<vec3<u32>>(ast::Builtin::kNumWorkgroups,
+                         ast::PipelineStage::kVertex,
+                         false),
+    ParamsFor<vec3<u32>>(ast::Builtin::kNumWorkgroups,
+                         ast::PipelineStage::kFragment,
+                         false),
+    ParamsFor<vec3<u32>>(ast::Builtin::kNumWorkgroups,
+                         ast::PipelineStage::kCompute,
+                         true),
+
     ParamsFor<u32>(ast::Builtin::kSampleIndex,
                    ast::PipelineStage::kVertex,
                    false),
@@ -480,6 +490,7 @@ TEST_F(ResolverBuiltinsValidationTest, ComputeBuiltin_Pass) {
   //   [[builtin(local_invocationIndex)]] li_index: u32,
   //   [[builtin(global_invocationId)]] gi: vec3<u32>,
   //   [[builtin(workgroup_id)]] wi: vec3<u32>,
+  //   [[builtin(num_workgroups)]] nwgs: vec3<u32>,
   // ) {}
 
   auto* li_id =
@@ -493,8 +504,11 @@ TEST_F(ResolverBuiltinsValidationTest, ComputeBuiltin_Pass) {
             ast::DecorationList{Builtin(ast::Builtin::kGlobalInvocationId)});
   auto* wi = Param("wi", ty.vec3<u32>(),
                    ast::DecorationList{Builtin(ast::Builtin::kWorkgroupId)});
+  auto* nwgs =
+      Param("nwgs", ty.vec3<u32>(),
+            ast::DecorationList{Builtin(ast::Builtin::kNumWorkgroups)});
 
-  Func("main", ast::VariableList{li_id, li_index, gi, wi}, ty.void_(), {},
+  Func("main", ast::VariableList{li_id, li_index, gi, wi, nwgs}, ty.void_(), {},
        ast::DecorationList{
            Stage(ast::PipelineStage::kCompute),
            WorkgroupSize(Expr(Source{Source::Location{12, 34}}, 2))});
@@ -517,6 +531,21 @@ TEST_F(ResolverBuiltinsValidationTest, ComputeBuiltin_WorkGroupIdNotVec3U32) {
             "'vec3<u32>'");
 }
 
+TEST_F(ResolverBuiltinsValidationTest, ComputeBuiltin_NumWorkgroupsNotVec3U32) {
+  auto* nwgs = Param("nwgs", ty.f32(),
+                     ast::DecorationList{Builtin(
+                         Source{{12, 34}}, ast::Builtin::kNumWorkgroups)});
+  Func("main", ast::VariableList{nwgs}, ty.void_(), {},
+       ast::DecorationList{
+           Stage(ast::PipelineStage::kCompute),
+           WorkgroupSize(Expr(Source{Source::Location{12, 34}}, 2))});
+
+  EXPECT_FALSE(r()->Resolve());
+  EXPECT_EQ(r()->error(),
+            "12:34 error: store type of builtin(num_workgroups) must be "
+            "'vec3<u32>'");
+}
+
 TEST_F(ResolverBuiltinsValidationTest,
        ComputeBuiltin_GlobalInvocationNotVec3U32) {
   auto* gi = Param("gi", ty.vec3<i32>(),
diff --git a/src/resolver/resolver.cc b/src/resolver/resolver.cc
index 1a1143f820..4abe0711d2 100644
--- a/src/resolver/resolver.cc
+++ b/src/resolver/resolver.cc
@@ -1259,6 +1259,7 @@ bool Resolver::ValidateBuiltinDecoration(const ast::BuiltinDecoration* deco,
       break;
     case ast::Builtin::kGlobalInvocationId:
     case ast::Builtin::kLocalInvocationId:
+    case ast::Builtin::kNumWorkgroups:
     case ast::Builtin::kWorkgroupId:
       if (stage != ast::PipelineStage::kNone &&
           !(stage == ast::PipelineStage::kCompute && is_input)) {
@@ -1338,10 +1339,6 @@ bool Resolver::ValidateBuiltinDecoration(const ast::BuiltinDecoration* deco,
         return false;
       }
       break;
-    case ast::Builtin::kNumWorkgroups:
-      // TODO(crbug.com/tint/752): Backend support (needs extra work for HLSL).
-      AddError("num_workgroups builtin is not yet implemented", deco->source());
-      return false;
     default:
       break;
   }