Fix atomics in workgroup storage

src/back/hlsl/writer.rs

@@ -682,7 +682,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
     /// Adds no trailing or leading whitespace
     pub(super) fn write_value_type(&mut self, module: &Module, inner: &TypeInner) -> BackendResult {
         match *inner {
-            TypeInner::Scalar { kind, width } => {
+            TypeInner::Scalar { kind, width } | TypeInner::Atomic { kind, width } => {
                 write!(self.out, "{}", kind.to_hlsl_str(width)?)?;
             }
             TypeInner::Vector { size, kind, width } => {

src/front/wgsl/mod.rs

@@ -3284,12 +3284,14 @@ impl Parser {
             }
             "atomicStore" => {
                 let _ = lexer.next();
+                emitter.start(context.expressions);
                 lexer.open_arguments()?;
                 let mut expression_ctx = context.as_expression(block, &mut emitter);
                 let pointer = self.parse_atomic_pointer(lexer, expression_ctx.reborrow())?;
                 lexer.expect(Token::Separator(','))?;
                 let value = self.parse_general_expression(lexer, expression_ctx)?;
                 lexer.close_arguments()?;
+                block.extend(emitter.finish(context.expressions));
                 Some(crate::Statement::Store { pointer, value })
             }
             "textureStore" => {

tests/in/globals.wgsl

@@ -3,8 +3,10 @@
 let Foo: bool = true;
 
 var<workgroup> wg : array<f32, 10u>;
+var<workgroup> at: atomic<u32>;
 
 [[stage(compute), workgroup_size(1)]]
 fn main() {
 	wg[3] = 1.0;
+	atomicStore(&at, 2u);
 }

tests/out/glsl/globals.main.Compute.glsl

@@ -7,9 +7,12 @@ layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
 
 shared float wg[10];
 
+shared uint at;
+
 
 void main() {
     wg[3] = 1.0;
+    at = 2u;
     return;
 }
 

tests/out/hlsl/globals.hlsl

@@ -1,10 +1,12 @@
 static const bool Foo = true;
 
 groupshared float wg[10];
+groupshared uint at;
 
 [numthreads(1, 1, 1)]
 void main()
 {
     wg[3] = 1.0;
+    at = 2u;
     return;
 }

tests/out/msl/globals.msl

@@ -9,7 +9,9 @@ struct type2 {
 
 kernel void main1(
   threadgroup type2& wg
+, threadgroup metal::atomic_uint& at
 ) {
     wg.inner[3] = 1.0;
+    metal::atomic_store_explicit(&at, 2u, metal::memory_order_relaxed);
     return;
 }

tests/out/spv/globals.spvasm

@@ -1,13 +1,13 @@
 ; SPIR-V
 ; Version: 1.0
 ; Generator: rspirv
-; Bound: 21
+; Bound: 26
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
-OpEntryPoint GLCompute %15 "main"
-OpExecutionMode %15 LocalSize 1 1 1
-OpDecorate %11 ArrayStride 4
+OpEntryPoint GLCompute %18 "main"
+OpExecutionMode %18 LocalSize 1 1 1
+OpDecorate %12 ArrayStride 4
 %2 = OpTypeVoid
 %4 = OpTypeBool
 %3 = OpConstantTrue  %4
@@ -17,17 +17,23 @@ OpDecorate %11 ArrayStride 4
 %7 = OpConstant  %8  3
 %10 = OpTypeFloat 32
 %9 = OpConstant  %10  1.0
-%11 = OpTypeArray %10 %5
-%13 = OpTypePointer Workgroup %11
-%12 = OpVariable  %13  Workgroup
-%16 = OpTypeFunction %2
-%18 = OpTypePointer Workgroup %10
-%19 = OpConstant  %6  3
-%15 = OpFunction  %2  None %16
-%14 = OpLabel
-OpBranch %17
+%11 = OpConstant  %6  2
+%12 = OpTypeArray %10 %5
+%14 = OpTypePointer Workgroup %12
+%13 = OpVariable  %14  Workgroup
+%16 = OpTypePointer Workgroup %6
+%15 = OpVariable  %16  Workgroup
+%19 = OpTypeFunction %2
+%21 = OpTypePointer Workgroup %10
+%22 = OpConstant  %6  3
+%24 = OpConstant  %8  2
+%25 = OpConstant  %6  256
+%18 = OpFunction  %2  None %19
 %17 = OpLabel
-%20 = OpAccessChain  %18  %12 %19
-OpStore %20 %9
+OpBranch %20
+%20 = OpLabel
+%23 = OpAccessChain  %21  %13 %22
+OpStore %23 %9
+OpAtomicStore %15 %24 %25 %11
 OpReturn
 OpFunctionEnd

tests/out/wgsl/globals.wgsl

@@ -1,9 +1,11 @@
 let Foo: bool = true;
 
 var<workgroup> wg: array<f32,10u>;
+var<workgroup> at: atomic<u32>;
 
 [[stage(compute), workgroup_size(1, 1, 1)]]
 fn main() {
     wg[3] = 1.0;
+    atomicStore(&at, 2u);
     return;
 }