[ir] Validate that subgroupBroadcast ID is const

Fixed: 445197082
Change-Id: Ic7b2085b9c985a1995e92b3ec667c01f754aa58f
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/262394
Reviewed-by: dan sinclair <[email protected]>
Auto-Submit: James Price <[email protected]>
Commit-Queue: dan sinclair <[email protected]>

Author: jrprice

src/tint/lang/core/ir/transform/builtin_polyfill_test.cc

@@ -2851,13 +2851,29 @@ TEST_F(IR_BuiltinPolyfillTest, Dot4U8Packed) {
     EXPECT_EQ(expect, str());
 }
 
-TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_NoPolyfill) {
-    Build(core::BuiltinFn::kSubgroupBroadcast, ty.f16(), Vector{ty.f16(), ty.u32()});
+class IR_SubgroupBroadcastPolyfillTest : public TransformTest {
+  protected:
+    /// Helper to build a function that calls subgroupBroadcast with the given type.
+    /// @param type the type
+    void Build(const core::type::Type* type) {
+        auto* param = b.FunctionParam("arg", type);
+        auto* func = b.Function("foo", type);
+        func->SetParams({param});
+        b.Append(func->Block(), [&] {
+            auto* result = b.Call(type, BuiltinFn::kSubgroupBroadcast, param, 1_u);
+            b.Return(func, result);
+            mod.SetName(result, "result");
+        });
+    }
+};
+
+TEST_F(IR_SubgroupBroadcastPolyfillTest, SubgroupBroadcastF16_NoPolyfill) {
+    Build(ty.f16());
 
     auto* src = R"(
-%foo = func(%arg:f16, %arg_1:u32):f16 {  # %arg_1: 'arg'
+%foo = func(%arg:f16):f16 {
   $B1: {
-    %result:f16 = subgroupBroadcast %arg, %arg_1
+    %result:f16 = subgroupBroadcast %arg, 1u
     ret %result
   }
 }
@@ -2873,13 +2889,13 @@ TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_NoPolyfill) {
     EXPECT_EQ(expect, str());
 }
 
-TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcast_F32_NotPolyfilled) {
-    Build(core::BuiltinFn::kSubgroupBroadcast, ty.f32(), Vector{ty.f32(), ty.u32()});
+TEST_F(IR_SubgroupBroadcastPolyfillTest, SubgroupBroadcastF32_NoPolyfill) {
+    Build(ty.f32());
 
     auto* src = R"(
-%foo = func(%arg:f32, %arg_1:u32):f32 {  # %arg_1: 'arg'
+%foo = func(%arg:f32):f32 {
   $B1: {
-    %result:f32 = subgroupBroadcast %arg, %arg_1
+    %result:f32 = subgroupBroadcast %arg, 1u
     ret %result
   }
 }
@@ -2895,27 +2911,27 @@ TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcast_F32_NotPolyfilled) {
     EXPECT_EQ(expect, str());
 }
 
-TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_Scalar) {
-    Build(core::BuiltinFn::kSubgroupBroadcast, ty.f16(), Vector{ty.f16(), ty.u32()});
+TEST_F(IR_SubgroupBroadcastPolyfillTest, SubgroupBroadcastF16_Scalar) {
+    Build(ty.f16());
 
     auto* src = R"(
-%foo = func(%arg:f16, %arg_1:u32):f16 {  # %arg_1: 'arg'
+%foo = func(%arg:f16):f16 {
   $B1: {
-    %result:f16 = subgroupBroadcast %arg, %arg_1
+    %result:f16 = subgroupBroadcast %arg, 1u
     ret %result
   }
 }
 )";
 
     auto* expect = R"(
-%foo = func(%arg:f16, %arg_1:u32):f16 {  # %arg_1: 'arg'
+%foo = func(%arg:f16):f16 {
   $B1: {
-    %4:vec2<f16> = construct %arg, 0.0h
-    %5:u32 = bitcast %4
-    %6:u32 = subgroupBroadcast %5, %arg_1
-    %7:vec2<f16> = bitcast %6
-    %8:f16 = access %7, 0u
-    ret %8
+    %3:vec2<f16> = construct %arg, 0.0h
+    %4:u32 = bitcast %3
+    %5:u32 = subgroupBroadcast %4, 1u
+    %6:vec2<f16> = bitcast %5
+    %7:f16 = access %6, 0u
+    ret %7
   }
 }
 )";
@@ -2929,25 +2945,25 @@ TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_Scalar) {
     EXPECT_EQ(expect, str());
 }
 
-TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_Vec2) {
-    Build(core::BuiltinFn::kSubgroupBroadcast, ty.vec2<f16>(), Vector{ty.vec2<f16>(), ty.u32()});
+TEST_F(IR_SubgroupBroadcastPolyfillTest, SubgroupBroadcastF16_Vec2) {
+    Build(ty.vec2<f16>());
 
     auto* src = R"(
-%foo = func(%arg:vec2<f16>, %arg_1:u32):vec2<f16> {  # %arg_1: 'arg'
+%foo = func(%arg:vec2<f16>):vec2<f16> {
   $B1: {
-    %result:vec2<f16> = subgroupBroadcast %arg, %arg_1
+    %result:vec2<f16> = subgroupBroadcast %arg, 1u
     ret %result
   }
 }
 )";
 
     auto* expect = R"(
-%foo = func(%arg:vec2<f16>, %arg_1:u32):vec2<f16> {  # %arg_1: 'arg'
+%foo = func(%arg:vec2<f16>):vec2<f16> {
   $B1: {
-    %4:u32 = bitcast %arg
-    %5:u32 = subgroupBroadcast %4, %arg_1
-    %6:vec2<f16> = bitcast %5
-    ret %6
+    %3:u32 = bitcast %arg
+    %4:u32 = subgroupBroadcast %3, 1u
+    %5:vec2<f16> = bitcast %4
+    ret %5
   }
 }
 )";
@@ -2961,27 +2977,27 @@ TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_Vec2) {
     EXPECT_EQ(expect, str());
 }
 
-TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_Vec3) {
-    Build(core::BuiltinFn::kSubgroupBroadcast, ty.vec3<f16>(), Vector{ty.vec3<f16>(), ty.u32()});
+TEST_F(IR_SubgroupBroadcastPolyfillTest, SubgroupBroadcastF16_Vec3) {
+    Build(ty.vec3<f16>());
 
     auto* src = R"(
-%foo = func(%arg:vec3<f16>, %arg_1:u32):vec3<f16> {  # %arg_1: 'arg'
+%foo = func(%arg:vec3<f16>):vec3<f16> {
   $B1: {
-    %result:vec3<f16> = subgroupBroadcast %arg, %arg_1
+    %result:vec3<f16> = subgroupBroadcast %arg, 1u
     ret %result
   }
 }
 )";
 
     auto* expect = R"(
-%foo = func(%arg:vec3<f16>, %arg_1:u32):vec3<f16> {  # %arg_1: 'arg'
+%foo = func(%arg:vec3<f16>):vec3<f16> {
   $B1: {
-    %4:vec4<f16> = construct %arg, 0.0h
-    %5:vec2<u32> = bitcast %4
-    %6:vec2<u32> = subgroupBroadcast %5, %arg_1
-    %7:vec4<f16> = bitcast %6
-    %8:vec3<f16> = swizzle %7, xyz
-    ret %8
+    %3:vec4<f16> = construct %arg, 0.0h
+    %4:vec2<u32> = bitcast %3
+    %5:vec2<u32> = subgroupBroadcast %4, 1u
+    %6:vec4<f16> = bitcast %5
+    %7:vec3<f16> = swizzle %6, xyz
+    ret %7
   }
 }
 )";
@@ -2995,25 +3011,25 @@ TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_Vec3) {
     EXPECT_EQ(expect, str());
 }
 
-TEST_F(IR_BuiltinPolyfillTest, SubgroupBroadcastF16_Vec4) {
-    Build(core::BuiltinFn::kSubgroupBroadcast, ty.vec4<f16>(), Vector{ty.vec4<f16>(), ty.u32()});
+TEST_F(IR_SubgroupBroadcastPolyfillTest, SubgroupBroadcastF16_Vec4) {
+    Build(ty.vec4<f16>());
 
     auto* src = R"(
-%foo = func(%arg:vec4<f16>, %arg_1:u32):vec4<f16> {  # %arg_1: 'arg'
+%foo = func(%arg:vec4<f16>):vec4<f16> {
   $B1: {
-    %result:vec4<f16> = subgroupBroadcast %arg, %arg_1
+    %result:vec4<f16> = subgroupBroadcast %arg, 1u
     ret %result
   }
 }
 )";
 
     auto* expect = R"(
-%foo = func(%arg:vec4<f16>, %arg_1:u32):vec4<f16> {  # %arg_1: 'arg'
+%foo = func(%arg:vec4<f16>):vec4<f16> {
   $B1: {
-    %4:vec2<u32> = bitcast %arg
-    %5:vec2<u32> = subgroupBroadcast %4, %arg_1
-    %6:vec4<f16> = bitcast %5
-    ret %6
+    %3:vec2<u32> = bitcast %arg
+    %4:vec2<u32> = subgroupBroadcast %3, 1u
+    %5:vec4<f16> = bitcast %4
+    ret %5
   }
 }
 )";

src/tint/lang/core/ir/validator.cc

@@ -3251,7 +3251,8 @@ void Validator::CheckBuiltinCall(const BuiltinCall* call) {
 
 void Validator::CheckCoreBuiltinCall(const CoreBuiltinCall* call,
                                      const core::intrinsic::Overload& overload) {
-    if (call->Func() == core::BuiltinFn::kQuadBroadcast) {
+    if (call->Func() == core::BuiltinFn::kQuadBroadcast ||
+        call->Func() == core::BuiltinFn::kSubgroupBroadcast) {
         TINT_ASSERT(call->Args().Length() == 2);
         constexpr uint32_t kIdArg = 1;
         auto* id = call->Args()[kIdArg];

src/tint/lang/core/ir/validator_call_test.cc

@@ -635,4 +635,23 @@ TEST_F(IR_ValidatorTest, CallBuiltinFn_QuadBroadcast_NonConstId) {
 )")) << res.Failure();
 }
 
+TEST_F(IR_ValidatorTest, CallBuiltinFn_SubgroupBroadcast_NonConstId) {
+    auto* func = b.Function("foo", ty.void_(), core::ir::Function::PipelineStage::kFragment);
+    b.Append(func->Block(), [&] {
+        auto* id = b.Let("a", 2_u);
+        b.Let("x", b.Call<i32>(core::BuiltinFn::kSubgroupBroadcast, 2_i, id));
+        b.Return(func);
+    });
+
+    auto res = ir::Validate(mod);
+    ASSERT_NE(res, Success);
+    EXPECT_THAT(res.Failure().reason,
+                testing::HasSubstr(
+                    R"(:4:36 error: subgroupBroadcast: non-constant ID provided
+    %3:i32 = subgroupBroadcast 2i, %a
+                                   ^^
+
+)")) << res.Failure();
+}
+
 }  // namespace tint::core::ir