wip

Author: BjornTheProgrammer

Cargo.lock

@@ -143,15 +143,12 @@ nvml-wrapper = "0.10.0"
 sysinfo = "0.30.13"
 systemstat = "0.2.3"
 
-portable-atomic = {version = "1.3.3", features = ["critical-section"]}
-portable-atomic-util = {version = "0.1.3", features = ["alloc"]} # alloc is for no_std
-
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "dd274b6515f853ab1c82db410d6fc7697060f69e" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "dd274b6515f853ab1c82db410d6fc7697060f69e" }
+# cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "dd274b6515f853ab1c82db410d6fc7697060f69e" }
+# cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "dd274b6515f853ab1c82db410d6fc7697060f69e" }
 ### For local development. ###
-# cubecl = { path = "../cubecl/crates/cubecl" }
-# cubecl-common = { path = "../cubecl/crates/cubecl-common" }
+cubecl = { path = "../cubecl/crates/cubecl", default-features = false, features=["portable-atomic-critical-section"] }
+cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }
 
 [profile.dev]
 debug = 0     # Speed up compilation time and not necessary.

Cargo.toml

@@ -120,8 +120,6 @@ rand = { workspace = true, features = ["std_rng"] } # Default enables std
 
 # Using in place of use std::sync::Mutex when std is disabled
 spin = { workspace = true, features = ["mutex", "spin_mutex", "portable_atomic"] }
-portable-atomic = {version = "1.3.3"}
-portable-atomic-util = {version = "0.1.3", features = ["alloc"]}
 
 # The same implementation of HashMap in std but with no_std support (only alloc crate is needed)
 hashbrown = { workspace = true, features = ["serde"] } # no_std compatible
@@ -138,6 +136,10 @@ thiserror = { workspace = true, optional = true }
 regex = { workspace = true, optional = true }
 num-traits = { workspace = true }
 
+[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies]
+portable-atomic = { version = "1.6.0" }
+portable-atomic-util = { version = "0.2.0", features = [ "alloc" ] }
+
 [dev-dependencies]
 tempfile = { workspace = true }
 burn-dataset = { path = "../burn-dataset", version = "0.14.0", features = [

crates/burn-core/Cargo.toml

@@ -6,10 +6,10 @@ use crate::module::{
 
 use alloc::string::ToString;
 
-#[cfg(not(use_portable_atomics))]
+#[cfg(target_has_atomic = "ptr")]
 use alloc::sync::Arc;
 
-#[cfg(use_portable_atomics)]
+#[cfg(not(target_has_atomic = "ptr"))]
 use portable_atomic_util::Arc;
 
 use alloc::vec::Vec;