Implement atomic compare and exchange for msl

shards-lang · Feb 16, 2024 · 5f43672 · 5f43672
1 parent dd699aa
commit 5f43672
Showing 1 changed file with 43 additions and 4 deletions.
diff --git a/naga/src/back/msl/writer.rs b/naga/src/back/msl/writer.rs
@@ -2921,10 +2921,38 @@ impl<W: Write> Writer<W> {
                                 &context.expression,
                             )?;
                         }
-                        crate::AtomicFunction::Exchange { .. } => {
-                            return Err(Error::FeatureNotImplemented(
-                                "atomic CompareExchange".to_string(),
-                            ));
+                        crate::AtomicFunction::Exchange { compare: Some(cmp) } => {
+                            let policy = context.expression.choose_bounds_check_policy(pointer);
+
+                            let result_type_name = match context.expression.info[result].ty {
+                                TypeResolution::Handle(ty_handle) => {
+                                    let ty_name = TypeContext {
+                                        handle: ty_handle,
+                                        gctx: context.expression.module.to_ctx(),
+                                        names: &self.names,
+                                        access: crate::StorageAccess::empty(),
+                                        binding: None,
+                                        first_time: false,
+                                    };
+                                    ty_name
+                                }
+                                _ => {
+                                    return Err(Error::FeatureNotImplemented(
+                                        "atomic CompareExchange".to_string(),
+                                    ));
+                                }
+                            };
+
+                            write!(
+                                self.out,
+                                "__make_atomic_cas_result<{result_type_name}>({ATOMIC_REFERENCE}"
+                            )?;
+                            self.put_access_chain(pointer, policy, &context.expression)?;
+                            write!(self.out, ", ")?;
+                            self.put_expression(cmp, &context.expression, true)?;
+                            write!(self.out, ", ")?;
+                            self.put_expression(value, &context.expression, true)?;
+                            write!(self.out, ")")?;
                         }
                     }
                     // done
@@ -3176,6 +3204,17 @@ impl<W: Write> Writer<W> {
         }
         writeln!(self.out)?;
 
+        writeln!(
+            self.out,
+            r#"
+        template<typename R, typename T, typename A> R __make_atomic_cas_result(device A* ptr, T cmp_, T v) {{
+            T cmp = cmp_;
+            bool exchanged = {NAMESPACE}::atomic_compare_exchange_weak_explicit(ptr, &cmp, v, {NAMESPACE}::memory_order_relaxed, {NAMESPACE}::memory_order_relaxed);
+            return R{{cmp, exchanged}};
+        }}
+        "#
+        )?;
+
         {
             let mut indices = vec![];
             for (handle, var) in module.global_variables.iter() {