From c50a11059c92ab26a349732afcf5e0fb2afda438 Mon Sep 17 00:00:00 2001
From: Natalie Chouinard <chouinard@google.com>
Date: Wed, 03 Jun 2026 19:30:48 -0700
Subject: [PATCH] [tint][spirv] Clamp subgroupShuffle to actual size

Replace the hardcoded max subgroup size mask in the SPIR-V
subgroupShuffle polyfill with the actual subgroup size from the builtin
value. Because SubgroupSize is guaranteed to be a power-of-2 in Vulkan,
it is still safe to subtract 1 and mask to clamp.

This polyfill is also now enabled for all SPIR-V generation, rather than
limited by a toggle, since the requirement of invocation ID to be less
than the subgroup size is a general SPIR-V spec requirement to avoid
poison values, not specific to a driver issue.

https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpGroupNonUniformShuffle

Fixes: 513947572
Change-Id: Id6e68170e467284152212e8c014d2c9688760175
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/313515
Commit-Queue: Natalie Chouinard <chouinard@google.com>
SLSA-Policy-Verified: SLSA Policy Verification Service <devtools-gerritcodereview-exitgate@google.com>
Reviewed-by: James Price <jrprice@google.com>
---
diff --git a/src/dawn/native/Toggles.cpp b/src/dawn/native/Toggles.cpp
index 58df017..a18adb6 100644
--- a/src/dawn/native/Toggles.cpp
+++ b/src/dawn/native/Toggles.cpp
@@ -622,10 +622,6 @@
     {Toggle::MetalPolyfillClampFloat,
      {"metal_polyfill_clamp_float", "Polyfill clamp function for floating point (metal).",
       "https://crbug.com/407109056", ToggleStage::Device}},
