[tint][ir][val] Improve check for @workgroup_size params

- Adds in tests that overrides are only module scoped if use for
  params here

Issue: 376624999
Change-Id: Ie138c623f2882da0a64ed776dcef2978779463e4
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/217877
Commit-Queue: Ryan Harrison <rharrison@chromium.org>
Auto-Submit: Ryan Harrison <rharrison@chromium.org>
Commit-Queue: James Price <jrprice@google.com>
Reviewed-by: James Price <jrprice@google.com>
diff --git a/src/tint/lang/core/ir/validator.cc b/src/tint/lang/core/ir/validator.cc
index 4d789b9..c96bda2 100644
--- a/src/tint/lang/core/ir/validator.cc
+++ b/src/tint/lang/core/ir/validator.cc
@@ -2160,7 +2160,17 @@
         }
 
         if (auto* r = size->As<ir::InstructionResult>()) {
-            if (r->Instruction() && r->Instruction()->Is<core::ir::Override>()) {
+            if (!r->Instruction()) {
+                AddError(func) << "instruction for @workgroup_size param is not defined";
+                return;
+            }
+
+            if (r->Instruction()->Block() != mod_.root_block) {
+                AddError(func) << "@workgroup_size param defined by non-module scope value";
+                return;
+            }
+
+            if (r->Instruction()->Is<core::ir::Override>()) {
                 continue;
             }
 
diff --git a/src/tint/lang/core/ir/validator_test.cc b/src/tint/lang/core/ir/validator_test.cc
index 74cf423..5e4dc7b 100644
--- a/src/tint/lang/core/ir/validator_test.cc
+++ b/src/tint/lang/core/ir/validator_test.cc
@@ -1362,6 +1362,32 @@
 )");
 }
 
+TEST_F(IR_ValidatorTest, Function_WorkgroupSize_NonRootBlockOverride) {
+    auto* f = ComputeEntryPoint();
+    Override* o;
+    b.Append(f->Block(), [&] {
+        o = b.Override(ty.u32());
+        b.Return(f);
+    });
+    f->SetWorkgroupSize({o->Result(0), o->Result(0), o->Result(0)});
+
+    auto res = ir::Validate(mod, Capabilities{Capability::kAllowOverrides});
+    ASSERT_NE(res, Success);
+    EXPECT_EQ(res.Failure().reason.Str(),
+              R"(:1:1 error: @workgroup_size param defined by non-module scope value
+%f = @compute @workgroup_size(%2, %2, %2) func():void {
+^^
+
+note: # Disassembly
+%f = @compute @workgroup_size(%2, %2, %2) func():void {
+  $B1: {
+    %2:u32 = override @id(0)
+    ret
+  }
+}
+)");
+}
+
 TEST_F(IR_ValidatorTest, Function_Vertex_BasicPosition) {
     auto* f = b.Function("my_func", ty.vec4<f32>(), Function::PipelineStage::kVertex);
     f->SetReturnBuiltin(BuiltinValue::kPosition);