Fix workgroupUniformLoad returning atomic types (gfx-rs#8791)

Author: cryvosh

CHANGELOG.md

@@ -68,6 +68,7 @@ Bottom level categories:
 
 - Reject zero-value construction of a runtime-sized array with a validation error. Previously it would crash in the HLSL backend. By @mooori in [#8741](https://github.com/gfx-rs/wgpu/pull/8741).
 - Reject splat vector construction if the argument type does not match the type of the vector's scalar. Previously it would succeed. By @mooori in [#8829](https://github.com/gfx-rs/wgpu/pull/8829).
+- Fixed `workgroupUniformLoad` incorrectly returning an atomic when called on an atomic, it now returns the inner `T` as per the spec. By @cryvosh in [#8791](https://github.com/gfx-rs/wgpu/pull/8791).
 
 ### Documentation
 

cts_runner/test.lst

@@ -220,6 +220,9 @@ webgpu:shader,execution,expression,call,builtin,textureSampleBaseClampToEdge:2d_
 // NOTE: This is supposed to be an exhaustive listing underneath
 // `webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:*`, so exceptions can be
 // worked around.
+webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="atomic%3Cu32%3E";*
+webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="atomic%3Ci32%3E";*
+webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="AtomicInStruct";*
 webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="bool";*
 webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="u32";*
 webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="vec4u";*
@@ -228,10 +231,6 @@ webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type=
 webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="SimpleStruct";*
 //FAIL: https://github.com/gfx-rs/wgpu/issues/8812
 // webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="ComplexStruct";*
-//FAIL: https://github.com/gfx-rs/wgpu/pull/8791
-// webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="atomic%3Cu32%3E";*
-// webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="atomic%3Ci32%3E";*
-// webgpu:shader,execution,expression,call,builtin,workgroupUniformLoad:types:type="AtomicInStruct";*
 webgpu:shader,execution,flow_control,return:*
 // Many other vertex_buffer_access subtests also passing, but there are too many to enumerate.
 // Fails on Metal in CI only, not when running locally.

naga/src/back/spv/block.rs

@@ -4122,43 +4122,15 @@ impl BlockContext<'_> {
                     self.writer
                         .write_control_barrier(crate::Barrier::WORK_GROUP, &mut block.body);
                     let result_type_id = self.get_expression_type_id(&self.fun_info[result].ty);
-                    // Embed the body of
-                    match self.write_access_chain(
+                    // Match `Expression::Load` behavior, including `OpAtomicLoad` when
+                    // loading from a pointer to `atomic<T>`.
+                    let id = self.write_checked_load(
                         pointer,
                         &mut block,
                         AccessTypeAdjustment::None,
-                    )? {
-                        ExpressionPointer::Ready { pointer_id } => {
-                            let id = self.gen_id();
-                            block.body.push(Instruction::load(
-                                result_type_id,
-                                id,
-                                pointer_id,
-                                None,
-                            ));
-                            self.cached[result] = id;
-                        }
-                        ExpressionPointer::Conditional { condition, access } => {
-                            self.cached[result] = self.write_conditional_indexed_load(
-                                result_type_id,
-                                condition,
-                                &mut block,
-                                move |id_gen, block| {
-                                    // The in-bounds path. Perform the access and the load.
-                                    let pointer_id = access.result_id.unwrap();
-                                    let value_id = id_gen.next();
-                                    block.body.push(access);
-                                    block.body.push(Instruction::load(
-                                        result_type_id,
-                                        value_id,
-                                        pointer_id,
-                                        None,
-                                    ));
-                                    value_id
-                                },
-                            )
-                        }
-                    }
+                        result_type_id,
+                    )?;
+                    self.cached[result] = id;
                     self.writer
                         .write_control_barrier(crate::Barrier::WORK_GROUP, &mut block.body);
                 }

naga/src/front/wgsl/lower/mod.rs

@@ -3041,9 +3041,33 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                                 ir::TypeInner::Pointer {
                                     base,
                                     space: ir::AddressSpace::WorkGroup,
-                                } => base,
-                                ref other => {
-                                    log::error!("Type {other:?} passed to workgroupUniformLoad");
+                                } => match ctx.module.types[base].inner {
+                                    // Match `Expression::Load` semantics:
+                                    // loading through a pointer to `atomic<T>` produces a `T`.
+                                    ir::TypeInner::Atomic(scalar) => ctx.module.types.insert(
+                                        ir::Type {
+                                            name: None,
+                                            inner: ir::TypeInner::Scalar(scalar),
+                                        },
+                                        span,
+                                    ),
+                                    _ => base,
+                                },
+                                ir::TypeInner::ValuePointer {
+                                    size,
+                                    scalar,
+                                    space: ir::AddressSpace::WorkGroup,
+                                } => ctx.module.types.insert(
+                                    ir::Type {
+                                        name: None,
+                                        inner: match size {
+                                            Some(size) => ir::TypeInner::Vector { size, scalar },
+                                            None => ir::TypeInner::Scalar(scalar),
+                                        },
+                                    },
+                                    span,
+                                ),
+                                _ => {
                                     let span = ctx.ast_expressions.get_span(expr);
                                     return Err(Box::new(Error::InvalidWorkGroupUniformLoad(span)));
                                 }

naga/src/valid/function.rs

@@ -1479,7 +1479,23 @@ impl super::Validator {
                         base: ty,
                         space: AddressSpace::WorkGroup,
                     };
-                    if !expected_pointer_inner.non_struct_equivalent(pointer_inner, context.types) {
+                    // workgroupUniformLoad on atomic<T> returns T, not atomic<T>.
+                    // Verify the pointer's atomic scalar matches the result scalar.
+                    let atomic_specialization_ok = match *pointer_inner {
+                        Ti::Pointer {
+                            base: pointer_base,
+                            space: AddressSpace::WorkGroup,
+                        } => match (&context.types[pointer_base].inner, &context.types[ty].inner) {
+                            (&Ti::Atomic(pointer_scalar), &Ti::Scalar(result_scalar)) => {
+                                pointer_scalar == result_scalar
+                            }
+                            _ => false,
+                        },
+                        _ => false,
+                    };
+                    if !expected_pointer_inner.non_struct_equivalent(pointer_inner, context.types)
+                        && !atomic_specialization_ok
+                    {
                         return Err(FunctionError::WorkgroupUniformLoadInvalidPointer(pointer)
                             .with_span_static(span, "WorkGroupUniformLoad"));
                     }

naga/tests/in/wgsl/workgroup-uniform-load-atomic.wgsl

@@ -0,0 +1,43 @@
+// Test workgroupUniformLoad specialization for atomic<T> -> T
+
+struct AtomicStruct {
+    atomic_scalar: atomic<u32>,
+    atomic_arr: array<atomic<i32>, 2>,
+}
+
+var<workgroup> wg_scalar: atomic<u32>;
+var<workgroup> wg_signed: atomic<i32>;
+var<workgroup> wg_struct: AtomicStruct;
+
+@compute @workgroup_size(64)
+fn test_atomic_workgroup_uniform_load(
+    @builtin(workgroup_id) workgroup_id: vec3u,
+    @builtin(local_invocation_id) local_id: vec3u
+) {
+    let active_tile_index = workgroup_id.x + workgroup_id.y * 32768;
+    
+    // Each thread may set the atomics
+    atomicOr(&wg_scalar, u32(active_tile_index >= 64));
+    atomicAdd(&wg_signed, 1i);
+    atomicStore(&wg_struct.atomic_scalar, 1u);
+    atomicAdd(&wg_struct.atomic_arr[0], 1i);
+    
+    workgroupBarrier();
+    
+    // workgroupUniformLoad on atomic<u32> should return u32
+    let scalar_val: u32 = workgroupUniformLoad(&wg_scalar);
+    
+    // workgroupUniformLoad on atomic<i32> should return i32
+    let signed_val: i32 = workgroupUniformLoad(&wg_signed);
+    
+    // workgroupUniformLoad on struct.atomic_scalar should return u32
+    let struct_scalar: u32 = workgroupUniformLoad(&wg_struct.atomic_scalar);
+    
+    // workgroupUniformLoad on struct.atomic_arr[i] should return i32
+    let struct_arr_val: i32 = workgroupUniformLoad(&wg_struct.atomic_arr[0]);
+    
+    // Should be able to use all results in comparisons
+    if scalar_val == 0u && signed_val > 0i && struct_scalar > 0u && struct_arr_val > 0i {
+        return;
+    }
+}

naga/tests/out/glsl/wgsl-workgroup-uniform-load-atomic.test_atomic_workgroup_uniform_load.Compute.glsl

@@ -0,0 +1,83 @@
+#version 310 es
+
+precision highp float;
+precision highp int;
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+struct AtomicStruct {
+    uint atomic_scalar;
+    int atomic_arr[2];
+};
+shared uint wg_scalar;
+
+shared int wg_signed;
+
+shared AtomicStruct wg_struct;
+
+
+void main() {
+    if (gl_LocalInvocationID == uvec3(0u)) {
+        wg_scalar = 0u;
+        wg_signed = 0;
+        wg_struct = AtomicStruct(0u, int[2](0, 0));
+    }
+    memoryBarrierShared();
+    barrier();
+    uvec3 workgroup_id = gl_WorkGroupID;
+    uvec3 local_id = gl_LocalInvocationID;
+    bool local = false;
+    bool local_1 = false;
+    bool local_2 = false;
+    uint active_tile_index = (workgroup_id.x + (workgroup_id.y * 32768u));
+    uint _e11 = atomicOr(wg_scalar, uint((active_tile_index >= 64u)));
+    int _e14 = atomicAdd(wg_signed, 1);
+    wg_struct.atomic_scalar = 1u;
+    int _e22 = atomicAdd(wg_struct.atomic_arr[0], 1);
+    memoryBarrierShared();
+    barrier();
+    memoryBarrierShared();
+    barrier();
+    uint _e24 = wg_scalar;
+    memoryBarrierShared();
+    barrier();
+    memoryBarrierShared();
+    barrier();
+    int _e26 = wg_signed;
+    memoryBarrierShared();
+    barrier();
+    memoryBarrierShared();
+    barrier();
+    uint _e29 = wg_struct.atomic_scalar;
+    memoryBarrierShared();
+    barrier();
+    memoryBarrierShared();
+    barrier();
+    int _e33 = wg_struct.atomic_arr[0];
+    memoryBarrierShared();
+    barrier();
+    if ((_e24 == 0u)) {
+        local = (_e26 > 0);
+    } else {
+        local = false;
+    }
+    bool _e41 = local;
+    if (_e41) {
+        local_1 = (_e29 > 0u);
+    } else {
+        local_1 = false;
+    }
+    bool _e47 = local_1;
+    if (_e47) {
+        local_2 = (_e33 > 0);
+    } else {
+        local_2 = false;
+    }
+    bool _e53 = local_2;
+    if (_e53) {
+        return;
+    } else {
+        return;
+    }
+}
+

naga/tests/out/hlsl/wgsl-workgroup-uniform-load-atomic.hlsl

@@ -0,0 +1,64 @@
+struct AtomicStruct {
+    uint atomic_scalar;
+    int atomic_arr[2];
+};
+
+groupshared uint wg_scalar;
+groupshared int wg_signed;
+groupshared AtomicStruct wg_struct;
+
+[numthreads(64, 1, 1)]
+void test_atomic_workgroup_uniform_load(uint3 workgroup_id : SV_GroupID, uint3 local_id : SV_GroupThreadID, uint3 __local_invocation_id : SV_GroupThreadID)
+{
+    if (all(__local_invocation_id == uint3(0u, 0u, 0u))) {
+        wg_scalar = (uint)0;
+        wg_signed = (int)0;
+        wg_struct = (AtomicStruct)0;
+    }
+    GroupMemoryBarrierWithGroupSync();
+    bool local = (bool)0;
+    bool local_1 = (bool)0;
+    bool local_2 = (bool)0;
+
+    uint active_tile_index = (workgroup_id.x + (workgroup_id.y * 32768u));
+    uint _e11; InterlockedOr(wg_scalar, uint((active_tile_index >= 64u)), _e11);
+    int _e14; InterlockedAdd(wg_signed, int(1), _e14);
+    wg_struct.atomic_scalar = 1u;
+    int _e22; InterlockedAdd(wg_struct.atomic_arr[0], int(1), _e22);
+    GroupMemoryBarrierWithGroupSync();
+    GroupMemoryBarrierWithGroupSync();
+    uint _e24 = wg_scalar;
+    GroupMemoryBarrierWithGroupSync();
+    GroupMemoryBarrierWithGroupSync();
+    int _e26 = wg_signed;
+    GroupMemoryBarrierWithGroupSync();
+    GroupMemoryBarrierWithGroupSync();
+    uint _e29 = wg_struct.atomic_scalar;
+    GroupMemoryBarrierWithGroupSync();
+    GroupMemoryBarrierWithGroupSync();
+    int _e33 = wg_struct.atomic_arr[0];
+    GroupMemoryBarrierWithGroupSync();
+    if ((_e24 == 0u)) {
+        local = (_e26 > int(0));
+    } else {
+        local = false;
+    }
+    bool _e41 = local;
+    if (_e41) {
+        local_1 = (_e29 > 0u);
+    } else {
+        local_1 = false;
+    }
+    bool _e47 = local_1;
+    if (_e47) {
+        local_2 = (_e33 > int(0));
+    } else {
+        local_2 = false;
+    }
+    bool _e53 = local_2;
+    if (_e53) {
+        return;
+    } else {
+        return;
+    }
+}

naga/tests/out/hlsl/wgsl-workgroup-uniform-load-atomic.ron

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