-    {Toggle::SubgroupShuffleClamped,
-     {"subgroup_shuffle_clamped",
-      "Polyfill subgroupShuffle by clamping the id param to within maximum possible subgroup size.",
-      "https://crbug.com/dawn/2502", ToggleStage::Device}},
     {Toggle::VulkanSampleCompareDepthCubeArrayWorkaround,
      {"vulkan_sample_compare_depth_cube_array_workaround",
       "Polyfill specific case of textureSampleCompare(Level) with workaround.",
diff --git a/src/dawn/native/Toggles.h b/src/dawn/native/Toggles.h
index 38a8a90..02fde65 100644
--- a/src/dawn/native/Toggles.h
+++ b/src/dawn/native/Toggles.h
@@ -152,7 +152,6 @@
     ScalarizeMaxMinClamp,
     SaturateAsMinMaxF16,
     MetalPolyfillClampFloat,
-    SubgroupShuffleClamped,
     VulkanSampleCompareDepthCubeArrayWorkaround,
     VulkanSampleCompare2DWorkaround,
     MetalDisableModuleConstantF16,
diff --git a/src/dawn/native/vulkan/PhysicalDeviceVk.cpp b/src/dawn/native/vulkan/PhysicalDeviceVk.cpp
index 2176cd0..82a60c5 100644
--- a/src/dawn/native/vulkan/PhysicalDeviceVk.cpp
+++ b/src/dawn/native/vulkan/PhysicalDeviceVk.cpp
@@ -1109,10 +1109,6 @@
         // chromium:448873316: Non-scalar (vector) saturate from uniform fails.
         deviceToggles->Default(Toggle::SaturateAsMinMaxF16, true);
 
-        // Polyfill a clamp of `id` param in subgroupShuffle to follow spec limitations.
-        // See crbug.com/435246627
-        deviceToggles->Default(Toggle::SubgroupShuffleClamped, true);
-
         if (gpu_info::IsIntelGen12LP(GetVendorId(), GetDeviceId())) {
             // dawn:1688: Intel Mesa driver has a bug about reusing the VkDeviceMemory that was
             // previously bound to a 2D VkImage. To work around that bug we have to disable the
diff --git a/src/dawn/native/vulkan/ShaderModuleVk.cpp b/src/dawn/native/vulkan/ShaderModuleVk.cpp
index 74fe3d7..a3c77c7 100644
--- a/src/dawn/native/vulkan/ShaderModuleVk.cpp
+++ b/src/dawn/native/vulkan/ShaderModuleVk.cpp
@@ -317,8 +317,6 @@
     req.tintOptions.extensions.use_uniform_buffers =
         !GetDevice()->IsToggleEnabled(Toggle::DecomposeUniformBuffers);
 
-    req.tintOptions.workarounds.subgroup_shuffle_clamped =
-        GetDevice()->IsToggleEnabled(Toggle::SubgroupShuffleClamped);
     req.tintOptions.workarounds.texture_sample_compare_depth_cube_array =
         GetDevice()->IsToggleEnabled(Toggle::VulkanSampleCompareDepthCubeArrayWorkaround);
     req.tintOptions.workarounds.texture_sample_compare_2d_polyfill =
diff --git a/src/tint/lang/spirv/writer/common/options.h b/src/tint/lang/spirv/writer/common/options.h
index 8c4d22d..c9a5dd9 100644
--- a/src/tint/lang/spirv/writer/common/options.h
+++ b/src/tint/lang/spirv/writer/common/options.h
@@ -87,10 +87,6 @@
         /// and `unpack4x8unorm` builtins
         bool polyfill_pack_unpack_4x8_norm = false;
 
-        /// Set to `true` to generate a polyfill clamp of `id` param of subgroupShuffle to within
-        /// the spec max subgroup size.
-        bool subgroup_shuffle_clamped = false;
-
         /// Set to 'true' to force workaround for 'textureSampleCompare(Level)' for texture arrays
         /// of cube depth.
         bool texture_sample_compare_depth_cube_array = false;
@@ -132,7 +128,6 @@
                      scalarize_max_min_clamp,
                      dva_transform_handle,
                      polyfill_pack_unpack_4x8_norm,
-                     subgroup_shuffle_clamped,
                      texture_sample_compare_depth_cube_array,
                      texture_sample_compare_2d_polyfill,
                      polyfill_subgroup_broadcast_f16,
diff --git a/src/tint/lang/spirv/writer/raise/builtin_polyfill.cc b/src/tint/lang/spirv/writer/raise/builtin_polyfill.cc
index 326b351..7d6c803 100644
--- a/src/tint/lang/spirv/writer/raise/builtin_polyfill.cc
+++ b/src/tint/lang/spirv/writer/raise/builtin_polyfill.cc
@@ -184,6 +184,9 @@
     /// The type manager.
     core::type::Manager& ty{ir.Types()};
 
+    /// Module-scoped variable for subgroup size mask.
+    core::ir::Var* subgroup_size_mask_ = nullptr;
+
     /// Process the module.
     void Process() {
         // Find the builtins that need replacing.
@@ -249,9 +252,7 @@
                     case core::BuiltinFn::kSubgroupShuffleDown:
                     case core::BuiltinFn::kSubgroupShuffleUp:
                     case core::BuiltinFn::kSubgroupShuffleXor: {
-                        bool clamped = config.subgroup_shuffle_clamped;
-                        worklist.push_back(
-                            [this, builtin, clamped] { SubgroupShuffle(builtin, clamped); });
+                        worklist.push_back([this, builtin] { SubgroupShuffle(builtin); });
                         break;
                     }
                     case core::BuiltinFn::kTextureDimensions:
@@ -364,6 +365,48 @@
             });
             construct->SetArg(0, value);
         }
+
+        if (subgroup_size_mask_) {
+            for (auto func : ir.functions) {
+                if (func->IsEntryPoint()) {
+                    SetSubgroupSizeMaskForEntryPoint(func);
+                }
+            }
+        }
+    }
+
+    /// Set the subgroup_size_mask variable from an entry point.
+    void SetSubgroupSizeMaskForEntryPoint(core::ir::Function* ep) {
+        b.InsertBefore(ep->Block()->Front(), [&] {
+            core::ir::Value* subgroup_size = nullptr;
+            for (auto* param : ep->Params()) {
+                if (param->Attributes().builtin == core::BuiltinValue::kSubgroupSize) {
+                    subgroup_size = param;
+                    break;
+                }
+                if (auto* str = param->Type()->As<core::type::Struct>()) {
+                    for (auto* member : str->Members()) {
+                        if (member->Attributes().builtin == core::BuiltinValue::kSubgroupSize) {
+                            subgroup_size =
+                                b.Access(ty.u32(), param, u32(member->Index()))->Result();
+                            break;
+                        }
+                    }
+                    if (subgroup_size) {
+                        break;
+                    }
+                }
+            }
+            if (!subgroup_size) {
+                auto* param = b.FunctionParam("tint_subgroup_size", ty.u32());
+                param->SetBuiltin(core::BuiltinValue::kSubgroupSize);
+                ep->AppendParam(param);
+                subgroup_size = param;
+            }
+
+            auto* mask = b.Subtract(subgroup_size, 1_u);
+            b.Store(subgroup_size_mask_, mask);
+        });
     }
 
     /// Create a literal operand.
