Support training in no-std (#2830)

* burn-core: support no-std optim

* burn-autodiff: support no-std

* Update xtask/src/commands/build.rs

---------

Co-authored-by: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>

Cargo.toml

@@ -154,6 +154,7 @@ torch-sys = "0.19.0"  # matches what tch is using, required for lib detection
 
 ahash = { version = "0.8.11", default-features = false }
 portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
+portable-atomic = { version = "1.10.0" }
 
 ### For the main burn branch. ###
 cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "b41cbd82d53f091e76f56cad58c277fe2481c48e" }

crates/burn-autodiff/Cargo.toml

@@ -25,6 +25,9 @@ burn-tensor-testgen = { path = "../burn-tensor-testgen", version = "0.17.0", opt
 derive-new = { workspace = true }
 spin = { workspace = true }
 log = { workspace = true }
+hashbrown = { workspace = true }
+num-traits = { workspace = true }
+portable-atomic = { workspace = true }
 
 [dev-dependencies]
 burn-tensor = { path = "../burn-tensor", version = "0.17.0", default-features = false, features = [

crates/burn-autodiff/src/backend.rs

@@ -4,6 +4,7 @@ use crate::{
     runtime::AutodiffClient,
     tensor::AutodiffTensor,
 };
+use alloc::{format, string::String};
 use burn_tensor::{
     backend::{AutodiffBackend, Backend},
     ops::{BoolTensor, IntTensor, QuantizedTensor},

crates/burn-autodiff/src/checkpoint/base.rs

@@ -2,8 +2,10 @@ use super::{
     retro_forward::RetroForwards,
     state::{BackwardStates, State},
 };
+use crate::collections::HashMap;
 use crate::graph::NodeID;
-use std::collections::HashMap;
+
+use alloc::{vec, vec::Vec};
 
 #[derive(new, Debug)]
 /// Links a [NodeID] to its autodiff graph [NodeRef]

crates/burn-autodiff/src/checkpoint/builder.rs

@@ -1,9 +1,11 @@
 use crate::{
+    collections::HashMap,
     graph::{ComputingProperty, NodeID, NodeSteps},
     tensor::AutodiffTensor,
 };
+use alloc::{boxed::Box, sync::Arc, vec::Vec};
 use burn_tensor::backend::Backend;
-use std::{any::Any, collections::HashMap, sync::Arc};
+use core::any::Any;
 
 use super::{
     base::{Checkpointer, NodeTree},

crates/burn-autodiff/src/checkpoint/retro_forward.rs

@@ -1,6 +1,8 @@
+use crate::collections::HashMap;
 use crate::graph::NodeID;
 
-use std::{collections::HashMap, fmt::Debug, sync::Arc};
+use alloc::sync::Arc;
+use core::fmt::Debug;
 
 use super::state::{BackwardStates, State};
 

crates/burn-autodiff/src/checkpoint/state.rs

@@ -1,6 +1,8 @@
-use std::{any::Any, collections::HashMap};
+use core::any::Any;
 
+use crate::collections::HashMap;
 use crate::graph::NodeID;
+use alloc::boxed::Box;
 
 /// In order to accept arbitrary node output in the same hashmap, we need to upcast them to any.
 pub(crate) type StateContent = Box<dyn Any + Send>;

crates/burn-autodiff/src/checkpoint/strategy.rs

@@ -1,9 +1,9 @@
 use core::fmt::Debug;
-use std::sync::Arc;
 
 use burn_tensor::backend::Backend;
 
 use crate::{graph::ComputingProperty, tensor::AutodiffTensor};
+use alloc::sync::Arc;
 
 use super::{
     builder::{ActionType, CheckpointerBuilder},

crates/burn-autodiff/src/graph/base.rs

@@ -1,9 +1,9 @@
 use super::NodeID;
-use crate::{checkpoint::base::Checkpointer, grads::Gradients};
-use std::collections::HashMap;
+use crate::{checkpoint::base::Checkpointer, collections::HashMap, grads::Gradients};
+use alloc::{boxed::Box, vec::Vec};
 
 /// Backward step for reverse mode autodiff.
-pub trait Step: Send + std::fmt::Debug {
+pub trait Step: Send + core::fmt::Debug {
     /// Executes the step and consumes it.
     fn step(self: Box<Self>, grads: &mut Gradients, checkpointer: &mut Checkpointer);
     /// Depth of the operation relative to the first node added to a graph.

crates/burn-autodiff/src/graph/node.rs

@@ -1,5 +1,9 @@
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::Arc;
+use alloc::{sync::Arc, vec::Vec};
+
+#[cfg(target_has_atomic = "64")]
+use core::sync::atomic::{AtomicU64, Ordering};
+#[cfg(not(target_has_atomic = "64"))]
+use portable_atomic::{AtomicU64, Ordering};
 
 use crate::checkpoint::retro_forward::RetroForward;
 use crate::runtime::AutodiffClientImpl;

crates/burn-autodiff/src/graph/traversal.rs

@@ -1,6 +1,9 @@
 use super::{Step, StepBoxed};
-use crate::NodeID;
-use std::collections::{HashMap, HashSet};
+use crate::{
+    collections::{HashMap, HashSet},
+    NodeID,
+};
+use alloc::vec::Vec;
 
 /// Breadth for search algorithm.
 pub struct BreadthFirstSearch;

crates/burn-autodiff/src/lib.rs

@@ -1,3 +1,4 @@
+#![cfg_attr(not(feature = "std"), no_std)]
 #![warn(missing_docs)]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 
@@ -34,3 +35,12 @@ pub use backend::*;
 
 #[cfg(feature = "export_tests")]
 mod tests;
+
+/// A facade around for HashMap and HashSet.
+/// This avoids elaborate import wrangling having to happen in every module.
+mod collections {
+    #[cfg(not(feature = "std"))]
+    pub use hashbrown::{HashMap, HashSet};
+    #[cfg(feature = "std")]
+    pub use std::collections::{HashMap, HashSet};
+}

crates/burn-autodiff/src/ops/activation.rs

@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use core::marker::PhantomData;
 
 use crate::{
     checkpoint::{

crates/burn-autodiff/src/ops/backward.rs

@@ -14,13 +14,13 @@ use burn_tensor::backend::Backend;
 /// Concrete types implementing this trait should not have any state.
 /// If a state is necessary during the backward pass,
 /// they should be declared with the associated type 'State'.
-pub trait Backward<B, const N: usize>: Send + std::fmt::Debug
+pub trait Backward<B, const N: usize>: Send + core::fmt::Debug
 where
     Self: Sized + 'static,
     B: Backend,
 {
     /// Associated type to compute the backward pass.
-    type State: Clone + Send + std::fmt::Debug + 'static;
+    type State: Clone + Send + core::fmt::Debug + 'static;
 
     /// The backward pass.
     fn backward(

crates/burn-autodiff/src/ops/base.rs

@@ -10,8 +10,9 @@ use crate::{
     graph::{ComputingProperty, NodeID, NodeRef, Requirement, Step},
     tensor::AutodiffTensor,
 };
+use alloc::{boxed::Box, vec::Vec};
 use burn_tensor::{backend::Backend, ops::FloatTensor, Shape, TensorMetadata};
-use std::marker::PhantomData;
+use core::marker::PhantomData;
 
 /// Operation in preparation.
 ///
@@ -134,7 +135,7 @@ where
 impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, ComputePropertyDone>
 where
     B: Backend,
-    S: Clone + Send + std::fmt::Debug + 'static,
+    S: Clone + Send + core::fmt::Debug + 'static,
     BO: Backward<B, N, State = S>,
 {
     /// Prepare an operation that requires a state during the backward pass.
@@ -161,7 +162,7 @@ where
 impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, UnTracked>
 where
     B: Backend,
-    S: Clone + Send + std::fmt::Debug + 'static,
+    S: Clone + Send + core::fmt::Debug + 'static,
     BO: Backward<B, N, State = S>,
 {
     /// Finish the preparation of an untracked operation and returns the output tensor.
@@ -184,7 +185,7 @@ where
 impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, Tracked>
 where
     B: Backend,
-    S: Clone + Send + std::fmt::Debug + 'static,
+    S: Clone + Send + core::fmt::Debug + 'static,
     BO: Backward<B, N, State = S>,
 {
     /// Finish the preparation of a tracked operation and returns the output tensor.
@@ -235,7 +236,7 @@ struct OpsStep<B, T, SB, const N: usize>
 where
     B: Backend,
     T: Backward<B, N, State = SB>,
-    SB: Clone + Send + std::fmt::Debug + 'static,
+    SB: Clone + Send + core::fmt::Debug + 'static,
 {
     ops: Ops<SB, N>,
     backward: T,
@@ -246,7 +247,7 @@ impl<B, T, SB, const N: usize> Step for OpsStep<B, T, SB, N>
 where
     B: Backend,
     T: Backward<B, N, State = SB>,
-    SB: Clone + Send + std::fmt::Debug + 'static,
+    SB: Clone + Send + core::fmt::Debug + 'static,
 {
     fn step(self: Box<Self>, grads: &mut Gradients, checkpointer: &mut Checkpointer) {
         self.backward.backward(self.ops, grads, checkpointer);

crates/burn-autodiff/src/ops/bool_tensor.rs

@@ -1,4 +1,5 @@
 use crate::{checkpoint::strategy::CheckpointStrategy, tensor::AutodiffTensor, Autodiff};
+use alloc::vec::Vec;
 
 use burn_tensor::{
     backend::Backend,
@@ -31,7 +32,7 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
         B::bool_reshape(tensor, shape)
     }
 
-    fn bool_slice(tensor: BoolTensor<B>, ranges: &[std::ops::Range<usize>]) -> BoolTensor<B> {
+    fn bool_slice(tensor: BoolTensor<B>, ranges: &[core::ops::Range<usize>]) -> BoolTensor<B> {
         B::bool_slice(tensor, ranges)
     }
 
@@ -41,7 +42,7 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
 
     fn bool_slice_assign(
         tensor: BoolTensor<Self>,
-        ranges: &[std::ops::Range<usize>],
+        ranges: &[core::ops::Range<usize>],
         value: BoolTensor<Self>,
     ) -> BoolTensor<Self> {
         B::bool_slice_assign(tensor, ranges, value)

crates/burn-autodiff/src/ops/int_tensor.rs

@@ -1,4 +1,5 @@
 use crate::{checkpoint::strategy::CheckpointStrategy, tensor::AutodiffTensor, Autodiff};
+use alloc::vec::Vec;
 
 use burn_tensor::{
     backend::Backend,
@@ -27,7 +28,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_reshape(tensor, shape)
     }
 
-    fn int_slice(tensor: IntTensor<B>, ranges: &[std::ops::Range<usize>]) -> IntTensor<B> {
+    fn int_slice(tensor: IntTensor<B>, ranges: &[core::ops::Range<usize>]) -> IntTensor<B> {
         B::int_slice(tensor, ranges)
     }
 
@@ -37,7 +38,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
 
     fn int_slice_assign(
         tensor: IntTensor<B>,
-        ranges: &[std::ops::Range<usize>],
+        ranges: &[core::ops::Range<usize>],
         value: IntTensor<B>,
     ) -> IntTensor<B> {
         B::int_slice_assign(tensor, ranges, value)
@@ -305,7 +306,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_random(shape, distribution, device)
     }
 
-    fn int_arange(range: std::ops::Range<i64>, device: &Device<Self>) -> IntTensor<Self> {
+    fn int_arange(range: core::ops::Range<i64>, device: &Device<Self>) -> IntTensor<Self> {
         B::int_arange(range, device)
     }
 

crates/burn-autodiff/src/ops/qtensor.rs

@@ -1,4 +1,4 @@
-use std::ops::Range;
+use core::ops::Range;
 
 use burn_tensor::{
     backend::Backend,

crates/burn-autodiff/src/ops/tensor.rs

@@ -1,5 +1,9 @@
-use alloc::{vec, vec::Vec};
-use std::marker::PhantomData;
+use alloc::{boxed::Box, vec, vec::Vec};
+use core::marker::PhantomData;
+
+#[cfg(not(feature = "std"))]
+#[allow(unused_imports, reason = "required on aarch64, unused on x86_64")]
+use num_traits::float::Float;
 
 use crate::{
     checkpoint::{
@@ -1127,15 +1131,15 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
 
     fn float_slice(
         tensor: FloatTensor<Self>,
-        ranges: &[std::ops::Range<usize>],
+        ranges: &[core::ops::Range<usize>],
     ) -> FloatTensor<Self> {
         #[derive(Debug)]
         struct Index;
 
         #[derive(new, Debug)]
         struct RetroSlice<B: Backend> {
             tensor_id: NodeID,
-            ranges: Vec<std::ops::Range<usize>>,
+            ranges: Vec<core::ops::Range<usize>>,
             _backend: PhantomData<B>,
         }
 
@@ -1148,7 +1152,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         }
 
         impl<B: Backend> Backward<B, 1> for Index {
-            type State = (Vec<std::ops::Range<usize>>, Shape, B::Device);
+            type State = (Vec<core::ops::Range<usize>>, Shape, B::Device);
 
             fn backward(
                 self,
@@ -1186,7 +1190,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
 
     fn float_slice_assign(
         tensor: FloatTensor<Self>,
-        ranges: &[std::ops::Range<usize>],
+        ranges: &[core::ops::Range<usize>],
         value: FloatTensor<Self>,
     ) -> FloatTensor<Self> {
         #[derive(Debug)]
@@ -1195,7 +1199,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         #[derive(new, Debug)]
         struct RetroSliceAssign<B: Backend> {
             tensor_id: NodeID,
-            ranges: Vec<std::ops::Range<usize>>,
+            ranges: Vec<core::ops::Range<usize>>,
             value_id: NodeID,
             _backend: PhantomData<B>,
         }
@@ -1210,7 +1214,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         }
 
         impl<B: Backend> Backward<B, 2> for SliceAssign {
-            type State = (Vec<std::ops::Range<usize>>, Shape, B::Device);
+            type State = (Vec<core::ops::Range<usize>>, Shape, B::Device);
 
             fn backward(
                 self,
@@ -2066,7 +2070,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
                     let ops = checkpointer.retrieve_node_output(ops.state);
                     let exponent = B::float_neg(B::float_powf_scalar(ops, 2.0));
                     let numerator = B::float_mul_scalar(B::float_exp(exponent), 2.0.elem());
-                    let denominator = std::f64::consts::PI.sqrt().elem();
+                    let denominator = core::f64::consts::PI.sqrt().elem();
                     let value = B::float_div_scalar(numerator, denominator);
 
                     B::float_mul(grad, value)

crates/burn-autodiff/src/ops/transaction.rs

@@ -8,7 +8,7 @@ use crate::{checkpoint::strategy::CheckpointStrategy, Autodiff};
 impl<B: Backend, C: CheckpointStrategy> TransactionOps<Self> for Autodiff<B, C> {
     fn tr_execute(
         transaction: TransactionPrimitive<Self>,
-    ) -> impl std::future::Future<Output = burn_tensor::ops::TransactionPrimitiveResult> + 'static + Send
+    ) -> impl core::future::Future<Output = burn_tensor::ops::TransactionPrimitiveResult> + 'static + Send
     {
         B::tr_execute(TransactionPrimitive {
             read_floats: transaction

crates/burn-autodiff/src/runtime/memory_management.rs

@@ -1,9 +1,10 @@
-use crate::{tensor::NodeRefCount, NodeID};
-use std::{
+use crate::{
     collections::{HashMap, HashSet},
-    mem,
-    sync::Arc,
+    tensor::NodeRefCount,
+    NodeID,
 };
+use alloc::{borrow::ToOwned, sync::Arc, vec, vec::Vec};
+use core::mem;
 
 #[derive(Default, Debug)]
 pub struct GraphMemoryManagement {
@@ -82,7 +83,10 @@ impl GraphMemoryManagement {
 
     fn clear_unused_roots(&mut self, to_delete: &mut Vec<NodeID>) {
         for (id, parents) in self.nodes.iter() {
-            let is_useful = matches!(self.statuses.get(id), Some(NodeMemoryStatus::Useful));
+            let is_useful = matches!(
+                self.statuses.get(id.as_ref()),
+                Some(NodeMemoryStatus::Useful)
+            );
 
             // Check if parents are either empty or absent from self.nodes
             let parents_absent = parents.iter().all(|p| !self.nodes.contains_key(p));

crates/burn-autodiff/src/runtime/mutex.rs

@@ -11,7 +11,7 @@ use burn_tensor::backend::Backend;
 pub struct MutexClient;
 
 impl core::fmt::Debug for MutexClient {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.write_str("MutexClient")
     }
 }