From d1d7dff8a6e1246b06e5acce077fb0cc2b8d5f11 Mon Sep 17 00:00:00 2001
From: Bjorn <75190918+BjornTheProgrammer@users.noreply.github.com>
Date: Wed, 7 Aug 2024 17:27:38 -0700
Subject: [PATCH] Made compatible with no_std and no atomic ptr

---
 Cargo.toml                         |  2 ++
 crates/cubecl-common/Cargo.toml    |  7 +++++++
 crates/cubecl-common/src/reader.rs | 15 ++++++++++++++-
 crates/cubecl-runtime/Cargo.toml   |  7 ++++++-
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c9c39e572..92b28754e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -68,5 +68,7 @@ derive_more = { version = "0.99.18", features = ["display"], default-features =
 env_logger = "0.11.3"
 strum = {version = "0.26.3", features = ["derive"]}
 
+portable-atomic-util = { version = "0.2.2", features = ["alloc"] } # alloc is for no_std
+
 [profile.dev]
 opt-level = 2
diff --git a/crates/cubecl-common/Cargo.toml b/crates/cubecl-common/Cargo.toml
index 30302b243..73530d095 100644
--- a/crates/cubecl-common/Cargo.toml
+++ b/crates/cubecl-common/Cargo.toml
@@ -26,5 +26,12 @@ serde = { workspace = true }
 rand = { workspace = true }
 pollster = { workspace = true, optional = true }
 
+[target.'cfg(target_has_atomic = "ptr")'.dependencies]
+spin = { workspace = true, features = ["mutex", "spin_mutex"] }
+
+[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies]
+portable-atomic-util = { workspace = true }
+spin = { workspace = true, features = ["mutex", "spin_mutex", "portable_atomic"] }
+
 [dev-dependencies]
 dashmap = { workspace = true }
diff --git a/crates/cubecl-common/src/reader.rs b/crates/cubecl-common/src/reader.rs
index d459a6e9c..3b32359d8 100644
--- a/crates/cubecl-common/src/reader.rs
+++ b/crates/cubecl-common/src/reader.rs
@@ -1,4 +1,10 @@
-use alloc::{boxed::Box, sync::Arc, task::Wake, vec::Vec};
+#[cfg(target_has_atomic = "ptr")]
+use alloc::{sync::Arc, task::Wake};
+
+#[cfg(not(target_has_atomic = "ptr"))]
+use portable_atomic_util::{task::Wake, Arc};
+
+use alloc::{boxed::Box, vec::Vec};
 use core::{
     future::Future,
     pin::Pin,
@@ -16,8 +22,15 @@ pub fn reader_from_concrete(val: Vec<u8>) -> Reader {
 struct DummyWaker;
 
 impl Wake for DummyWaker {
+    #[cfg(target_has_atomic = "ptr")]
     fn wake(self: Arc<Self>) {}
+    #[cfg(target_has_atomic = "ptr")]
     fn wake_by_ref(self: &Arc<Self>) {}
+
+    #[cfg(not(target_has_atomic = "ptr"))]
+    fn wake(_this: Arc<Self>) {}
+    #[cfg(not(target_has_atomic = "ptr"))]
+    fn wake_by_ref(_this: &Arc<Self>) {}
 }
 
 /// Read a future synchronously.
diff --git a/crates/cubecl-runtime/Cargo.toml b/crates/cubecl-runtime/Cargo.toml
index 1a7f8374e..d82739c2e 100644
--- a/crates/cubecl-runtime/Cargo.toml
+++ b/crates/cubecl-runtime/Cargo.toml
@@ -30,7 +30,6 @@ autotune-persistent-cache = ["dirs", "md5", "serde", "serde_json"] # Assume std
 [dependencies]
 cubecl-common = { path = "../cubecl-common", version = "0.1.1", default-features = false }
 derive-new = { workspace = true }
-spin = { workspace = true }
 log = { workspace = true }
 hashbrown = { workspace = true }
 dirs = { workspace = true, optional = true }
@@ -40,6 +39,12 @@ md5 = { workspace = true, optional = true }
 pollster = { workspace = true, optional = true }
 async-channel = { workspace = true, optional = true }
 
+[target.'cfg(target_has_atomic = "ptr")'.dependencies]
+spin = { workspace = true, features = ["mutex", "spin_mutex"] }
+
+[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies]
+spin = { workspace = true, features = ["mutex", "spin_mutex", "portable_atomic"] }
+
 [target.'cfg(target_family = "wasm")'.dependencies]
 web-time = { workspace = true }