From bb604fd6a3ad9cb9dbe2ff579be55d008e624b18 Mon Sep 17 00:00:00 2001
From: glalonde <lalondegeoffrey@gmail.com>
Date: Tue, 1 Feb 2022 16:37:09 -0800
Subject: [PATCH] Mitigation for MSL atomic bounds check. (#1703)

[msl-out] Correct output for bounds-checked atomic accesses.
---
 src/back/msl/writer.rs                      | 20 ++++++++-
 tests/in/bounds-check-zero-atomic.param.ron |  6 +++
 tests/in/bounds-check-zero-atomic.wgsl      | 25 +++++++++++
 tests/out/msl/bounds-check-zero-atomic.msl  | 49 +++++++++++++++++++++
 tests/snapshots.rs                          |  1 +
 5 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 tests/in/bounds-check-zero-atomic.param.ron
 create mode 100644 tests/in/bounds-check-zero-atomic.wgsl
 create mode 100644 tests/out/msl/bounds-check-zero-atomic.msl

diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs
index 90db79f2f..f2fd1bd89 100644
--- a/src/back/msl/writer.rs
+++ b/src/back/msl/writer.rs
@@ -791,15 +791,33 @@ impl<W: Write> Writer<W> {
         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
+        // the pointer operand should be unchecked.
+        let policy = context.choose_bounds_check_policy(pointer);
+        let checked = policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
+            && self.put_bounds_checks(pointer, context, back::Level(0), "")?;
+
+        // If requested and successfully put bounds checks, continue the ternary expression.
+        if checked {
+            write!(self.out, " ? ")?;
+        }
+
         write!(
             self.out,
             "{}::atomic_fetch_{}_explicit({}",
             NAMESPACE, key, ATOMIC_REFERENCE
         )?;
-        self.put_expression(pointer, context, true)?;
+        self.put_access_chain(pointer, policy, context)?;
         write!(self.out, ", ")?;
         self.put_expression(value, context, true)?;
         write!(self.out, ", {}::memory_order_relaxed)", NAMESPACE)?;
+
+        // Finish the ternary expression.
+        if checked {
+            write!(self.out, " : DefaultConstructible()")?;
+        }
+
         Ok(())
     }
 
diff --git a/tests/in/bounds-check-zero-atomic.param.ron b/tests/in/bounds-check-zero-atomic.param.ron
new file mode 100644
index 000000000..3ca0053af
--- /dev/null
+++ b/tests/in/bounds-check-zero-atomic.param.ron
@@ -0,0 +1,6 @@
+(
+	bounds_check_policies: (
+		index: ReadZeroSkipWrite,
+		buffer: ReadZeroSkipWrite,
+	),
+)
diff --git a/tests/in/bounds-check-zero-atomic.wgsl b/tests/in/bounds-check-zero-atomic.wgsl
new file mode 100644
index 000000000..ad4012dbf
--- /dev/null
+++ b/tests/in/bounds-check-zero-atomic.wgsl
@@ -0,0 +1,25 @@
+// Tests for `naga::back::BoundsCheckPolicy::ReadZeroSkipWrite` for atomic types.
+
+// These are separate from `bounds-check-zero.wgsl because SPIR-V does not yet
+// support `ReadZeroSkipWrite` for atomics. Once it does, the test files could
+// be combined.
+
+struct Globals {
+    a: atomic<u32>;
+    b: array<atomic<u32>, 10>;
+    c: array<atomic<u32>>;
+};
+
+@group(0) @binding(0) var<storage, read_write> globals: Globals;
+
+fn fetch_add_atomic() -> u32 {
+   return atomicAdd(&globals.a, 1u);
+}
+
+fn fetch_add_atomic_static_sized_array(i: i32) -> u32 {
+   return atomicAdd(&globals.b[i], 1u);
+}
+
+fn fetch_add_atomic_dynamic_sized_array(i: i32) -> u32 {
+   return atomicAdd(&globals.c[i], 1u);
+}
diff --git a/tests/out/msl/bounds-check-zero-atomic.msl b/tests/out/msl/bounds-check-zero-atomic.msl
new file mode 100644
index 000000000..d6911b1aa
--- /dev/null
+++ b/tests/out/msl/bounds-check-zero-atomic.msl
@@ -0,0 +1,49 @@
+// language: metal1.1
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+struct DefaultConstructible {
+    template<typename T>
+    operator T() && {
+        return T {};
+    }
+};
+struct _mslBufferSizes {
+    metal::uint size0;
+};
+
+struct type_1 {
+    metal::atomic_uint inner[10];
+};
+typedef metal::atomic_uint type_2[1];
+struct Globals {
+    metal::atomic_uint a;
+    type_1 b;
+    type_2 c;
+};
+
+metal::uint fetch_add_atomic(
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    metal::uint _e3 = metal::atomic_fetch_add_explicit(&globals.a, 1u, metal::memory_order_relaxed);
+    return _e3;
+}
+
+metal::uint fetch_add_atomic_static_sized_array(
+    int i,
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    metal::uint _e5 = metal::uint(i) < 10 ? metal::atomic_fetch_add_explicit(&globals.b.inner[i], 1u, metal::memory_order_relaxed) : DefaultConstructible();
+    return _e5;
+}
+
+metal::uint fetch_add_atomic_dynamic_sized_array(
+    int i_1,
+    device Globals& globals,
+    constant _mslBufferSizes& _buffer_sizes
+) {
+    metal::uint _e5 = metal::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;
+}
diff --git a/tests/snapshots.rs b/tests/snapshots.rs
index 56378e2ac..65439888f 100644
--- a/tests/snapshots.rs
+++ b/tests/snapshots.rs
@@ -477,6 +477,7 @@ fn convert_wgsl() {
             Targets::SPIRV | Targets::METAL | Targets::GLSL | Targets::HLSL | Targets::WGSL,
         ),
         ("bounds-check-zero", Targets::SPIRV | Targets::METAL),
+        ("bounds-check-zero-atomic", Targets::METAL),
         ("bounds-check-restrict", Targets::SPIRV | Targets::METAL),
         ("bounds-check-image-restrict", Targets::SPIRV),
         ("bounds-check-image-rzsw", Targets::SPIRV),