fix(hlsl-out): use Interlocked<op> intrinsic for atomic integers (#…

…2294)

We currently assume that we are using raw `RWByteAddressBuffer` methods for all atomic operations (`<pointer>.Interlocked<op>(<raw_byte_offset>, …)`), which is only true when we use `var<storage, read_write>` globals. For `var<workgroup>` globals, we need `Interlocked<op>(<pointer>, …)`, using the original expression as the first argument.

Fix this by branching on the `pointer`'s address space in `Atomic` statements, and implementing the workgroup address space case with intrinsics.

Remove atomic ops from `access`, add new `atomicOps` test.

Fixes #2284

src/back/hlsl/writer.rs

@@ -1856,14 +1856,37 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     }
                 };
 
-                let var_handle = self.fill_access_chain(module, pointer, func_ctx)?;
-                // working around the borrow checker in `self.write_expr`
-                let chain = mem::take(&mut self.temp_access_chain);
-                let var_name = &self.names[&NameKey::GlobalVariable(var_handle)];
+                // Validation ensures that `pointer` has a `Pointer` type.
+                let pointer_space = func_ctx.info[pointer]
+                    .ty
+                    .inner_with(&module.types)
+                    .pointer_space()
+                    .unwrap();
 
                 let fun_str = fun.to_hlsl_suffix();
