From ddcd5d3121150b2b1beee6e54e9125ff31aaa9a2 Mon Sep 17 00:00:00 2001
From: Raph Levien <raph.levien@gmail.com>
Date: Mon, 24 Oct 2022 09:20:57 -0700
Subject: [PATCH] Fix incorrect atomic bounds check on metal back-end (#2099)

* Fix incorrect atomic bounds check on metal back-end

Generalize put_atomic_fetch to handle `exchange` as well, rather than special-cased code which didn't do the bounds check (the check handling as fixed in #1703 but only for the fetch cases, exchange was skipped).

Fixes #1848

* Add tests for atomic exchange
---
 src/back/msl/writer.rs                     | 29 ++++++++++++++--------
 tests/in/bounds-check-zero-atomic.wgsl     | 13 ++++++++++
 tests/out/msl/bounds-check-zero-atomic.msl | 26 +++++++++++++++++++
 3 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs
index 9a59289bef..a832c348ab 100644
--- a/src/back/msl/writer.rs
+++ b/src/back/msl/writer.rs
@@ -1187,6 +1187,17 @@ impl<W: Write> Writer<W> {
         key: &str,
         value: Handle<crate::Expression>,
         context: &ExpressionContext,
+    ) -> BackendResult {
+        self.put_atomic_operation(pointer, "fetch_", key, value, context)
+    }
+
+    fn put_atomic_operation(
+        &mut self,
+        pointer: Handle<crate::Expression>,
+        key1: &str,
+        key2: &str,
+        value: Handle<crate::Expression>,
+        context: &ExpressionContext,
     ) -> BackendResult {
         // If the pointer we're passing to the atomic operation needs to be conditional
         // for `ReadZeroSkipWrite`, the condition needs to *surround* the atomic op, and
@@ -1202,8 +1213,8 @@ impl<W: Write> Writer<W> {
 
         write!(
             self.out,
-            "{}::atomic_fetch_{}_explicit({}",
-            NAMESPACE, key, ATOMIC_REFERENCE
+            "{}::atomic_{}{}_explicit({}",
+            NAMESPACE, key1, key2, ATOMIC_REFERENCE
         )?;
         self.put_access_chain(pointer, policy, context)?;
         write!(self.out, ", ")?;
@@ -2725,15 +2736,13 @@ impl<W: Write> Writer<W> {
                             self.put_atomic_fetch(pointer, "max", value, &context.expression)?;
                         }
                         crate::AtomicFunction::Exchange { compare: None } => {
-                            write!(
-                                self.out,
-                                "{}::atomic_exchange_explicit({}",
-                                NAMESPACE, ATOMIC_REFERENCE,
+                            self.put_atomic_operation(
+                                pointer,
+                                "exchange",
+                                "",
+                                value,
+                                &context.expression,
                             )?;
-                            self.put_expression(pointer, &context.expression, true)?;
-                            write!(self.out, ", ")?;
-                            self.put_expression(value, &context.expression, true)?;
-                            write!(self.out, ", {}::memory_order_relaxed)", NAMESPACE)?;
                         }
                         crate::AtomicFunction::Exchange { .. } => {
                             return Err(Error::FeatureNotImplemented(
diff --git a/tests/in/bounds-check-zero-atomic.wgsl b/tests/in/bounds-check-zero-atomic.wgsl
index 3cdbea14a7..004f08a0a5 100644
--- a/tests/in/bounds-check-zero-atomic.wgsl
+++ b/tests/in/bounds-check-zero-atomic.wgsl
@@ -23,3 +23,16 @@ fn fetch_add_atomic_static_sized_array(i: i32) -> u32 {
 fn fetch_add_atomic_dynamic_sized_array(i: i32) -> u32 {
    return atomicAdd(&globals.c[i], 1u);
 }
+
+fn exchange_atomic() -> u32 {
+   return atomicExchange(&globals.a, 1u);
+}
+
+fn exchange_atomic_static_sized_array(i: i32) -> u32 {
+   return atomicExchange(&globals.b[i], 1u);
+}
+
+fn exchange_atomic_dynamic_sized_array(i: i32) -> u32 {
+   return atomicExchange(&globals.c[i], 1u);
+}
+
diff --git a/tests/out/msl/bounds-check-zero-atomic.msl b/tests/out/msl/bounds-check-zero-atomic.msl
index c76ba1b549..95028ee796 100644
--- a/tests/out/msl/bounds-check-zero-atomic.msl
+++ b/tests/out/msl/bounds-check-zero-atomic.msl
@@ -49,3 +49,29 @@ uint fetch_add_atomic_dynamic_sized_array(
     uint _e5 = uint(i_1) < 1 + (_buffer_sizes.size0 - 44 - 4) / 4 ? metal::atomic_fetch_add_explicit(&globals.c[i_1], 1u, metal::memory_order_relaxed) : DefaultConstructible();
     return _e5;
 }
+
+uint exchange_atomic(
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    uint _e3 = metal::atomic_exchange_explicit(&globals.a, 1u, metal::memory_order_relaxed);
+    return _e3;
+}
+
+uint exchange_atomic_static_sized_array(
+    int i_2,
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    uint _e5 = uint(i_2) < 10 ? metal::atomic_exchange_explicit(&globals.b.inner[i_2], 1u, metal::memory_order_relaxed) : DefaultConstructible();
+    return _e5;
+}
+
+uint exchange_atomic_dynamic_sized_array(
+    int i_3,
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    uint _e5 = uint(i_3) < 1 + (_buffer_sizes.size0 - 44 - 4) / 4 ? metal::atomic_exchange_explicit(&globals.c[i_3], 1u, metal::memory_order_relaxed) : DefaultConstructible();
+    return _e5;
+}