[naga wgsl-in] Properly convert arguments to atomic operations. (#7573)

Author: jimblandy

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

@@ -2504,15 +2504,16 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         }
                         "atomicLoad" => {
                             let mut args = ctx.prepare_args(arguments, 1, span);
-                            let pointer = self.atomic_pointer(args.next()?, ctx)?;
+                            let (pointer, _scalar) = self.atomic_pointer(args.next()?, ctx)?;
                             args.finish()?;
 
                             ir::Expression::Load { pointer }
                         }
                         "atomicStore" => {
                             let mut args = ctx.prepare_args(arguments, 2, span);
-                            let pointer = self.atomic_pointer(args.next()?, ctx)?;
-                            let value = self.expression(args.next()?, ctx)?;
+                            let (pointer, scalar) = self.atomic_pointer(args.next()?, ctx)?;
+                            let value =
+                                self.expression_with_leaf_scalar(args.next()?, scalar, ctx)?;
                             args.finish()?;
 
                             let rctx = ctx.runtime_expression_ctx(span)?;
@@ -2526,13 +2527,14 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         "atomicCompareExchangeWeak" => {
                             let mut args = ctx.prepare_args(arguments, 3, span);
 
-                            let pointer = self.atomic_pointer(args.next()?, ctx)?;
+                            let (pointer, scalar) = self.atomic_pointer(args.next()?, ctx)?;
 
-                            let compare = self.expression(args.next()?, ctx)?;
+                            let compare =
+                                self.expression_with_leaf_scalar(args.next()?, scalar, ctx)?;
 
                             let value = args.next()?;
                             let value_span = ctx.ast_expressions.get_span(value);
-                            let value = self.expression(value, ctx)?;
+                            let value = self.expression_with_leaf_scalar(value, scalar, ctx)?;
 
                             args.finish()?;
 
@@ -3200,13 +3202,13 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         &mut self,
         expr: Handle<ast::Expression<'source>>,
         ctx: &mut ExpressionContext<'source, '_, '_>,
-    ) -> Result<'source, Handle<ir::Expression>> {
+    ) -> Result<'source, (Handle<ir::Expression>, ir::Scalar)> {
         let span = ctx.ast_expressions.get_span(expr);
         let pointer = self.expression(expr, ctx)?;
 
         match *resolve_inner!(ctx, pointer) {
             ir::TypeInner::Pointer { base, .. } => match ctx.module.types[base].inner {
-                ir::TypeInner::Atomic { .. } => Ok(pointer),
+                ir::TypeInner::Atomic(scalar) => Ok((pointer, scalar)),
                 ref other => {
                     log::error!("Pointer type to {:?} passed to atomic op", other);
                     Err(Box::new(Error::InvalidAtomicPointer(span)))
@@ -3229,8 +3231,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
     ) -> Result<'source, Option<Handle<ir::Expression>>> {
         let mut args = ctx.prepare_args(args, 2, span);
 
-        let pointer = self.atomic_pointer(args.next()?, ctx)?;
-        let value = self.expression(args.next()?, ctx)?;
+        let (pointer, scalar) = self.atomic_pointer(args.next()?, ctx)?;
+        let value = self.expression_with_leaf_scalar(args.next()?, scalar, ctx)?;
         let value_inner = resolve_inner!(ctx, value);
         args.finish()?;
 

naga/tests/in/wgsl/abstract-types-atomic.toml

@@ -0,0 +1 @@
+targets = "WGSL"

naga/tests/in/wgsl/abstract-types-atomic.wgsl

@@ -0,0 +1,34 @@
+@group(0) @binding(0)
+var<storage,read_write> atomic_i32: atomic<i32>;
+@group(0) @binding(1)
+var<storage,read_write> atomic_u32: atomic<u32>;
+
+fn test_atomic_i32() {
+  atomicStore(&atomic_i32, 1);
+  _ = atomicCompareExchangeWeak(&atomic_i32, 1, 1i);
+  _ = atomicCompareExchangeWeak(&atomic_i32, 1i, 1);
+
+  _ = atomicAdd(&atomic_i32, 1);
+  _ = atomicSub(&atomic_i32, 1);
+  _ = atomicAnd(&atomic_i32, 1);
+  _ = atomicXor(&atomic_i32, 1);
+  _ = atomicOr(&atomic_i32, 1);
+  _ = atomicMin(&atomic_i32, 1);
+  _ = atomicMax(&atomic_i32, 1);
+  _ = atomicExchange(&atomic_i32, 1);
+}
+
+fn test_atomic_u32() {
+  atomicStore(&atomic_u32, 1);
+  _ = atomicCompareExchangeWeak(&atomic_u32, 1, 1u);
+  _ = atomicCompareExchangeWeak(&atomic_u32, 1u, 1);
+
+  _ = atomicAdd(&atomic_u32, 1);
+  _ = atomicSub(&atomic_u32, 1);
+  _ = atomicAnd(&atomic_u32, 1);
+  _ = atomicXor(&atomic_u32, 1);
+  _ = atomicOr(&atomic_u32, 1);
+  _ = atomicMin(&atomic_u32, 1);
+  _ = atomicMax(&atomic_u32, 1);
+  _ = atomicExchange(&atomic_u32, 1);
+}

naga/tests/out/wgsl/wgsl-abstract-types-atomic.wgsl

@@ -0,0 +1,35 @@
+@group(0) @binding(0) 
+var<storage, read_write> atomic_i32_: atomic<i32>;
+@group(0) @binding(1) 
+var<storage, read_write> atomic_u32_: atomic<u32>;
+
+fn test_atomic_i32_() {
+    atomicStore((&atomic_i32_), 1i);
+    let _e5 = atomicCompareExchangeWeak((&atomic_i32_), 1i, 1i);
+    let _e9 = atomicCompareExchangeWeak((&atomic_i32_), 1i, 1i);
+    let _e12 = atomicAdd((&atomic_i32_), 1i);
+    let _e15 = atomicSub((&atomic_i32_), 1i);
+    let _e18 = atomicAnd((&atomic_i32_), 1i);
+    let _e21 = atomicXor((&atomic_i32_), 1i);
+    let _e24 = atomicOr((&atomic_i32_), 1i);
+    let _e27 = atomicMin((&atomic_i32_), 1i);
+    let _e30 = atomicMax((&atomic_i32_), 1i);
+    let _e33 = atomicExchange((&atomic_i32_), 1i);
+    return;
+}
+
+fn test_atomic_u32_() {
+    atomicStore((&atomic_u32_), 1u);
+    let _e5 = atomicCompareExchangeWeak((&atomic_u32_), 1u, 1u);
+    let _e9 = atomicCompareExchangeWeak((&atomic_u32_), 1u, 1u);
+    let _e12 = atomicAdd((&atomic_u32_), 1u);
+    let _e15 = atomicSub((&atomic_u32_), 1u);
+    let _e18 = atomicAnd((&atomic_u32_), 1u);
+    let _e21 = atomicXor((&atomic_u32_), 1u);
+    let _e24 = atomicOr((&atomic_u32_), 1u);
+    let _e27 = atomicMin((&atomic_u32_), 1u);
+    let _e30 = atomicMax((&atomic_u32_), 1u);
+    let _e33 = atomicExchange((&atomic_u32_), 1u);
+    return;
+}
+