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);

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)