use threadgroup pointers instead of references in metal by 39ali · Pull Request #9380 · gfx-rs/wgpu

39ali · 2026-04-06T16:29:34Z

Connections
fixes #4500

Description
threadgroups are broken in metal when referenced instead of using pointers(i suspect because of compiler reordering)
this seems to fix the issue

Testing
ran tests, and local test that uses threadgroups

Squash or Rebase?
either

Checklist

Run cargo fmt.
Run taplo format.
Run cargo clippy --tests. If applicable, add:
- --target wasm32-unknown-unknown
Run cargo xtask test to run tests.
If this contains user-facing changes, add a CHANGELOG.md entry.

ErichDonGubler · 2026-04-07T22:29:51Z

@39ali: Before we have somebody do review, this needs to be rebased, and CI errors need to be fixed.

inner-daemons

Some initial thoughts, nothing major but will need to be addressed

inner-daemons · 2026-04-12T02:35:00Z

                };
-                (coherent, space, access, "&")
+
+                let (suffix, reference) = if let crate::AddressSpace::WorkGroup = var.space {


If let is probably less clear than simply var.space == .... Also, do you know whether this should affect AddressSpace::TaskPayload, or any ray payload space? If not that warrants a comment explaining.

just saw the mesh shader changes, i'll look into it

inner-daemons · 2026-04-12T02:35:24Z

+            // but for it to work with the rest of the code we reference it in a temp var in the function body:
+            // threadgroup type& temp = *temp_ptr;
+            for (handle, var) in module.global_variables.iter() {
+                if var.space == crate::AddressSpace::WorkGroup && !fun_info[handle].is_empty() {


Again, TaskPayload?

inner-daemons · 2026-04-12T02:36:10Z

+            // for threadgroup, we use pointer and not a reference to disable compiler reordering
+            // but for it to work with the rest of the code we reference it in a temp var in the function body:
+            // threadgroup type& temp = *temp_ptr;


This should probably link to an issue. "Compiler reordering" is not perfectly clear on what the problem actually is. And its not obvious why we can redeclare as a refernece later without issue.

inner-daemons · 2026-04-12T02:37:39Z

        Ok(write!(
            out,
-            "{}{}{}{}{}{}{} {}",
+            "{}{}{}{}{}{}{} {}{}",


This does not adequately handle cases where there might be another variable with this name. You have to use the writer's namer.

inner-daemons · 2026-04-12T02:37:56Z

+
+                    writeln!(
+                        self.out,
+                        "    threadgroup {}& {} = *{}_ptr;",


This _ptr will have to change according to the comment above.

39ali · 2026-04-12T10:40:19Z

@inner-daemons i didn't like that local reference hack so i changed it to actually use the pointer

teoxoy · 2026-04-20T14:48:22Z

@39ali could you resolve the CI failures? I will take a look at the PR after.

39ali · 2026-04-20T20:32:33Z

@teoxoy done

inner-daemons · 2026-04-20T20:38:13Z

Gonna block this on my review again since it touches mesh shader stuff

teoxoy

Looks good overall, but I think we can remove 2 sections.

39ali · 2026-04-21T17:37:21Z

@teoxoy done 👍

teoxoy · 2026-04-22T15:08:01Z

Thanks!

inner-daemons · 2026-04-23T05:50:47Z

 ) {
-    uint3 nagaGridSize = _ts_main(__local_invocation_index, taskPayload, workgroupData);
+    threadgroup float workgroupData;


Why did you make this change

I'll give you some more context. There is a difference between workgroup variables declared as function parameters and those declared as variables in the function, because you have to manually allocate those declared in the function parameters (and then bind them). But there are benefits so we usually use those.

Mesh shaders don't support using this kind of workgroup variables because the grids can be dispatched by another task shader (which obviously can't allocate and bind a workgroup variable for each dispatched mesh shader grid). So mesh shaders intentionally declare them as variables in the wrapper function, but the inner function doesn't notice since it just takes a reference anyway. Task shaders don't have this limitation so they are treated like compute shaders.

Good catch! I wonder why none of the tests failed.

@teoxoy There was an edge case here to protect mesh shaders. It was expanded to task shaders in this change, not removed entirely. So it shouldn't have broken anything but it might've worsened performance or reduced the maximum size of the payload unexpectedly, I'm not entirely sure.

inner-daemons · 2026-04-23T05:52:07Z

  uint __local_invocation_index
 , object_data TaskPayload& taskPayload
-, threadgroup float& workgroupData
+, threadgroup float* workgroupData


Is this really necessary? Keep in mind this right here isn't actually declaring such a workgroup variable, its merely taking a reference to an existing variable

inner-daemons · 2026-04-23T05:56:13Z

@@ -442,7 +442,7 @@ impl TypedGlobalVariable<'_> {
        };
        let (coherent, space, access, reference) = match (var.space.to_msl_name(), var.space) {
            (Some(space), crate::AddressSpace::WorkGroup) => {


Should this also match for TaskPayload address space?

I don't think it should, it doesn't map to the threadgroup address space.

wgpu/naga/src/back/msl/writer.rs

Line 688 in 253477b

Self::TaskPayload => Some("object_data"),

Task payload storage class is defined to be basically identical to threadgroup storage class, except that in mesh shaders it is immutable. My point here is that this fix should therefore also apply to object_data variables so that they can benefit.

inner-daemons

I left some comments above, didn't realize there would be so many so I didn't put it in a review. Here, I mostly want you to avoid making changes to mesh/task shaders that aren't necessary.

Another thing that came up a ton is that TaskPayload address space functions very similarly to workgroup address space (though it is only mutable in task shaders). So it would probably be best to apply all of your fixes to that too.

inner-daemons · 2026-04-23T05:56:37Z

+        match context.function.expressions[chain] {
+            crate::Expression::GlobalVariable(handle) => {
+                let var = &context.module.global_variables[handle];
+                var.space == crate::AddressSpace::WorkGroup


Same here, should this match for task payload variables?

inner-daemons · 2026-04-23T05:57:55Z

                    _ => {
                        if var.space == crate::AddressSpace::WorkGroup
-                            && ep.stage == crate::ShaderStage::Mesh
+                            && (ep.stage == crate::ShaderStage::Mesh


This is the kind of stuff that I highlighted in the mesh shader ouptut but I don't know why it needed to change

inner-daemons · 2026-04-23T05:58:25Z

+            match *self {
+                Access::GlobalVariable(handle) => {
+                    let var = &module.global_variables[handle];
+                    var.space == crate::AddressSpace::WorkGroup


Again here, should this match a task payload pointer

inner-daemons · 2026-04-23T05:58:37Z

+        fn root_is_workgroup_pointer(&self, module: &crate::Module) -> bool {
+            if let Some(&Access::GlobalVariable(handle)) = self.stack.first() {
+                let var = &module.global_variables[handle];
+                return var.space == crate::AddressSpace::WorkGroup;


You know the deal

39ali changed the title ~~use threadgroup pointers instead of references~~ use threadgroup pointers instead of references in metal Apr 6, 2026

inner-daemons self-requested a review April 8, 2026 15:28

inner-daemons self-assigned this Apr 8, 2026

inner-daemons suggested changes Apr 12, 2026

View reviewed changes

39ali force-pushed the metal-threadgroup branch from f8258f1 to f615b3a Compare April 12, 2026 20:34

cwfitzgerald assigned teoxoy Apr 15, 2026

39ali force-pushed the metal-threadgroup branch from f615b3a to df6ba31 Compare April 20, 2026 20:16

use threadgroup pointers instead of references

a834739

39ali force-pushed the metal-threadgroup branch from df6ba31 to a834739 Compare April 20, 2026 20:16

fix clippy

229913f

inner-daemons self-requested a review April 20, 2026 20:38

teoxoy reviewed Apr 21, 2026

View reviewed changes

Comment thread naga/src/back/msl/writer.rs Outdated

Comment thread naga/src/back/msl/writer.rs Outdated

remove workgroup pointer check in function arg and structs

0ebe07d

teoxoy approved these changes Apr 22, 2026

View reviewed changes

inner-daemons reviewed Apr 23, 2026

View reviewed changes

inner-daemons suggested changes Apr 23, 2026

View reviewed changes

Conversation

39ali commented Apr 6, 2026 • edited by teoxoy Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ErichDonGubler commented Apr 7, 2026

Uh oh!

inner-daemons left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

39ali commented Apr 12, 2026

Uh oh!

teoxoy commented Apr 20, 2026

Uh oh!

39ali commented Apr 20, 2026

Uh oh!

inner-daemons commented Apr 20, 2026

Uh oh!

teoxoy left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

39ali commented Apr 21, 2026

Uh oh!

teoxoy commented Apr 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

inner-daemons Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

inner-daemons left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

39ali commented Apr 6, 2026 •

edited by teoxoy

Loading

inner-daemons left a comment •

edited

Loading

inner-daemons Apr 23, 2026 •

edited

Loading