Add a sync feature to common, core, and tensor (#893)

Author: AlexErrant

burn-common/Cargo.toml

@@ -15,6 +15,7 @@ default = ["std"]
 
 std = ["rand/std"]
 
+wasm-sync = []
 
 [target.'cfg(target_family = "wasm")'.dependencies]
 async-trait = { workspace = true }

burn-common/src/reader.rs

@@ -1,7 +1,7 @@
 use alloc::boxed::Box;
 use core::marker::PhantomData;
 
-#[cfg(target_family = "wasm")]
+#[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
 #[async_trait::async_trait]
 /// Allows to create async reader.
 pub trait AsyncReader<T>: Send {
@@ -15,10 +15,10 @@ pub enum Reader<T> {
     Concrete(T),
     /// Sync data variant.
     Sync(Box<dyn SyncReader<T>>),
-    #[cfg(target_family = "wasm")]
+    #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
     /// Async data variant.
     Async(Box<dyn AsyncReader<T>>),
-    #[cfg(target_family = "wasm")]
+    #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
     /// Future data variant.
     Future(core::pin::Pin<Box<dyn core::future::Future<Output = T> + Send>>),
 }
@@ -52,7 +52,7 @@ where
     }
 }
 
-#[cfg(target_family = "wasm")]
+#[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
 #[async_trait::async_trait]
 impl<I, O, F> AsyncReader<O> for MappedReader<I, O, F>
 where
@@ -67,7 +67,7 @@ where
 }
 
 impl<T> Reader<T> {
-    #[cfg(target_family = "wasm")]
+    #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
     /// Read the data.
     pub async fn read(self) -> T {
         match self {
@@ -78,7 +78,7 @@ impl<T> Reader<T> {
         }
     }
 
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     /// Read the data.
     pub fn read(self) -> T {
         match self {
@@ -92,9 +92,9 @@ impl<T> Reader<T> {
         match self {
             Self::Concrete(data) => Some(data),
             Self::Sync(reader) => Some(reader.read()),
-            #[cfg(target_family = "wasm")]
+            #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
             Self::Async(_func) => return None,
-            #[cfg(target_family = "wasm")]
+            #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
             Self::Future(_future) => return None,
         }
     }
@@ -106,10 +106,10 @@ impl<T> Reader<T> {
         O: 'static + Send,
         F: 'static + Send,
     {
-        #[cfg(target_family = "wasm")]
+        #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
         return Reader::Async(Box::new(MappedReader::new(self, mapper)));
 
-        #[cfg(not(target_family = "wasm"))]
+        #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
         Reader::Sync(Box::new(MappedReader::new(self, mapper)))
     }
 }

burn-core/Cargo.toml

@@ -29,6 +29,8 @@ dataset-minimal = ["burn-dataset"]
 dataset-sqlite = ["burn-dataset/sqlite"]
 dataset-sqlite-bundled = ["burn-dataset/sqlite-bundled"]
 
+wasm-sync = ["burn-tensor/wasm-sync", "burn-common/wasm-sync"]
+
 # Backend
 autodiff = ["burn-autodiff"]
 

burn-core/src/grad_clipping/base.rs

@@ -68,7 +68,7 @@ impl GradientClipping {
         clipped_grad.mask_fill(lower_mask, -threshold)
     }
 
-    #[cfg(target_family = "wasm")]
+    #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
     fn clip_by_norm<B: Backend, const D: usize>(
         &self,
         _grad: Tensor<B, D>,
@@ -77,7 +77,7 @@ impl GradientClipping {
         todo!("Not yet supported on wasm");
     }
 
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn clip_by_norm<B: Backend, const D: usize>(
         &self,
         grad: Tensor<B, D>,
@@ -96,7 +96,7 @@ impl GradientClipping {
         }
     }
 
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn l2_norm<B: Backend, const D: usize>(tensor: Tensor<B, D>) -> Tensor<B, 1> {
         let squared = tensor.powf(2.0);
         let sum = squared.sum();

burn-core/src/record/tensor.rs

@@ -44,7 +44,6 @@ impl<'de, S: PrecisionSettings> Deserialize<'de> for FloatTensorSerde<S> {
     }
 }
 
-// #[cfg(not(target_family = "wasm"))]
 impl<S: PrecisionSettings> Serialize for IntTensorSerde<S> {
     fn serialize<Se>(&self, serializer: Se) -> Result<Se::Ok, Se::Error>
     where
@@ -90,10 +89,10 @@ impl<B: Backend, const D: usize> Record for Tensor<B, D> {
     type Item<S: PrecisionSettings> = FloatTensorSerde<S>;
 
     fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
-        #[cfg(target_family = "wasm")]
+        #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
         todo!("Recording float tensors isn't yet supported on wasm.");
 
-        #[cfg(not(target_family = "wasm"))]
+        #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
         FloatTensorSerde::new(self.into_data().convert().serialize())
     }
 
@@ -106,10 +105,10 @@ impl<B: Backend, const D: usize> Record for Tensor<B, D, Int> {
     type Item<S: PrecisionSettings> = IntTensorSerde<S>;
 
     fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
-        #[cfg(target_family = "wasm")]
+        #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
         todo!("Recording int tensors isn't yet supported on wasm.");
 
-        #[cfg(not(target_family = "wasm"))]
+        #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
         IntTensorSerde::new(self.into_data().convert().serialize())
     }
 
@@ -122,10 +121,10 @@ impl<B: Backend, const D: usize> Record for Tensor<B, D, Bool> {
     type Item<S: PrecisionSettings> = BoolTensorSerde;
 
     fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
-        #[cfg(target_family = "wasm")]
+        #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
         todo!("Recording bool tensors isn't yet supported on wasm.");
 
-        #[cfg(not(target_family = "wasm"))]
+        #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
         BoolTensorSerde::new(self.into_data().serialize())
     }
 

burn-tensor/Cargo.toml

@@ -16,6 +16,7 @@ experimental-named-tensor = []
 export_tests = ["burn-tensor-testgen"]
 std = ["rand/std", "half/std"]
 benchmark = []
+wasm-sync = []
 
 [dependencies]
 burn-common = { path = "../burn-common", version = "0.10.0", default-features = false }

burn-tensor/src/tensor/api/base.rs

@@ -2,11 +2,11 @@
 
 use alloc::vec::Vec;
 
-#[cfg(not(target_family = "wasm"))]
+#[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
 use alloc::format;
-#[cfg(not(target_family = "wasm"))]
+#[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
 use alloc::string::String;
-#[cfg(not(target_family = "wasm"))]
+#[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
 use alloc::vec;
 
 use burn_common::{reader::Reader, stub::Mutex};
@@ -325,25 +325,25 @@ where
         Self::new(K::to_device(self.primitive, device))
     }
 
-    #[cfg(target_family = "wasm")]
+    #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
     /// Returns the data of the current tensor.
     pub async fn into_data(self) -> Data<K::Elem, D> {
         K::into_data(self.primitive).read().await
     }
 
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     /// Returns the data of the current tensor.
     pub fn into_data(self) -> Data<K::Elem, D> {
         K::into_data(self.primitive).read()
     }
 
-    #[cfg(target_family = "wasm")]
+    #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
     /// Returns the data of the current tensor.
     pub async fn to_data(&self) -> Data<K::Elem, D> {
         K::into_data(self.primitive.clone()).read().await
     }
 
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     /// Returns the data of the current tensor without taking ownership.
     pub fn to_data(&self) -> Data<K::Elem, D> {
         Self::into_data(self.clone())
@@ -467,7 +467,7 @@ where
     K: BasicOps<B>,
     <K as BasicOps<B>>::Elem: Debug,
 {
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     #[inline]
     fn push_newline_indent(acc: &mut String, indent: usize) {
         acc.push('\n');
@@ -476,7 +476,7 @@ where
         }
     }
 
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn fmt_inner_tensor(
         &self,
         acc: &mut String,
@@ -498,7 +498,7 @@ where
         }
     }
 
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn fmt_outer_tensor(
         &self,
         acc: &mut String,
@@ -533,7 +533,7 @@ where
     /// * `acc` - A mutable reference to a `String` used as an accumulator for the formatted output.
     /// * `depth` - The current depth of the tensor dimensions being processed.
     /// * `multi_index` - A mutable slice of `usize` representing the current indices in each dimension.
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn display_recursive(
         &self,
         acc: &mut String,
@@ -644,7 +644,7 @@ where
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         writeln!(f, "Tensor {{")?;
 
-        #[cfg(not(target_family = "wasm"))]
+        #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
         {
             let po = PRINT_OPTS.lock().unwrap();
             let mut acc = String::new();

burn-tensor/src/tensor/api/numeric.rs

@@ -9,7 +9,7 @@ where
     K: Numeric<B>,
     K::Elem: Element,
 {
-    #[cfg(not(target_family = "wasm"))]
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     /// Convert the tensor into a scalar.
     ///
     /// # Panics
@@ -21,7 +21,7 @@ where
         data.value[0]
     }
 
-    #[cfg(target_family = "wasm")]
+    #[cfg(all(not(feature = "wasm-sync"), target_family = "wasm"))]
     /// Convert the tensor into a scalar.
     ///
     /// # Panics

burn/Cargo.toml

@@ -18,6 +18,9 @@ std = ["burn-core/std"]
 # Training with full features
 train = ["burn-train/default", "autodiff", "dataset"]
 
+# Useful when targeting WASM and not using WGPU.
+wasm-sync = ["burn-core/wasm-sync"]
+
 ## Include nothing
 train-minimal = ["burn-train"]
 

examples/mnist-inference-web/build-for-web.sh

@@ -1,3 +1,4 @@
+#!/usr/bin/env bash
 
 # Add wasm32 target for compiler.
 rustup target add wasm32-unknown-unknown