Fixed situations with local_invocation_id and/or local_invocation_index being written incorrectly

Author: inner-daemons

CHANGELOG.md

@@ -175,6 +175,7 @@ depth_stencil: Some(wgpu::DepthStencilState::stencil(
 - Renamed `EXPERIMENTAL_PASSTHROUGH_SHADERS` to `PASSTHROUGH_SHADERS` and made this no longer an experimental feature. by @inner-daemons in [#9054](https://github.com/gfx-rs/wgpu/pull/9054).
 - BREAKING: End offsets in trace and `player` commands are now represented using `offset` + `size` instead. By @ErichDonGubler in [9073](https://github.com/gfx-rs/wgpu/pull/9073).
 - Validate some uncaught cases where buffer transfer operations could overflow when computing an end offset. By @ErichDonGubler in [9073](https://github.com/gfx-rs/wgpu/pull/9073).
+- Fix various issues relating to `local_invocation_index` and `local_invocation_id` in HLSL and MSL. By @inner-daemons in [#9099](https://github.com/gfx-rs/wgpu/pull/9099).
 - Added internal labels to validation GPU objects and timestamp normalization code to improve clarity in graphics debuggers. By @szostid in [9094](https://github.com/gfx-rs/wgpu/pull/9094)
 - Fix multi-planar texture copying. By @noituri [#9069](https://github.com/gfx-rs/wgpu/pull/9069)
 

naga/src/back/hlsl/writer.rs

@@ -74,6 +74,7 @@ struct EntryPointBinding {
     ty_name: String,
     /// Members of generated structure
     members: Vec<EpStructMember>,
+    local_invocation_index_name: Option<String>,
 }
 
 pub(super) struct EntryPointInterface {
@@ -619,11 +620,23 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
         write!(self.out, "struct {struct_name}")?;
         writeln!(self.out, " {{")?;
+        let mut local_invocation_index_name = None;
+        let mut subgroup_id_used = false;
         for m in members.iter() {
             // Sanity check that each IO member is a built-in or is assigned a
             // location. Also see note about nesting in `write_ep_input_struct`.
             debug_assert!(m.binding.is_some());
 
+            match m.binding {
+                Some(crate::Binding::BuiltIn(crate::BuiltIn::SubgroupId)) => {
+                    subgroup_id_used = true;
+                }
+                Some(crate::Binding::BuiltIn(crate::BuiltIn::LocalInvocationIndex)) => {
+                    local_invocation_index_name = Some(m.name.clone());
+                }
+                _ => (),
+            }
+
             if is_subgroup_builtin_binding(&m.binding) {
                 continue;
             }
@@ -636,17 +649,10 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             self.write_semantic(&m.binding, Some(shader_stage))?;
             writeln!(self.out, ";")?;
         }
-        if members.iter().any(|arg| {
-            matches!(
-                arg.binding,
-                Some(crate::Binding::BuiltIn(crate::BuiltIn::SubgroupId))
-            )
-        }) {
-            writeln!(
-                self.out,
-                "{}uint __local_invocation_index : SV_GroupIndex;",
-                back::INDENT
-            )?;
+        if subgroup_id_used && local_invocation_index_name.is_none() {
+            let name = self.namer.call("local_invocation_index");
+            writeln!(self.out, "{}uint {name} : SV_GroupIndex;", back::INDENT)?;
+            local_invocation_index_name = Some(name);
         }
         writeln!(self.out, "}};")?;
         writeln!(self.out)?;
@@ -666,6 +672,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             arg_name: self.namer.call(struct_name.to_lowercase().as_str()),
             ty_name: struct_name,
             members,
+            local_invocation_index_name,
         })
     }
 
@@ -845,8 +852,10 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             Some(crate::Binding::BuiltIn(crate::BuiltIn::SubgroupId)) => {
                 write!(
                     self.out,
-                    "{}.__local_invocation_index / WaveGetLaneCount()",
-                    ep_input.arg_name
+                    "{}.{} / WaveGetLaneCount()",
+                    ep_input.arg_name,
+                    // When writing SubgroupId, we always guarantee that local_invocation_index_name is written
+                    ep_input.local_invocation_index_name.as_ref().unwrap()
                 )?;
             }
             _ => {
@@ -1587,6 +1596,8 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         let need_workgroup_variables_initialization =
             self.need_workgroup_variables_initialization(func_ctx, module);
 
+        let needs_local_invocation_id_name = need_workgroup_variables_initialization;
+        let mut local_invocation_id_name = None;
         // Write function arguments for non entry point functions
         match func_ctx.ty {
             back::FunctionType::Function(handle) => {
@@ -1614,6 +1625,12 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                         let argument_name =
                             &self.names[&NameKey::EntryPointArgument(ep_index, index as u32)];
 
+                        if arg.binding
+                            == Some(crate::Binding::BuiltIn(crate::BuiltIn::LocalInvocationId))
+                        {
+                            local_invocation_id_name = Some(argument_name.clone());
+                        }
+
                         write!(self.out, " {argument_name}")?;
                         if let TypeInner::Array { base, size, .. } = module.types[arg.ty].inner {
                             self.write_array_size(module, base, size)?;
@@ -1622,7 +1639,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                         self.write_semantic(&arg.binding, Some((stage, Io::Input)))?;
                     }
                 }
-                if need_workgroup_variables_initialization {
+                if needs_local_invocation_id_name && local_invocation_id_name.is_none() {
                     if self
                         .entry_point_io
                         .get(&(ep_index as usize))
@@ -1633,7 +1650,9 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     {
                         write!(self.out, ", ")?;
                     }
-                    write!(self.out, "uint3 __local_invocation_id : SV_GroupThreadID")?;
+                    let var_name = self.namer.call("local_invocation_id");
+                    write!(self.out, "uint3 {var_name} : SV_GroupThreadID")?;
+                    local_invocation_id_name = Some(var_name);
                 }
             }
         }
@@ -1653,7 +1672,13 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         writeln!(self.out, "{{")?;
 
         if need_workgroup_variables_initialization {
-            self.write_workgroup_variables_initialization(func_ctx, module)?;
+            self.write_workgroup_variables_initialization(
+                func_ctx,
+                module,
+                // need_workgroup_variables_initialization forces this to be written
+                // if the user doesn't specify it (so this must be Some())
+                local_invocation_id_name.unwrap(),
+            )?;
         }
 
         if let back::FunctionType::EntryPoint(index) = func_ctx.ty {
@@ -1804,12 +1829,13 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         &mut self,
         func_ctx: &back::FunctionCtx,
         module: &Module,
+        local_invocation_id_name: String,
     ) -> BackendResult {
         let level = back::Level(1);
 
         writeln!(
             self.out,
-            "{level}if (all(__local_invocation_id == uint3(0u, 0u, 0u))) {{"
+            "{level}if (all({local_invocation_id_name} == uint3(0u, 0u, 0u))) {{"
         )?;
 
         let vars = module.global_variables.iter().filter(|&(handle, var)| {

naga/src/back/msl/keywords.rs

@@ -343,6 +343,8 @@ const RESERVED: &[&str] = &[
     "M_SQRT1_2",
     // Naga utilities
     "DefaultConstructible",
+    // Naga builtin names
+    "__local_invocation_id",
     super::writer::FREXP_FUNCTION,
     super::writer::MODF_FUNCTION,
     super::writer::ABS_FUNCTION,

naga/tests/in/wgsl/8820-multiple-local-invocation-index-id.wgsl

@@ -0,0 +1,15 @@
+var<workgroup> wg_var: u32;
+
+struct Input {
+    @builtin(local_invocation_id)
+    local_invocation_id: vec3<u32>,
+    @builtin(local_invocation_index)
+    local_invocation_index: u32,
+}
+
+@compute
+@workgroup_size(1)
+fn compute1(input: Input) {
+    wg_var = input.local_invocation_index * 2;
+    wg_var += input.local_invocation_id.x;
+}

naga/tests/out/glsl/wgsl-8820-multiple-local-invocation-index-id.compute1.Compute.glsl

@@ -0,0 +1,27 @@
+#version 310 es
+
+precision highp float;
+precision highp int;
+
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+
+struct Input {
+    uvec3 local_invocation_id;
+    uint local_invocation_index;
+};
+shared uint wg_var;
+
+
+void main() {
+    if (gl_LocalInvocationID == uvec3(0u)) {
+        wg_var = 0u;
+    }
+    memoryBarrierShared();
+    barrier();
+    Input input_ = Input(gl_LocalInvocationID, gl_LocalInvocationIndex);
+    wg_var = (input_.local_invocation_index * 2u);
+    uint _e8 = wg_var;
+    wg_var = (_e8 + input_.local_invocation_id.x);
+    return;
+}
+

naga/tests/out/hlsl/spv-subgroup-operations-s.hlsl

@@ -4,7 +4,7 @@ static uint global_2 = (uint)0;
 static uint global_3 = (uint)0;
 
 struct ComputeInput_main {
-    uint __local_invocation_index : SV_GroupIndex;
+    uint local_invocation_index : SV_GroupIndex;
 };
 
 void function()
@@ -39,7 +39,7 @@ void function()
 void main(ComputeInput_main computeinput_main)
 {
     uint param = (1u + WaveGetLaneCount() - 1u) / WaveGetLaneCount();
-    uint param_1 = computeinput_main.__local_invocation_index / WaveGetLaneCount();
+    uint param_1 = computeinput_main.local_invocation_index / WaveGetLaneCount();
     uint param_2 = WaveGetLaneCount();
     uint param_3 = WaveGetLaneIndex();
     global = param;

naga/tests/out/hlsl/wgsl-8820-multiple-local-invocation-index-id.hlsl

@@ -0,0 +1,19 @@
+struct Input {
+    uint3 local_invocation_id : SV_GroupThreadID;
+    uint local_invocation_index : SV_GroupIndex;
+};
+
+groupshared uint wg_var;
+
+[numthreads(1, 1, 1)]
+void compute1_(Input input, uint3 local_invocation_id : SV_GroupThreadID)
+{
+    if (all(local_invocation_id == uint3(0u, 0u, 0u))) {
+        wg_var = (uint)0;
+    }
+    GroupMemoryBarrierWithGroupSync();
+    wg_var = (input.local_invocation_index * 2u);
+    uint _e8 = wg_var;
+    wg_var = (_e8 + input.local_invocation_id.x);
+    return;
+}

naga/tests/out/hlsl/wgsl-8820-multiple-local-invocation-index-id.ron

@@ -0,0 +1,12 @@
+(
+    vertex:[
+    ],
+    fragment:[
+    ],
+    compute:[
+        (
+            entry_point:"compute1_",
+            target_profile:"cs_5_1",
+        ),
+    ],
+)

naga/tests/out/hlsl/wgsl-atomicOps-int64.hlsl

@@ -30,9 +30,9 @@ groupshared int64_t workgroup_atomic_arr[2];
 groupshared Struct workgroup_struct;
 
 [numthreads(2, 1, 1)]
-void cs_main(uint3 id : SV_GroupThreadID, uint3 __local_invocation_id : SV_GroupThreadID)
+void cs_main(uint3 id : SV_GroupThreadID)
 {
-    if (all(__local_invocation_id == uint3(0u, 0u, 0u))) {
+    if (all(id == uint3(0u, 0u, 0u))) {
         workgroup_atomic_scalar = (uint64_t)0;
         workgroup_atomic_arr = (int64_t[2])0;
         workgroup_struct = (Struct)0;

naga/tests/out/hlsl/wgsl-atomicOps.hlsl

@@ -21,9 +21,9 @@ groupshared int workgroup_atomic_arr[2];
 groupshared Struct workgroup_struct;
 
 [numthreads(2, 1, 1)]
-void cs_main(uint3 id : SV_GroupThreadID, uint3 __local_invocation_id : SV_GroupThreadID)
+void cs_main(uint3 id : SV_GroupThreadID)
 {
-    if (all(__local_invocation_id == uint3(0u, 0u, 0u))) {
+    if (all(id == uint3(0u, 0u, 0u))) {
         workgroup_atomic_scalar = (uint)0;
         workgroup_atomic_arr = (int[2])0;
         workgroup_struct = (Struct)0;