@@ -1248,10 +1291,10 @@
     /// Handles SubgroupShuffle(), SubgroupShuffleDown(), SubgroupShuffleUp(), SubgroupShuffleXor()
     /// builtins.
     /// @param builtin the builtin call instruction
-    void SubgroupShuffle(core::ir::CoreBuiltinCall* builtin, bool clamp_subgroup_shuffle) {
+    void SubgroupShuffle(core::ir::CoreBuiltinCall* builtin) {
         TINT_IR_ASSERT(ir, builtin->Args().size() == 2);
         // The second argument is either 'id' , 'delta', or 'mask'.
-        // All must be bound by [0, 128)
+        // All must be bound by [0, subgroup_size)
         auto* arg2 = builtin->Args()[1];
         // arg2 must be an unsigned integer scalar, so bitcast if necessary.
         if (arg2->Type()->IsSignedIntegerScalar()) {
@@ -1261,15 +1304,18 @@
         }
 
         /// Polyfill a `subgroupShuffleX` builtin call with one that has clamped the arg2 param
-        if (clamp_subgroup_shuffle) {
-            auto* shuffle_id = builtin->Args()[1];
-            auto* mask_max_subgroup_size =
-                b.Constant(core::u32(tint::internal_limits::kMaxSubgroupSize - 1));
-            b.InsertBefore(builtin, [&] {
-                auto* clamp_via_masking_and = b.And(shuffle_id, mask_max_subgroup_size);
-                builtin->SetArg(1, clamp_via_masking_and->Result());
+        auto* shuffle_id = builtin->Args()[1];
+        if (!subgroup_size_mask_) {
+            b.Append(ir.root_block, [&] {
+                subgroup_size_mask_ =
+                    b.Var<core::AddressSpace::kPrivate, u32>("tint_subgroup_size_mask");
             });
         }
+        b.InsertBefore(builtin, [&] {
+            auto* subgroup_size_mask = b.Load(subgroup_size_mask_);
+            auto* clamp_via_masking_and = b.And(shuffle_id, subgroup_size_mask);
+            builtin->SetArg(1, clamp_via_masking_and->Result());
+        });
     }
 
     /// Handle a SubgroupBroadcast() builtin.
diff --git a/src/tint/lang/spirv/writer/raise/builtin_polyfill.h b/src/tint/lang/spirv/writer/raise/builtin_polyfill.h
index c03f556..4e19350 100644
--- a/src/tint/lang/spirv/writer/raise/builtin_polyfill.h
+++ b/src/tint/lang/spirv/writer/raise/builtin_polyfill.h
@@ -41,7 +41,6 @@
 struct PolyfillConfig {
     bool use_vulkan_memory_model = false;
     SpvVersion version = SpvVersion::kSpv13;
-    bool subgroup_shuffle_clamped = false;
     bool texture_sample_compare_depth_cube_array = false;
     bool texture_sample_compare_2d_polyfill = false;
     bool cooperative_matrix_stride_is_matrix_elements = false;
diff --git a/src/tint/lang/spirv/writer/raise/raise.cc b/src/tint/lang/spirv/writer/raise/raise.cc
index cd58980..1f41cc4 100644
--- a/src/tint/lang/spirv/writer/raise/raise.cc
+++ b/src/tint/lang/spirv/writer/raise/raise.cc
@@ -214,7 +214,6 @@
     raise::PolyfillConfig config = {
         .use_vulkan_memory_model = options.extensions.use_vulkan_memory_model,
         .version = options.spirv_version,
-        .subgroup_shuffle_clamped = options.workarounds.subgroup_shuffle_clamped,
         .texture_sample_compare_depth_cube_array =
             options.workarounds.texture_sample_compare_depth_cube_array,
         .texture_sample_compare_2d_polyfill =
diff --git a/src/tint/lang/spirv/writer/writer_fuzz.cc b/src/tint/lang/spirv/writer/writer_fuzz.cc
index 130b0b0..5b4d36b 100644
--- a/src/tint/lang/spirv/writer/writer_fuzz.cc
+++ b/src/tint/lang/spirv/writer/writer_fuzz.cc
@@ -312,7 +312,6 @@
     options.workarounds.dva_transform_handle = fuzzed_options.dva_transform_handle;
     options.workarounds.polyfill_pack_unpack_4x8_norm =
         fuzzed_options.polyfill_pack_unpack_4x8_norm;
-    options.workarounds.subgroup_shuffle_clamped = fuzzed_options.subgroup_shuffle_clamped;
     options.workarounds.polyfill_subgroup_broadcast_f16 =
         fuzzed_options.polyfill_subgroup_broadcast_f16;
     options.workarounds.pass_matrix_by_pointer = fuzzed_options.pass_matrix_by_pointer;
diff --git a/test/tint/builtins/gen/literal/subgroupShuffle/030422.wgsl.expected.spvasm b/test/tint/builtins/gen/literal/subgroupShuffle/030422.wgsl.expected.spvasm
index 619fb81..fc8da8f 100644
--- a/test/tint/builtins/gen/literal/subgroupShuffle/030422.wgsl.expected.spvasm
+++ b/test/tint/builtins/gen/literal/subgroupShuffle/030422.wgsl.expected.spvasm
@@ -4,52 +4,77 @@
 ; SPIR-V
 ; Version: 1.3
 ; Generator: Google Tint Compiler; 1
-; Bound: 26
+; Bound: 41
 ; Schema: 0
                OpCapability Shader
+               OpCapability GroupNonUniform
                OpCapability GroupNonUniformShuffle
                OpMemoryModel Logical GLSL450
-               OpEntryPoint Fragment %fragment_main "fragment_main"
+               OpEntryPoint Fragment %fragment_main "fragment_main" %fragment_main_subgroup_size_Input
                OpExecutionMode %fragment_main OriginUpperLeft
                OpMemberName %prevent_dce_block 0 "inner"
                OpName %prevent_dce_block "prevent_dce_block"
+               OpName %tint_subgroup_size_mask "tint_subgroup_size_mask"
+               OpName %fragment_main_subgroup_size_Input "fragment_main_subgroup_size_Input"
                OpName %subgroupShuffle_030422 "subgroupShuffle_030422"
                OpName %res "res"
+               OpName %fragment_main_inner "fragment_main_inner"
+               OpName %tint_subgroup_size "tint_subgroup_size"
                OpName %fragment_main "fragment_main"
                OpMemberDecorate %prevent_dce_block 0 Offset 0
                OpDecorate %prevent_dce_block Block
                OpDecorate %1 DescriptorSet 0
                OpDecorate %1 Binding 0
                OpDecorate %1 Coherent
+               OpDecorate %fragment_main_subgroup_size_Input Flat
+               OpDecorate %fragment_main_subgroup_size_Input BuiltIn SubgroupSize
       %float = OpTypeFloat 32
 %prevent_dce_block = OpTypeStruct %float
 %_ptr_StorageBuffer_prevent_dce_block = OpTypePointer StorageBuffer %prevent_dce_block
           %1 = OpVariable %_ptr_StorageBuffer_prevent_dce_block StorageBuffer
-          %6 = OpTypeFunction %float
        %uint = OpTypeInt 32 0
+%_ptr_Private_uint = OpTypePointer Private %uint
+          %8 = OpConstantNull %uint
+%tint_subgroup_size_mask = OpVariable %_ptr_Private_uint Private %8
+%_ptr_Input_uint = OpTypePointer Input %uint
+%fragment_main_subgroup_size_Input = OpVariable %_ptr_Input_uint Input
+         %12 = OpTypeFunction %float
         %int = OpTypeInt 32 1
       %int_1 = OpConstant %int 1
      %uint_3 = OpConstant %uint 3
     %float_1 = OpConstant %float 1
 %_ptr_Function_float = OpTypePointer Function %float
        %void = OpTypeVoid
-         %20 = OpTypeFunction %void
+         %28 = OpTypeFunction %void %uint
+     %uint_1 = OpConstant %uint 1
 %_ptr_StorageBuffer_float = OpTypePointer StorageBuffer %float
      %uint_0 = OpConstant %uint 0
-%subgroupShuffle_030422 = OpFunction %float None %6
-          %7 = OpLabel
+         %37 = OpTypeFunction %void
+%subgroupShuffle_030422 = OpFunction %float None %12
+         %13 = OpLabel
         %res = OpVariable %_ptr_Function_float Function
-          %9 = OpBitcast %uint %int_1
-         %12 = OpGroupNonUniformShuffle %float %uint_3 %float_1 %9
-               OpStore %res %12
-         %17 = OpLoad %float %res None
-               OpReturnValue %17
+         %14 = OpBitcast %uint %int_1
+         %17 = OpLoad %uint %tint_subgroup_size_mask None
+         %18 = OpBitwiseAnd %uint %14 %17
+         %19 = OpGroupNonUniformShuffle %float %uint_3 %float_1 %18
+               OpStore %res %19
+         %24 = OpLoad %float %res None
+               OpReturnValue %24
                OpFunctionEnd
-%fragment_main = OpFunction %void None %20
-         %21 = OpLabel
-         %22 = OpFunctionCall %float %subgroupShuffle_030422
-         %23 = OpAccessChain %_ptr_StorageBuffer_float %1 %uint_0
-               OpStore %23 %22 None
+%fragment_main_inner = OpFunction %void None %28
+%tint_subgroup_size = OpFunctionParameter %uint
+         %29 = OpLabel
+         %30 = OpISub %uint %tint_subgroup_size %uint_1
+               OpStore %tint_subgroup_size_mask %30 None
+         %32 = OpFunctionCall %float %subgroupShuffle_030422
+         %33 = OpAccessChain %_ptr_StorageBuffer_float %1 %uint_0
+               OpStore %33 %32 None
+               OpReturn
+               OpFunctionEnd
+%fragment_main = OpFunction %void None %37
+         %38 = OpLabel
+         %39 = OpLoad %uint %fragment_main_subgroup_size_Input None
+         %40 = OpFunctionCall %void %fragment_main_inner %39
                OpReturn
                OpFunctionEnd
 ;
@@ -58,51 +83,75 @@
 ; SPIR-V
 ; Version: 1.3
 ; Generator: Google Tint Compiler; 1
-; Bound: 26
+; Bound: 41
 ; Schema: 0
                OpCapability Shader
+               OpCapability GroupNonUniform
                OpCapability GroupNonUniformShuffle
                OpMemoryModel Logical GLSL450
-               OpEntryPoint GLCompute %compute_main "compute_main"
+               OpEntryPoint GLCompute %compute_main "compute_main" %compute_main_subgroup_size_Input
                OpExecutionMode %compute_main LocalSize 1 1 1
                OpMemberName %prevent_dce_block 0 "inner"
                OpName %prevent_dce_block "prevent_dce_block"
+               OpName %tint_subgroup_size_mask "tint_subgroup_size_mask"
+               OpName %compute_main_subgroup_size_Input "compute_main_subgroup_size_Input"
                OpName %subgroupShuffle_030422 "subgroupShuffle_030422"
                OpName %res "res"
+               OpName %compute_main_inner "compute_main_inner"
+               OpName %tint_subgroup_size "tint_subgroup_size"
                OpName %compute_main "compute_main"
                OpMemberDecorate %prevent_dce_block 0 Offset 0
                OpDecorate %prevent_dce_block Block
                OpDecorate %1 DescriptorSet 0
                OpDecorate %1 Binding 0
                OpDecorate %1 Coherent
+               OpDecorate %compute_main_subgroup_size_Input BuiltIn SubgroupSize
       %float = OpTypeFloat 32
 %prevent_dce_block = OpTypeStruct %float
 %_ptr_StorageBuffer_prevent_dce_block = OpTypePointer StorageBuffer %prevent_dce_block
           %1 = OpVariable %_ptr_StorageBuffer_prevent_dce_block StorageBuffer
-          %6 = OpTypeFunction %float
        %uint = OpTypeInt 32 0
+%_ptr_Private_uint = OpTypePointer Private %uint
+          %8 = OpConstantNull %uint
+%tint_subgroup_size_mask = OpVariable %_ptr_Private_uint Private %8
+%_ptr_Input_uint = OpTypePointer Input %uint
+%compute_main_subgroup_size_Input = OpVariable %_ptr_Input_uint Input
+         %12 = OpTypeFunction %float
         %int = OpTypeInt 32 1
       %int_1 = OpConstant %int 1
      %uint_3 = OpConstant %uint 3
     %float_1 = OpConstant %float 1
 %_ptr_Function_float = OpTypePointer Function %float
        %void = OpTypeVoid
-         %20 = OpTypeFunction %void
+         %28 = OpTypeFunction %void %uint
+     %uint_1 = OpConstant %uint 1
 %_ptr_StorageBuffer_float = OpTypePointer StorageBuffer %float
... (file diff truncated)

diff --git a/src/tint/lang/spirv/writer/raise/builtin_polyfill_test.cc b/src/tint/lang/spirv/writer/raise/builtin_polyfill_test.cc
index 6b406ed..5004c46 100644
--- a/src/tint/lang/spirv/writer/raise/builtin_polyfill_test.cc
+++ b/src/tint/lang/spirv/writer/raise/builtin_polyfill_test.cc
@@ -1387,19 +1387,22 @@
 }
 
 TEST_F(SpirvWriter_BuiltinPolyfillTest, SubgroupShuffleDown_Clamped) {
-    auto* val = b.FunctionParam("val", ty.i32());
-    auto* delta = b.FunctionParam("delta", ty.u32());
     auto* func = b.Function("foo", ty.i32());
-    func->SetParams({val, delta});
+    func->SetStage(core::ir::Function::PipelineStage::kFragment);
+    func->SetReturnLocation(0);
 
     b.Append(func->Block(), [&] {
+        auto* val = b.Let("val", 1_i);
+        auto* delta = b.Let("delta", 1_u);
         auto* result = b.Call(ty.i32(), core::BuiltinFn::kSubgroupShuffleDown, val, delta);
         b.Return(func, result);
     });
 
     auto* src = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
+%foo = @fragment func():i32 [@location(0)] {
   $B1: {
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
     %4:i32 = subgroupShuffleDown %val, %delta
     ret %4
   }
@@ -1408,35 +1411,46 @@
     EXPECT_EQ(src, str());
 
     auto* expect = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
-  $B1: {
-    %4:u32 = and %delta, 127u
-    %5:i32 = subgroupShuffleDown %val, %4
-    ret %5
+$B1: {  # root
+  %tint_subgroup_size_mask:ptr<private, u32, read_write> = var undef
+}
+
+%foo = @fragment func(%tint_subgroup_size:u32 [@subgroup_size]):i32 [@location(0)] {
+  $B2: {
+    %4:u32 = sub %tint_subgroup_size, 1u
+    store %tint_subgroup_size_mask, %4
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
+    %7:u32 = load %tint_subgroup_size_mask
+    %8:u32 = and %delta, %7
+    %9:i32 = subgroupShuffleDown %val, %8
+    ret %9
   }
 }
 )";
 
-    PolyfillConfig config{.subgroup_shuffle_clamped = true};
-    Run(BuiltinPolyfill, config);
+    Run(BuiltinPolyfill, PolyfillConfig{});
 
     EXPECT_EQ(expect, str());
 }
 
 TEST_F(SpirvWriter_BuiltinPolyfillTest, SubgroupShuffleDown_SignedDelta_Clamped) {
-    auto* val = b.FunctionParam("val", ty.i32());
-    auto* delta = b.FunctionParam("delta", ty.u32());
     auto* func = b.Function("foo", ty.i32());
-    func->SetParams({val, delta});
+    func->SetStage(core::ir::Function::PipelineStage::kFragment);
+    func->SetReturnLocation(0);
 
     b.Append(func->Block(), [&] {
+        auto* val = b.Let("val", 1_i);
+        auto* delta = b.Let("delta", 1_u);
         auto* result = b.Call(ty.i32(), core::BuiltinFn::kSubgroupShuffleDown, val, delta);
         b.Return(func, result);
     });
 
     auto* src = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
+%foo = @fragment func():i32 [@location(0)] {
   $B1: {
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
     %4:i32 = subgroupShuffleDown %val, %delta
     ret %4
   }
@@ -1445,98 +1459,46 @@
     EXPECT_EQ(src, str());
 
     auto* expect = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
-  $B1: {
-    %4:u32 = and %delta, 127u
-    %5:i32 = subgroupShuffleDown %val, %4
-    ret %5
+$B1: {  # root
+  %tint_subgroup_size_mask:ptr<private, u32, read_write> = var undef
+}
+
+%foo = @fragment func(%tint_subgroup_size:u32 [@subgroup_size]):i32 [@location(0)] {
+  $B2: {
+    %4:u32 = sub %tint_subgroup_size, 1u
+    store %tint_subgroup_size_mask, %4
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
+    %7:u32 = load %tint_subgroup_size_mask
+    %8:u32 = and %delta, %7
+    %9:i32 = subgroupShuffleDown %val, %8
+    ret %9
   }
 }
 )";
 
-    PolyfillConfig config{.subgroup_shuffle_clamped = true};
-    Run(BuiltinPolyfill, config);
-
-    EXPECT_EQ(expect, str());
-}
-
-TEST_F(SpirvWriter_BuiltinPolyfillTest, SubgroupShuffleDown_Unclamped) {
-    auto* val = b.FunctionParam("val", ty.i32());
-    auto* delta = b.FunctionParam("delta", ty.u32());
-    auto* func = b.Function("foo", ty.i32());
-    func->SetParams({val, delta});
-
-    b.Append(func->Block(), [&] {
-        auto* result = b.Call(ty.i32(), core::BuiltinFn::kSubgroupShuffleDown, val, delta);
-        b.Return(func, result);
-    });
-
-    auto* src = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
-  $B1: {
-    %4:i32 = subgroupShuffleDown %val, %delta
-    ret %4
-  }
-}
-)";
-    EXPECT_EQ(src, str());
-
-    PolyfillConfig config{.subgroup_shuffle_clamped = false};
-    Run(BuiltinPolyfill, config);
-
-    EXPECT_EQ(src, str());
-}
-
-TEST_F(SpirvWriter_BuiltinPolyfillTest, SubgroupShuffleDown_SignedDelta_Unclamped) {
-    auto* val = b.FunctionParam("val", ty.i32());
-    auto* delta = b.FunctionParam("delta", ty.u32());
-    auto* func = b.Function("foo", ty.i32());
-    func->SetParams({val, delta});
-
-    b.Append(func->Block(), [&] {
-        auto* result = b.Call(ty.i32(), core::BuiltinFn::kSubgroupShuffleDown, val, delta);
-        b.Return(func, result);
-    });
-
-    auto* src = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
-  $B1: {
-    %4:i32 = subgroupShuffleDown %val, %delta
-    ret %4
-  }
-}
-)";
-    EXPECT_EQ(src, str());
-
-    auto* expect = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
-  $B1: {
-    %4:i32 = subgroupShuffleDown %val, %delta
-    ret %4
-  }
-}
-)";
-
-    PolyfillConfig config{.subgroup_shuffle_clamped = false};
-    Run(BuiltinPolyfill, config);
+    Run(BuiltinPolyfill, PolyfillConfig{});
 
     EXPECT_EQ(expect, str());
 }
 
 TEST_F(SpirvWriter_BuiltinPolyfillTest, SubgroupShuffleUp_Clamped) {
-    auto* val = b.FunctionParam("val", ty.i32());
-    auto* delta = b.FunctionParam("delta", ty.u32());
     auto* func = b.Function("foo", ty.i32());
-    func->SetParams({val, delta});
+    func->SetStage(core::ir::Function::PipelineStage::kFragment);
+    func->SetReturnLocation(0);
 
     b.Append(func->Block(), [&] {
+        auto* val = b.Let("val", 1_i);
+        auto* delta = b.Let("delta", 1_u);
         auto* result = b.Call(ty.i32(), core::BuiltinFn::kSubgroupShuffleUp, val, delta);
         b.Return(func, result);
     });
 
     auto* src = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
+%foo = @fragment func():i32 [@location(0)] {
   $B1: {
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
     %4:i32 = subgroupShuffleUp %val, %delta
     ret %4
   }
@@ -1545,35 +1507,46 @@
     EXPECT_EQ(src, str());
 
     auto* expect = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
-  $B1: {
-    %4:u32 = and %delta, 127u
-    %5:i32 = subgroupShuffleUp %val, %4
-    ret %5
+$B1: {  # root
+  %tint_subgroup_size_mask:ptr<private, u32, read_write> = var undef
+}
+
+%foo = @fragment func(%tint_subgroup_size:u32 [@subgroup_size]):i32 [@location(0)] {
+  $B2: {
+    %4:u32 = sub %tint_subgroup_size, 1u
+    store %tint_subgroup_size_mask, %4
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
+    %7:u32 = load %tint_subgroup_size_mask
+    %8:u32 = and %delta, %7
+    %9:i32 = subgroupShuffleUp %val, %8
+    ret %9
   }
 }
 )";
 
-    PolyfillConfig config{.subgroup_shuffle_clamped = true};
-    Run(BuiltinPolyfill, config);
+    Run(BuiltinPolyfill, PolyfillConfig{});
 
     EXPECT_EQ(expect, str());
 }
 
 TEST_F(SpirvWriter_BuiltinPolyfillTest, SubgroupShuffleUp_SignedDelta_Clamped) {
-    auto* val = b.FunctionParam("val", ty.i32());
-    auto* delta = b.FunctionParam("delta", ty.u32());
     auto* func = b.Function("foo", ty.i32());
-    func->SetParams({val, delta});
+    func->SetStage(core::ir::Function::PipelineStage::kFragment);
+    func->SetReturnLocation(0);
 
     b.Append(func->Block(), [&] {
+        auto* val = b.Let("val", 1_i);
+        auto* delta = b.Let("delta", 1_u);
         auto* result = b.Call(ty.i32(), core::BuiltinFn::kSubgroupShuffleUp, val, delta);
         b.Return(func, result);
     });
 
     auto* src = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
+%foo = @fragment func():i32 [@location(0)] {
   $B1: {
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
     %4:i32 = subgroupShuffleUp %val, %delta
     ret %4
   }
@@ -1582,62 +1555,46 @@
     EXPECT_EQ(src, str());
 
     auto* expect = R"(
-%foo = func(%val:i32, %delta:u32):i32 {
-  $B1: {
-    %4:u32 = and %delta, 127u
-    %5:i32 = subgroupShuffleUp %val, %4
-    ret %5
+$B1: {  # root
+  %tint_subgroup_size_mask:ptr<private, u32, read_write> = var undef
+}
+
+%foo = @fragment func(%tint_subgroup_size:u32 [@subgroup_size]):i32 [@location(0)] {
+  $B2: {
+    %4:u32 = sub %tint_subgroup_size, 1u
+    store %tint_subgroup_size_mask, %4
+    %val:i32 = let 1i
+    %delta:u32 = let 1u
+    %7:u32 = load %tint_subgroup_size_mask
+    %8:u32 = and %delta, %7
+    %9:i32 = subgroupShuffleUp %val, %8
+    ret %9
   }
 }
 )";
 
-    PolyfillConfig config{.subgroup_shuffle_clamped = true};
-    Run(BuiltinPolyfill, config);
+    Run(BuiltinPolyfill, PolyfillConfig{});
 
     EXPECT_EQ(expect, str());
... (truncated)