Migration/cubecl (#2041)

Cargo.toml

@@ -16,7 +16,7 @@ members = [
 
 exclude = [
     "examples/notebook",
-    # "crates/burn-cuda"     # comment this line to work on burn-cuda
+    "crates/burn-cuda",     # comment this line to work on burn-cuda
 ]
 
 [workspace.package]
@@ -27,7 +27,7 @@ license = "MIT OR Apache-2.0"
 
 [workspace.dependencies]
 bytemuck = "1.16.1"
-candle-core = { version = "0.5.1" }
+candle-core = { version = "0.6.0" }
 clap = { version = "4.5.9", features = ["derive"] }
 colored = "2.1.0"
 console_error_panic_hook = "0.1.7"
@@ -140,6 +140,9 @@ nvml-wrapper = "0.10.0"
 sysinfo = "0.30.13"
 systemstat = "0.2.3"
 
+cubecl = { version = "0.1.1", git = "https://github.com/tracel-ai/cubecl", default-features = false }
+cubecl-common = { version = "0.1.1", git = "https://github.com/tracel-ai/cubecl", default-features = false }
 
 [profile.dev]
 debug = 0 # Speed up compilation time and not necessary.
+opt-level = 2

backend-comparison/Cargo.toml

@@ -24,14 +24,14 @@ tch-cpu = ["burn/tch"]
 tch-gpu = ["burn/tch"]
 wgpu = ["burn/wgpu", "burn/autotune"]
 wgpu-fusion = ["wgpu", "burn/fusion"]
-cuda-jit = ["burn-cuda"]
+# cuda-jit = ["burn-cuda"]
 
 [dependencies]
 arboard = { workspace = true }
 burn = { path = "../crates/burn", default-features = false }
 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 }
+burn-wgpu = { path = "../crates/burn-wgpu", default-features = false, version = "0.14.0", optional = true }
+# burn-cuda = { path = "../crates/burn-cuda", version = "0.14.0", optional = true }
 clap = { workspace = true }
 colored = { workspace = true }
 derive-new = { workspace = true }
@@ -49,6 +49,7 @@ strum_macros = { workspace = true }
 sysinfo = { workspace = true, features = ["serde"] }
 wgpu = { workspace = true }
 wsl = { workspace = true }
+cubecl = { workspace = true, features = ["wgpu"] }
 
 [dev-dependencies]
 rstest = { workspace = true }

backend-comparison/benches/matmul.rs

@@ -29,7 +29,7 @@ impl<B: Backend, const D: usize> Benchmark for MatmulBenchmark<B, D> {
     }
 
     fn execute(&self, (lhs, rhs): Self::Args) {
-        lhs.clone().transpose().matmul(rhs.clone());
+        lhs.clone().matmul(rhs.clone());
     }
 
     fn prepare(&self) -> Self::Args {
@@ -52,11 +52,11 @@ fn bench<B: Backend>(
     token: Option<&str>,
 ) {
     const D: usize = 3;
-    let batch_size = 32;
-    let m = 256;
-    let k = 1024;
-    let n = 256;
-    let shape_lhs = [batch_size, k, m].into();
+    let batch_size = 8;
+    let m = 2048;
+    let k = 2048;
+    let n = 2048;
+    let shape_lhs = [batch_size, m, k].into();
     let shape_rhs = [batch_size, k, n].into();
 
     let benchmark = MatmulBenchmark::<B, D>::new(shape_lhs, shape_rhs, device.clone());

backend-comparison/src/persistence/system_info.rs

@@ -1,5 +1,5 @@
 use burn::serde::{Deserialize, Serialize};
-use burn_wgpu::GraphicsApi;
+use cubecl::wgpu::GraphicsApi;
 use std::collections::HashSet;
 use sysinfo;
 use wgpu;
@@ -51,7 +51,7 @@ impl BenchmarkSystemInfo {
     fn enumerate_gpus() -> Vec<String> {
         let instance = wgpu::Instance::default();
         let adapters: Vec<wgpu::Adapter> = instance
-            .enumerate_adapters(burn_wgpu::AutoGraphicsApi::backend().into())
+            .enumerate_adapters(cubecl::wgpu::AutoGraphicsApi::backend().into())
             .into_iter()
             .filter(|adapter| {
                 let info = adapter.get_info();

crates/burn-common/Cargo.toml

@@ -11,8 +11,8 @@ repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-common"
 version.workspace = true
 
 [features]
-default = ["std"]
-std = ["rand/std", "data-encoding/std", "dep:pollster"]
+default = ["std", "cubecl-common/default"]
+std = ["cubecl-common/std"]
 doc = ["default"]
 network = ["dep:indicatif", "dep:reqwest", "dep:tokio"]
 rayon = ["dep:rayon"]
@@ -23,13 +23,7 @@ web-time = { version = "1.1.0" }
 
 
 [dependencies]
-# ** Please make sure all dependencies support no_std when std is disabled **
-rand = { workspace = true }
-spin = { workspace = true }                      # using in place of use std::sync::Mutex;
-derive-new = { workspace = true }
-serde = { workspace = true }
 data-encoding = { workspace = true }
-pollster = { workspace = true, optional = true }
 
 # Network downloader
 indicatif = { workspace = true, optional = true }
@@ -38,6 +32,7 @@ tokio = { workspace = true, optional = true }
 
 # Parallel
 rayon = { workspace = true, optional = true }
+cubecl-common = { workspace = true, default-features = false }
 
 [dev-dependencies]
 dashmap = { workspace = true }