-                write!(self.out, " {res_name}; {var_name}.Interlocked{fun_str}(")?;
-                self.write_storage_address(module, &chain, func_ctx)?;
+                write!(self.out, " {res_name}; ")?;
+                match pointer_space {
+                    crate::AddressSpace::WorkGroup => {
+                        write!(self.out, "Interlocked{fun_str}(")?;
+                        self.write_expr(module, pointer, func_ctx)?;
+                    }
+                    crate::AddressSpace::Storage { .. } => {
+                        let var_handle = self.fill_access_chain(module, pointer, func_ctx)?;
+                        // The call to `self.write_storage_address` wants
+                        // mutable access to all of `self`, so temporarily take
+                        // ownership of our reusable access chain buffer.
+                        let chain = mem::take(&mut self.temp_access_chain);
+                        let var_name = &self.names[&NameKey::GlobalVariable(var_handle)];
+                        write!(self.out, "{var_name}.Interlocked{fun_str}(")?;
+                        self.write_storage_address(module, &chain, func_ctx)?;
+                        self.temp_access_chain = chain;
+                    }
+                    ref other => {
+                        return Err(Error::Custom(format!(
+                            "invalid address space {other:?} for atomic statement"
+                        )))
+                    }
+                }
                 write!(self.out, ", ")?;
                 // handle the special cases
                 match *fun {
@@ -1878,7 +1901,6 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
                 self.write_expr(module, value, func_ctx)?;
                 writeln!(self.out, ", {res_name});")?;
-                self.temp_access_chain = chain;
                 self.named_expressions.insert(result, res_name);
             }
             Statement::Switch {

tests/in/access.wgsl

@@ -151,23 +151,6 @@ fn foo_frag() -> @location(0) vec4<f32> {
 	return vec4<f32>(0.0);
 }
 
-@compute @workgroup_size(1)
-fn atomics() {
-	var tmp: i32;
-	let value = atomicLoad(&bar.atom);
-	tmp = atomicAdd(&bar.atom, 5);
-	tmp = atomicSub(&bar.atom, 5);
-	tmp = atomicAnd(&bar.atom, 5);
-	tmp = atomicOr(&bar.atom, 5);
-	tmp = atomicXor(&bar.atom, 5);
-	tmp = atomicMin(&bar.atom, 5);
-	tmp = atomicMax(&bar.atom, 5);
-	tmp = atomicExchange(&bar.atom, 5);
-	// https://github.com/gpuweb/gpuweb/issues/2021
-	// tmp = atomicCompareExchangeWeak(&bar.atom, 5, 5);
-	atomicStore(&bar.atom, value);
-}
-
 var<workgroup> val: u32;
 
 fn assign_through_ptr_fn(p: ptr<workgroup, u32>) {
@@ -184,4 +167,4 @@ fn assign_through_ptr() {
 
     assign_through_ptr_fn(&val);
     assign_array_through_ptr_fn(&arr);
-}
+}

tests/in/atomicOps.wgsl

@@ -0,0 +1,141 @@
+// This test covers the cross product of:
+//
+// * All atomic operations.
+// * On all applicable scopes (storage read-write, workgroup).
+// * For all shapes of modeling atomic data.
+
+struct Struct {
+    atomic_scalar: atomic<u32>,
+    atomic_arr: array<atomic<i32>, 2>,
+}
+
+@group(0) @binding(0)
+var<storage, read_write> storage_atomic_scalar: atomic<u32>;
+@group(0) @binding(1)
+var<storage, read_write> storage_atomic_arr: array<atomic<i32>, 2>;
+@group(0) @binding(2)
+var<storage, read_write> storage_struct: Struct;
+
+var<workgroup> workgroup_atomic_scalar: atomic<u32>;
+var<workgroup> workgroup_atomic_arr: array<atomic<i32>, 2>;
+var<workgroup> workgroup_struct: Struct;
+
+@compute
+@workgroup_size(2)
+fn cs_main(@builtin(local_invocation_id) id: vec3<u32>) {
+    atomicStore(&storage_atomic_scalar, 1u);
+    atomicStore(&storage_atomic_arr[1], 1i);
+    atomicStore(&storage_struct.atomic_scalar, 1u);
+    atomicStore(&storage_struct.atomic_arr[1], 1i);
+    atomicStore(&workgroup_atomic_scalar, 1u);
+    atomicStore(&workgroup_atomic_arr[1], 1i);
+    atomicStore(&workgroup_struct.atomic_scalar, 1u);
+    atomicStore(&workgroup_struct.atomic_arr[1], 1i);
+
+    workgroupBarrier();
+
+    atomicLoad(&storage_atomic_scalar);
+    atomicLoad(&storage_atomic_arr[1]);
+    atomicLoad(&storage_struct.atomic_scalar);
+    atomicLoad(&storage_struct.atomic_arr[1]);
+    atomicLoad(&workgroup_atomic_scalar);
+    atomicLoad(&workgroup_atomic_arr[1]);
+    atomicLoad(&workgroup_struct.atomic_scalar);
+    atomicLoad(&workgroup_struct.atomic_arr[1]);
+
+    workgroupBarrier();
+
+    atomicAdd(&storage_atomic_scalar, 1u);
+    atomicAdd(&storage_atomic_arr[1], 1i);
+    atomicAdd(&storage_struct.atomic_scalar, 1u);
+    atomicAdd(&storage_struct.atomic_arr[1], 1i);
+    atomicAdd(&workgroup_atomic_scalar, 1u);
+    atomicAdd(&workgroup_atomic_arr[1], 1i);
+    atomicAdd(&workgroup_struct.atomic_scalar, 1u);
+    atomicAdd(&workgroup_struct.atomic_arr[1], 1i);
+
+    workgroupBarrier();
+
+    atomicSub(&storage_atomic_scalar, 1u);
+    atomicSub(&storage_atomic_arr[1], 1i);
+    atomicSub(&storage_struct.atomic_scalar, 1u);
+    atomicSub(&storage_struct.atomic_arr[1], 1i);
+    atomicSub(&workgroup_atomic_scalar, 1u);
+    atomicSub(&workgroup_atomic_arr[1], 1i);
+    atomicSub(&workgroup_struct.atomic_scalar, 1u);
+    atomicSub(&workgroup_struct.atomic_arr[1], 1i);
+
+    workgroupBarrier();
+
+    atomicMax(&storage_atomic_scalar, 1u);
+    atomicMax(&storage_atomic_arr[1], 1i);
+    atomicMax(&storage_struct.atomic_scalar, 1u);
+    atomicMax(&storage_struct.atomic_arr[1], 1i);
+    atomicMax(&workgroup_atomic_scalar, 1u);
+    atomicMax(&workgroup_atomic_arr[1], 1i);
+    atomicMax(&workgroup_struct.atomic_scalar, 1u);
+    atomicMax(&workgroup_struct.atomic_arr[1], 1i);
+
+    workgroupBarrier();
+
+    atomicMin(&storage_atomic_scalar, 1u);
+    atomicMin(&storage_atomic_arr[1], 1i);
+    atomicMin(&storage_struct.atomic_scalar, 1u);
+    atomicMin(&storage_struct.atomic_arr[1], 1i);
+    atomicMin(&workgroup_atomic_scalar, 1u);
+    atomicMin(&workgroup_atomic_arr[1], 1i);
+    atomicMin(&workgroup_struct.atomic_scalar, 1u);
+    atomicMin(&workgroup_struct.atomic_arr[1], 1i);
+
+    workgroupBarrier();
+
+    atomicAnd(&storage_atomic_scalar, 1u);
+    atomicAnd(&storage_atomic_arr[1], 1i);
+    atomicAnd(&storage_struct.atomic_scalar, 1u);
+    atomicAnd(&storage_struct.atomic_arr[1], 1i);
+    atomicAnd(&workgroup_atomic_scalar, 1u);
+    atomicAnd(&workgroup_atomic_arr[1], 1i);
+    atomicAnd(&workgroup_struct.atomic_scalar, 1u);
+    atomicAnd(&workgroup_struct.atomic_arr[1], 1i);
+
+    workgroupBarrier();
+
+    atomicOr(&storage_atomic_scalar, 1u);
+    atomicOr(&storage_atomic_arr[1], 1i);
+    atomicOr(&storage_struct.atomic_scalar, 1u);
+    atomicOr(&storage_struct.atomic_arr[1], 1i);
+    atomicOr(&workgroup_atomic_scalar, 1u);
+    atomicOr(&workgroup_atomic_arr[1], 1i);
+    atomicOr(&workgroup_struct.atomic_scalar, 1u);
+    atomicOr(&workgroup_struct.atomic_arr[1], 1i);
+
+    workgroupBarrier();
+
+    atomicXor(&storage_atomic_scalar, 1u);
+    atomicXor(&storage_atomic_arr[1], 1i);
+    atomicXor(&storage_struct.atomic_scalar, 1u);
+    atomicXor(&storage_struct.atomic_arr[1], 1i);
+    atomicXor(&workgroup_atomic_scalar, 1u);
+    atomicXor(&workgroup_atomic_arr[1], 1i);
+    atomicXor(&workgroup_struct.atomic_scalar, 1u);
+    atomicXor(&workgroup_struct.atomic_arr[1], 1i);
+
+    atomicExchange(&storage_atomic_scalar, 1u);
+    atomicExchange(&storage_atomic_arr[1], 1i);
+    atomicExchange(&storage_struct.atomic_scalar, 1u);
+    atomicExchange(&storage_struct.atomic_arr[1], 1i);
+    atomicExchange(&workgroup_atomic_scalar, 1u);
+    atomicExchange(&workgroup_atomic_arr[1], 1i);
+    atomicExchange(&workgroup_struct.atomic_scalar, 1u);
+    atomicExchange(&workgroup_struct.atomic_arr[1], 1i);
+
+    // // TODO: https://github.com/gpuweb/gpuweb/issues/2021
+    // atomicCompareExchangeWeak(&storage_atomic_scalar, 1u);
+    // atomicCompareExchangeWeak(&storage_atomic_arr[1], 1i);
+    // atomicCompareExchangeWeak(&storage_struct.atomic_scalar, 1u);
+    // atomicCompareExchangeWeak(&storage_struct.atomic_arr[1], 1i);
+    // atomicCompareExchangeWeak(&workgroup_atomic_scalar, 1u);
+    // atomicCompareExchangeWeak(&workgroup_atomic_arr[1], 1i);
+    // atomicCompareExchangeWeak(&workgroup_struct.atomic_scalar, 1u);
+    // atomicCompareExchangeWeak(&workgroup_struct.atomic_arr[1], 1i);
+}