Merge pull request #2203 from Try/msl-atomics-fix

MSL: fix extraction of global variables, in case of atomics

3 files changed, 160 insertions, 0 deletions

diff --git a/reference/shaders-msl-no-opt/comp/extract-atomics-from-function.comp b/reference/shaders-msl-no-opt/comp/extract-atomics-from-function.comp
new file mode 100644
index 00000000..0d082e2f
--- /dev/null
+++ b/reference/shaders-msl-no-opt/comp/extract-atomics-from-function.comp
@@ -0,0 +1,82 @@
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+#pragma clang diagnostic ignored "-Wunused-variable"
+
+#include <metal_stdlib>
+#include <simd/simd.h>
+#include <metal_atomic>
+
+using namespace metal;
+
+constant uint3 gl_WorkGroupSize [[maybe_unused]] = uint3(64u, 1u, 1u);
+
+static inline __attribute__((always_inline))
+void testAdd(threadgroup uint& var)
+{
+    uint _29 = atomic_fetch_add_explicit((threadgroup atomic_uint*)&var, 1u, memory_order_relaxed);
+}
+
+static inline __attribute__((always_inline))
+void testMin(threadgroup uint& var)
+{
+    uint _31 = atomic_fetch_min_explicit((threadgroup atomic_uint*)&var, 2u, memory_order_relaxed);
+}
+
+static inline __attribute__((always_inline))
+void testMax(threadgroup uint& var)
+{
+    uint _33 = atomic_fetch_max_explicit((threadgroup atomic_uint*)&var, 3u, memory_order_relaxed);
+}
+
+static inline __attribute__((always_inline))
+void testOr(threadgroup uint& var)
+{
+    uint _35 = atomic_fetch_or_explicit((threadgroup atomic_uint*)&var, 5u, memory_order_relaxed);
+}
+
+static inline __attribute__((always_inline))
+void testXor(threadgroup uint& var)
+{
+    uint _37 = atomic_fetch_xor_explicit((threadgroup atomic_uint*)&var, 6u, memory_order_relaxed);
+}
+
+static inline __attribute__((always_inline))
+void testExchange(threadgroup uint& var)
+{
+    uint _39 = atomic_exchange_explicit((threadgroup atomic_uint*)&var, 7u, memory_order_relaxed);
+}
+
+static inline __attribute__((always_inline))
+void testCompSwap(threadgroup uint& var)
+{
+    uint _42;
+    do
+    {
+        _42 = 8u;
+    } while (!atomic_compare_exchange_weak_explicit((threadgroup atomic_uint*)&var, &_42, 9u, memory_order_relaxed, memory_order_relaxed) && _42 == 8u);
+}
+
+static inline __attribute__((always_inline))
+void testStore(threadgroup uint& var)
+{
+    atomic_store_explicit((threadgroup atomic_uint*)&var, 10u, memory_order_relaxed);
+}
+
+static inline __attribute__((always_inline))
+void foo(threadgroup uint& var)
+{
+    testAdd(var);
+    testMin(var);
+    testMax(var);
+    testOr(var);
+    testXor(var);
+    testExchange(var);
+    testCompSwap(var);
+    testStore(var);
+}
+
+kernel void main0()
+{
+    threadgroup uint var;
+    foo(var);
+}
+
diff --git a/shaders-msl-no-opt/comp/extract-atomics-from-function.comp b/shaders-msl-no-opt/comp/extract-atomics-from-function.comp
new file mode 100644
index 00000000..c1b7632f
--- /dev/null
+++ b/shaders-msl-no-opt/comp/extract-atomics-from-function.comp
@@ -0,0 +1,69 @@
+#version 460
+
+#extension GL_KHR_memory_scope_semantics : enable
+
+layout(local_size_x = 64) in;
+
+shared uint var;
+
+void testAdd()
+{
+  atomicAdd(var, 1);
+}
+
+void testMin()
+{
+  atomicMin(var, 2);
+}
+
+void testMax()
+{
+  atomicMax(var, 3);
+}
+
+void testAnd()
+{
+  atomicAnd(var, 4);
+}
+
+void testOr()
+{
+  atomicOr(var, 5);
+}
+
+void testXor()
+{
+  atomicXor(var, 6);
+}
+
+void testExchange()
+{
+  atomicExchange(var, 7);
+}
+
+void testCompSwap()
+{
+  atomicCompSwap(var, 8, 9);
+}
+
+void testStore()
+{
+  atomicStore(var, 10u, gl_ScopeDevice, gl_StorageSemanticsShared, gl_SemanticsRelaxed);
+}
+
+void foo()
+{
+  testAdd();
+  testMin();
+  testMax();
+  testOr();
+  testXor();
+  testExchange();
+  testCompSwap();
+  testStore();
+}
+
+void main()
+{
+  foo();
+}
diff --git a/spirv_msl.cpp b/spirv_msl.cpp
index e80f0ef3..9f9fcfc3 100644
--- a/spirv_msl.cpp
+++ b/spirv_msl.cpp
@@ -1898,9 +1898,18 @@ void CompilerMSL::extract_global_variables_from_function(uint32_t func_id, std::
 			case OpAtomicOr:
 			case OpAtomicXor:
 			case OpImageWrite:
+			{
 				if (needs_frag_discard_checks())
 					added_arg_ids.insert(builtin_helper_invocation_id);
+				uint32_t ptr = 0;
+				if (op == OpAtomicStore || op == OpImageWrite)
+					ptr = ops[0];
+				else
+					ptr = ops[2];
+				if (global_var_ids.find(ptr) != global_var_ids.end())
+					added_arg_ids.insert(ptr);
 				break;
+			}
 
 			// Emulate texture2D atomic operations
 			case OpImageTexelPointer: