Merge branch 'main' into jp/range_onnx_import

Cargo.toml

@@ -21,7 +21,7 @@ exclude = [
 
 [workspace.package]
 edition = "2021"
-version = "0.15.0"
+version = "0.14.0"
 readme = "README.md"
 license = "MIT OR Apache-2.0"
 
@@ -64,7 +64,6 @@ rusqlite = { version = "0.31.0" }
 rust-format = { version = "0.3.4" }
 sanitize-filename = "0.5.0"
 serde_rusqlite = "0.35.0"
-serde-wasm-bindgen = "0.6.5"
 spin = { version = "0.9.8", features = ["mutex", "spin_mutex"] }
 strum = "0.26.2"
 strum_macros = "0.26.2"
@@ -75,16 +74,12 @@ tokio = { version = "1.37.0", features = ["rt", "macros"] }
 tracing-appender = "0.2.3"
 tracing-core = "0.1.32"
 tracing-subscriber = "0.3.18"
-wasm-bindgen = "0.2.92"
-wasm-bindgen-futures = "0.4.42"
-wasm-logger = "0.2.0"
-wasm-timer = "0.2.5"
 md5 = "0.7.0"
 serial_test = "3.1.1"
 web-time = "1.1.0"
 hound = "3.5.1"
 image = "0.25.1"
-zip = "1.3.0"
+zip = "2.1.1"
 
 # Terminal UI
 ratatui = "0.26.3"

backend-comparison/Cargo.toml

@@ -29,9 +29,9 @@ cuda-jit = ["burn-cuda"]
 [dependencies]
 arboard = { workspace = true }
 burn = { path = "../crates/burn", default-features = false }
-burn-common = { path = "../crates/burn-common", version = "0.15.0" }
-burn-wgpu = { path = "../crates/burn-wgpu", default-features = false, version = "0.15.0" }
-burn-cuda = { path = "../crates/burn-cuda", version = "0.15.0", optional = true }
+burn-common = { path = "../crates/burn-common", version = "0.14.0" }
+burn-wgpu = { path = "../crates/burn-wgpu", default-features = false, version = "0.14.0" }
+burn-cuda = { path = "../crates/burn-cuda", version = "0.14.0", optional = true }
 clap = { workspace = true }
 colored = { workspace = true }
 derive-new = { workspace = true }

burn-book/src/basic-workflow/README.md

@@ -14,7 +14,7 @@ automatically add the missing imports as you add the code snippets to your code.
 Be sure to checkout the git branch corresponding to the version of Burn you are using to follow
 this guide.
 
-The current version of Burn is `0.15` and the corresponding branch to checkout is `main`.
+The current version of Burn is `0.14` and the corresponding branch to checkout is `main`.
 </div>
 
 The code for this demo can be executed from Burn's base directory using the command:

burn-book/src/basic-workflow/model.md

@@ -20,7 +20,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-burn = { version = "0.15.0", features = ["train", "wgpu", "vision"] }
+burn = { version = "0.14.0", features = ["train", "wgpu", "vision"] }
 ```
 
 Our goal will be to create a basic convolutional neural network used for image classification. We

burn-book/src/building-blocks/tensor.md

@@ -288,6 +288,7 @@ Those operations are only available for `Int` tensors.
 | `tensor.float()`                                 | `tensor.to(torch.float)`                                |
 | `tensor.from_ints(ints)`                         | N/A                                                     |
 | `tensor.int_random(shape, distribution, device)` | N/A                                                     |
+| `tensor.cartesian_grid(shape, device)`           | N/A                                                     |
 
 # Bool Operations
 

crates/burn-autodiff/Cargo.toml

@@ -17,15 +17,15 @@ std = []
 async = [] # Require std
 
 [dependencies]
-burn-common = { path = "../burn-common", version = "0.15.0" }
-burn-tensor = { path = "../burn-tensor", version = "0.15.0", default-features = false }
-burn-tensor-testgen = { path = "../burn-tensor-testgen", version = "0.15.0", optional = true }
+burn-common = { path = "../burn-common", version = "0.14.0" }
+burn-tensor = { path = "../burn-tensor", version = "0.14.0", default-features = false }
+burn-tensor-testgen = { path = "../burn-tensor-testgen", version = "0.14.0", optional = true }
 
 derive-new = { workspace = true }
 spin = { workspace = true }
 log = { workspace = true }
 
 [dev-dependencies]
-burn-tensor = { path = "../burn-tensor", version = "0.15.0", default-features = false, features = [
+burn-tensor = { path = "../burn-tensor", version = "0.14.0", default-features = false, features = [
   "export_tests",
 ] }

crates/burn-candle/Cargo.toml

@@ -21,17 +21,17 @@ accelerate = ["candle-core/accelerate"]
 
 [dependencies]
 derive-new = { workspace = true }
-burn-tensor = { path = "../burn-tensor", version = "0.15.0", default-features = false }
+burn-tensor = { path = "../burn-tensor", version = "0.14.0", default-features = false }
 half = { workspace = true }
 candle-core = { workspace = true }
 
 [dev-dependencies]
-burn-autodiff = { path = "../burn-autodiff", version = "0.15.0", default-features = false, features = [
+burn-autodiff = { path = "../burn-autodiff", version = "0.14.0", default-features = false, features = [
     "export_tests",
 ] }
-burn-tch = { path = "../burn-tch", version = "0.15.0", default-features = false, features = [
+burn-tch = { path = "../burn-tch", version = "0.14.0", default-features = false, features = [
 ] }
-burn-tensor = { path = "../burn-tensor", version = "0.15.0", default-features = false, features = [
+burn-tensor = { path = "../burn-tensor", version = "0.14.0", default-features = false, features = [
     "export_tests",
 ] }
 

crates/burn-compute/Cargo.toml

@@ -34,7 +34,7 @@ autotune-persistent-cache = [
 ] # Assume std
 
 [dependencies]
-burn-common = { path = "../burn-common", version = "0.15.0", default-features = false }
+burn-common = { path = "../burn-common", version = "0.14.0", default-features = false }
 derive-new = { workspace = true }
 spin = { workspace = true }
 log = { workspace = true }

crates/burn-core/Cargo.toml

@@ -99,17 +99,17 @@ test-wgpu = ["wgpu"] # To use wgpu during testing, default uses ndarray.
 
 # ** Please make sure all dependencies support no_std when std is disabled **
 
-burn-common = { path = "../burn-common", version = "0.15.0", default-features = false }
-burn-dataset = { path = "../burn-dataset", version = "0.15.0", optional = true, default-features = false }
-burn-derive = { path = "../burn-derive", version = "0.15.0" }
-burn-tensor = { path = "../burn-tensor", version = "0.15.0", default-features = false }
+burn-common = { path = "../burn-common", version = "0.14.0", default-features = false }
+burn-dataset = { path = "../burn-dataset", version = "0.14.0", optional = true, default-features = false }
+burn-derive = { path = "../burn-derive", version = "0.14.0" }
+burn-tensor = { path = "../burn-tensor", version = "0.14.0", default-features = false }
 
 # Backends
-burn-ndarray = { path = "../burn-ndarray", version = "0.15.0", optional = true, default-features = false }
-burn-wgpu = { path = "../burn-wgpu", version = "0.15.0", optional = true, default-features = false }
-burn-autodiff = { path = "../burn-autodiff", version = "0.15.0", optional = true }
-burn-tch = { path = "../burn-tch", version = "0.15.0", optional = true }
-burn-candle = { path = "../burn-candle", version = "0.15.0", optional = true }
+burn-ndarray = { path = "../burn-ndarray", version = "0.14.0", optional = true, default-features = false }
+burn-wgpu = { path = "../burn-wgpu", version = "0.14.0", optional = true, default-features = false }
+burn-autodiff = { path = "../burn-autodiff", version = "0.14.0", optional = true }
+burn-tch = { path = "../burn-tch", version = "0.14.0", optional = true }
+burn-candle = { path = "../burn-candle", version = "0.14.0", optional = true }
 
 derive-new = { workspace = true }
 log = { workspace = true, optional = true }
@@ -135,12 +135,12 @@ num-traits = { workspace = true }
 
 [dev-dependencies]
 tempfile = { workspace = true }
-burn-dataset = { path = "../burn-dataset", version = "0.15.0", features = [
+burn-dataset = { path = "../burn-dataset", version = "0.14.0", features = [
     "fake",
 ] }
 
-burn-ndarray = { path = "../burn-ndarray", version = "0.15.0", default-features = false }
-burn-autodiff = { path = "../burn-autodiff", version = "0.15.0" }
+burn-ndarray = { path = "../burn-ndarray", version = "0.14.0", default-features = false }
+burn-autodiff = { path = "../burn-autodiff", version = "0.14.0" }
 
 [package.metadata.docs.rs]
 features = ["doc"]

crates/burn-cube-macros/src/analysis.rs

@@ -4,6 +4,8 @@ use syn::{PathArguments, Stmt};
 
 use crate::VariableKey;
 
+pub const KEYWORDS: [&str; 1] = ["ABSOLUTE_INDEX"];
+
 #[derive(Debug)]
 /// Information about a single variable's use in Cube code
 /// Information about a single variable's use in Cube code
@@ -200,8 +202,10 @@ impl CodeAnalysisBuilder {
                     .get_ident()
                     .expect("Analysis: only ident path are supported.");
 
-                // Use
-                self.var_uses.push(ident.into());
+                if !KEYWORDS.contains(&ident.to_string().as_str()) {
+                    // Use
+                    self.var_uses.push(ident.into());
+                }
             }
             syn::Expr::Binary(expr) => {
                 self.find_occurrences_in_expr(&expr.left, depth);

crates/burn-cube-macros/src/codegen/variable.rs

@@ -2,7 +2,10 @@ use proc_macro2::TokenStream;
 use quote::ToTokens;
 use syn::Lit;
 
-use crate::{analysis::CodeAnalysis, codegen::base::codegen_expr};
+use crate::{
+    analysis::{CodeAnalysis, KEYWORDS},
+    codegen::base::codegen_expr,
+};
 
 /// Codegen for literals
 pub(crate) fn codegen_lit(lit: &syn::ExprLit) -> TokenStream {
@@ -11,10 +14,10 @@ pub(crate) fn codegen_lit(lit: &syn::ExprLit) -> TokenStream {
         Lit::Float(_) => {
             let lit_str = lit.lit.to_token_stream().to_string();
             let float_lit = lit_str.parse::<f32>().unwrap();
-            quote::quote! { #float_lit.into() }
+            quote::quote! { #float_lit }
         }
         _ => {
-            quote::quote! { #lit.into() }
+            quote::quote! { #lit }
         }
     }
 }
@@ -140,15 +143,21 @@ pub(crate) fn codegen_path_rhs(
         .get_ident()
         .expect("Codegen: Only ident path are supported.");
 
-    let will_be_used_again = variable_analyses.should_clone(ident, loop_level);
-
-    if will_be_used_again {
+    if KEYWORDS.contains(&ident.to_string().as_str()) {
         quote::quote! {
-            #ident.clone()
+            #ident :: expand(context)
         }
     } else {
-        quote::quote! {
-            #ident
+        let will_be_used_again = variable_analyses.should_clone(ident, loop_level);
+
+        if will_be_used_again {
+            quote::quote! {
+                #ident.clone()
+            }
+        } else {
+            quote::quote! {
+                #ident
+            }
         }
     }
 }

crates/burn-cube/Cargo.toml

@@ -20,13 +20,13 @@ template = []
 tensor = ["burn-tensor"]
 
 [dependencies]
-burn-compute = { path = "../burn-compute", version = "0.15.0", default-features = false }
-burn-tensor = { path = "../burn-tensor", version = "0.15.0", default-features = false, optional = true }
+burn-compute = { path = "../burn-compute", version = "0.14.0", default-features = false }
+burn-tensor = { path = "../burn-tensor", version = "0.14.0", default-features = false, optional = true }
 
 bytemuck = { workspace = true }
 half = { workspace = true, features = ["bytemuck"] }
 serde = { workspace = true }
-burn-cube-macros = { path = "../burn-cube-macros", version = "0.15.0" }
+burn-cube-macros = { path = "../burn-cube-macros", version = "0.14.0" }
 derive-new = { workspace = true }
 
 log = { workspace = true }

crates/burn-cube/src/language/branch.rs

@@ -16,15 +16,15 @@ where
     (start.val..end.val).map(UInt::new)
 }
 
-pub fn range_expand<F>(
-    context: &mut CubeContext,
-    start: ExpandElement,
-    end: ExpandElement,
-    unroll: bool,
-    mut func: F,
-) where
+pub fn range_expand<F, S, E>(context: &mut CubeContext, start: S, end: E, unroll: bool, mut func: F)
+where
     F: FnMut(&mut CubeContext, ExpandElement),
+    S: Into<ExpandElement>,
+    E: Into<ExpandElement>,
 {
+    let start: ExpandElement = start.into();
+    let end: ExpandElement = end.into();
+
     if unroll {
         let start = match start.deref() {
             Variable::ConstantScalar(val, _) => *val as usize,

crates/burn-cube/src/language/comptime.rs

@@ -5,7 +5,7 @@ use crate::{unexpanded, CubeContext, CubeType};
 ///
 /// Use `Comptime<Option<T>>` to have an alternate runtime behaviour if the compilation time value is not present
 pub struct Comptime<T> {
-    inner: T,
+    pub(crate) inner: T,
 }
 
 impl<T> Comptime<T> {

crates/burn-cube/src/language/context.rs

@@ -109,6 +109,10 @@ impl CubeContext {
         new
     }
 
+    pub fn create_shared(&mut self, item: Item, size: u32) -> ExpandElement {
+        ExpandElement::Plain(self.root.borrow_mut().create_shared(item, size))
+    }
+
     /// Obtain the index-th input
     pub fn input(&mut self, index: u16, item: Item) -> ExpandElement {
         ExpandElement::Plain(crate::dialect::Variable::GlobalInputArray(index, item))

crates/burn-cube/src/language/element/cast.rs

@@ -1,16 +1,20 @@
 use crate::dialect::Item;
 use crate::language::{assign, CubeContext, CubeElem, CubeType};
-use crate::unexpanded;
+use crate::{unexpanded, ExpandElement};
 
 /// Enable elegant casting from any to any CubeElem
 pub trait Cast: CubeElem {
     fn cast_from<From: CubeElem>(value: From) -> Self;
-    fn cast_from_expand(
+
+    fn cast_from_expand<From>(
         context: &mut CubeContext,
-        value: <Self as CubeType>::ExpandType,
-    ) -> <Self as CubeType>::ExpandType {
+        value: From,
+    ) -> <Self as CubeType>::ExpandType
+    where
+        From: Into<ExpandElement>,
+    {
         let new_var = context.create_local(Item::new(<Self as CubeElem>::as_elem()));
-        assign::expand(context, value, new_var.clone());
+        assign::expand(context, value.into(), new_var.clone());
         new_var
     }
 }

crates/burn-cube/src/language/element/float.rs

@@ -6,15 +6,15 @@ use crate::{Ceil, Cos, Erf, Exp, Floor, Log, Log1p, Powf, Recip, Sin, Sqrt, Tanh
 pub trait Float:
     Numeric + Exp + Log + Log1p + Cos + Sin + Tanh + Powf + Sqrt + Floor + Ceil + Erf + Recip
 {
-    fn new(val: f64) -> Self;
-    fn new_expand(context: &mut CubeContext, val: f64) -> <Self as CubeType>::ExpandType;
+    fn new(val: f32) -> Self;
+    fn new_expand(context: &mut CubeContext, val: f32) -> <Self as CubeType>::ExpandType;
 }
 
 macro_rules! impl_float {
     ($type:ident) => {
         #[derive(Clone, Copy)]
         pub struct $type {
-            pub val: f64,
+            pub val: f32,
             pub vectorization: u8,
         }
 
@@ -32,15 +32,15 @@ macro_rules! impl_float {
         impl Numeric for $type {}
 
         impl Float for $type {
-            fn new(val: f64) -> Self {
+            fn new(val: f32) -> Self {
                 Self {
                     val,
                     vectorization: 1,
                 }
             }
 
-            fn new_expand(_context: &mut CubeContext, val: f64) -> <Self as CubeType>::ExpandType {
-                let new_var = Variable::ConstantScalar(val, Self::as_elem());
+            fn new_expand(_context: &mut CubeContext, val: f32) -> <Self as CubeType>::ExpandType {
+                let new_var = Variable::ConstantScalar(val as f64, Self::as_elem());
                 ExpandElement::Plain(new_var)
             }
         }

crates/burn-cube/src/language/element/mod.rs

@@ -6,6 +6,7 @@ mod cube_elem;
 mod float;
 mod int;
 mod numeric;
+mod shared_memory;
 mod tensor;
 mod uint;
 
@@ -16,5 +17,6 @@ pub use cube_elem::*;
 pub use float::*;
 pub use int::*;
 pub use numeric::*;
+pub use shared_memory::*;
 pub use tensor::*;
 pub use uint::*;

crates/burn-cube/src/language/element/numeric.rs

@@ -50,7 +50,7 @@ pub trait Numeric:
             vectorization: (vec.len() as u8),
         });
         for (i, element) in vec.iter().enumerate() {
-            new_var = index_assign::expand(context, new_var, i.into(), (*element).into());
+            new_var = index_assign::expand(context, new_var, i, *element);
         }
 
         new_var