Fix workgroupUniformLoad returning an atomic type

Author: cryvosh

naga/src/back/spv/block.rs

@@ -3710,43 +3710,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

@@ -3059,7 +3059,32 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                                 ir::TypeInner::Pointer {
                                     base,
                                     space: ir::AddressSpace::WorkGroup,
-                                } => base,
+                                } => 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,
+                                ),
                                 ref other => {
                                     log::error!("Type {other:?} passed to workgroupUniformLoad");
                                     let span = ctx.ast_expressions.get_span(expr);

naga/src/valid/function.rs

@@ -1479,7 +1479,24 @@ impl super::Validator {
                         base: ty,
                         space: AddressSpace::WorkGroup,
                     };
-                    if !expected_pointer_inner.non_struct_equivalent(pointer_inner, context.types) {
+                    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,30 @@
+// Test workgroupUniformLoad specialization for atomic<T> -> T
+// Issue: https://github.com/gfx-rs/wgpu/issues/8785
+
+var<workgroup> wg_scalar: atomic<u32>;
+var<workgroup> wg_signed: atomic<i32>;
+
+@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 atomic
+    atomicOr(&wg_scalar, u32(active_tile_index >= 64));
+    atomicAdd(&wg_signed, 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);
+    
+    // Should be able to use the result in comparisons
+    if scalar_val == 0u && signed_val > 0i {
+        return;
+    }
+}

naga/tests/naga/wgsl_errors.rs

@@ -47,6 +47,37 @@ fn check_success(input: &str) {
     }
 }
 
+#[test]
+fn workgroup_uniform_load_atomic_returns_scalar() {
+    let input = r#"
+var<workgroup> wg_scratch: atomic<u32>;
+
+@compute @workgroup_size(4, 4, 4)
+fn interval_tile_main(
+    @builtin(workgroup_id) workgroup_id: vec3u,
+    @builtin(local_invocation_id) local_id: vec3u
+) {
+    let active_tile_index = workgroup_id.x + workgroup_id.y * 32768u;
+    atomicOr(&wg_scratch, u32(active_tile_index >= 64u));
+    workgroupBarrier();
+    if workgroupUniformLoad(&wg_scratch) == 0 {
+        return;
+    }
+}
+"#;
+
+    let module = naga::front::wgsl::parse_str(input).unwrap_or_else(|err| {
+        panic!(
+            "expected success, but parsing failed with:\n{}",
+            err.emit_to_string(input)
+        )
+    });
+
+    naga::valid::Validator::new(valid::ValidationFlags::default(), Capabilities::all())
+        .validate(&module)
+        .unwrap();
+}
+
 #[test]
 fn very_negative_integers() {
     // wgpu#4492

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

@@ -0,0 +1,50 @@
+#version 310 es
+
+precision highp float;
+precision highp int;
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+shared uint wg_scalar;
+
+shared int wg_signed;
+
+
+void main() {
+    if (gl_LocalInvocationID == uvec3(0u)) {
+        wg_scalar = 0u;
+        wg_signed = 0;
+    }
+    memoryBarrierShared();
+    barrier();
+    uvec3 workgroup_id = gl_WorkGroupID;
+    uvec3 local_id = gl_LocalInvocationID;
+    bool local = 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);
+    memoryBarrierShared();
+    barrier();
+    memoryBarrierShared();
+    barrier();
+    uint _e16 = wg_scalar;
+    memoryBarrierShared();
+    barrier();
+    memoryBarrierShared();
+    barrier();
+    int _e18 = wg_signed;
+    memoryBarrierShared();
+    barrier();
+    if ((_e16 == 0u)) {
+        local = (_e18 > 0);
+    } else {
+        local = false;
+    }
+    bool _e26 = local;
+    if (_e26) {
+        return;
+    } else {
+        return;
+    }
+}
+

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

@@ -0,0 +1,35 @@
+groupshared uint wg_scalar;
+groupshared int wg_signed;
+
+[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;
+    }
+    GroupMemoryBarrierWithGroupSync();
+    bool local = (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);
+    GroupMemoryBarrierWithGroupSync();
+    GroupMemoryBarrierWithGroupSync();
+    uint _e16 = wg_scalar;
+    GroupMemoryBarrierWithGroupSync();
+    GroupMemoryBarrierWithGroupSync();
+    int _e18 = wg_signed;
+    GroupMemoryBarrierWithGroupSync();
+    if ((_e16 == 0u)) {
+        local = (_e18 > int(0));
+    } else {
+        local = false;
+    }
+    bool _e26 = local;
+    if (_e26) {
+        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",
+        ),
+    ],
+)

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

@@ -0,0 +1,43 @@
+// language: metal1.0
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using metal::uint;
+
+
+struct test_atomic_workgroup_uniform_loadInput {
+};
+kernel void test_atomic_workgroup_uniform_load(
+  metal::uint3 workgroup_id [[threadgroup_position_in_grid]]
+, metal::uint3 local_id [[thread_position_in_threadgroup]]
+, threadgroup metal::atomic_uint& wg_scalar
+, threadgroup metal::atomic_int& wg_signed
+) {
+    if (metal::all(local_id == metal::uint3(0u))) {
+        metal::atomic_store_explicit(&wg_scalar, 0, metal::memory_order_relaxed);
+        metal::atomic_store_explicit(&wg_signed, 0, metal::memory_order_relaxed);
+    }
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    bool local = {};
+    uint active_tile_index = workgroup_id.x + (workgroup_id.y * 32768u);
+    uint _e11 = metal::atomic_fetch_or_explicit(&wg_scalar, static_cast<uint>(active_tile_index >= 64u), metal::memory_order_relaxed);
+    int _e14 = metal::atomic_fetch_add_explicit(&wg_signed, 1, metal::memory_order_relaxed);
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    uint unnamed = metal::atomic_load_explicit(&wg_scalar, metal::memory_order_relaxed);
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    int unnamed_1 = metal::atomic_load_explicit(&wg_signed, metal::memory_order_relaxed);
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    if (unnamed == 0u) {
+        local = unnamed_1 > 0;
+    } else {
+        local = false;
+    }
+    bool _e26 = local;
+    if (_e26) {
+        return;
+    } else {
+        return;
+    }
+}