Chore update cubecl (tracel-ai#3764)

nathanielsimard · web-flow · commit 3a2a1a8941d4 · 2025-09-22T10:30:48.000-04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -169,9 +169,9 @@ portable-atomic = { version = "1.11.1" }
 portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "d4aec5dd7ef73068347eb688f2327be5605bd75b" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "d4aec5dd7ef73068347eb688f2327be5605bd75b" }
-cubecl-quant = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "d4aec5dd7ef73068347eb688f2327be5605bd75b" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "ad03f7892f4069d3406f936fd9fa3ac3b090c1fe" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "ad03f7892f4069d3406f936fd9fa3ac3b090c1fe" }
+cubecl-quant = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "ad03f7892f4069d3406f936fd9fa3ac3b090c1fe" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }
diff --git a/crates/burn-cpu/Cargo.toml b/crates/burn-cpu/Cargo.toml
@@ -21,7 +21,6 @@ default = ["std", "fusion", "autotune", "burn-cubecl/default", "cubecl/default"]
 doc = ["burn-cubecl/doc"]
 fusion = ["burn-fusion", "burn-cubecl/fusion"]
 std = ["burn-cubecl/std", "cubecl/std"]
-compilation-cache = ["cubecl/compilation-cache"]
 
 [dependencies]
 burn-fusion = { path = "../burn-fusion", version = "0.19.0", optional = true }
diff --git a/crates/burn-cubecl-fusion/src/matmul/optimization.rs b/crates/burn-cubecl-fusion/src/matmul/optimization.rs
@@ -22,6 +22,8 @@ use burn_ir::TensorIr;
 use cubecl::features::TypeUsage;
 use cubecl::matmul::components::AccG;
 use cubecl::matmul::components::AccS;
+use cubecl::matmul::components::tile::loader::Filled;
+use cubecl::matmul::kernels::layered::Selection;
 use cubecl::matmul::kernels::layered::double_buffering::CyclicDoubleBufferingAlgorithm;
 use cubecl::matmul::kernels::layered::double_buffering::DoubleBufferingArgs;
 use cubecl::matmul::kernels::layered::double_unit::DoubleUnitAlgorithm;
@@ -33,7 +35,6 @@ use cubecl::matmul::kernels::layered::simple::SimpleArgs;
 use cubecl::matmul::kernels::layered::simple_unit::SimpleUnitAlgorithm;
 use cubecl::matmul::kernels::layered::vecmat::DoubleVecMatAlgorithm;
 use cubecl::matmul::kernels::layered::vecmat::SimpleVecMatAlgorithm;
-use cubecl::matmul::{components::tile::reader::Filled, kernels::layered::Selection};
 use cubecl::matmul::{
     components::{LhsS, MatmulLineSizes, MatmulPrecision},
     kernels::layered::Algorithm,
diff --git a/crates/burn-cubecl-fusion/src/shared/io.rs b/crates/burn-cubecl-fusion/src/shared/io.rs
@@ -248,7 +248,8 @@ pub fn input_as_linear_view<C: CubePrimitive>(
     #[comptime] pos: u32,
 ) -> View<Line<C>, u32> {
     let slice = input_as_slice::<Line<C>>(inputs, pos);
-    let layout = LinearLayout::new_Plain(PlainLayout::new(slice.len()));
+    let line_size = slice.line_size();
+    let layout = LinearLayout::new_Plain(PlainLayout::new(slice.len(), comptime!(line_size as u8)));
     View::new::<Slice<Line<C>>, u32>(&slice, layout)
 }
 
diff --git a/crates/burn-cubecl/src/kernel/matmul/tune/base.rs b/crates/burn-cubecl/src/kernel/matmul/tune/base.rs
@@ -1,7 +1,7 @@
 use burn_tensor::DType;
 use cubecl::{
     matmul::{
-        MatmulInputHandleRef, Strategy, SyncPartialReadingStrategy, SyncReadingStrategy,
+        MatmulInputHandleRef, Strategy, SyncLoadingStrategy, SyncPartialLoadingStrategy,
         components::{AccG, MatmulKind},
         kernels::layered::{
             Selection, TileSizeSelection, double_buffering::DoubleBufferingArgs,
@@ -161,7 +161,7 @@ fn matmul_simple<R: CubeRuntime, E: MatmulElement>(
 ) -> Result<(), String> {
     cubecl::matmul::launch_ref::<R, E>(
         &Strategy::Simple(
-            SyncReadingStrategy::Cyclic,
+            SyncLoadingStrategy::Cyclic,
             Selection::Inferred(SimpleArgs { multi_rows: false }),
         ),
         &lhs.client,
@@ -179,7 +179,7 @@ fn matmul_simple_multi_rows<R: CubeRuntime, E: MatmulElement>(
 ) -> Result<(), String> {
     cubecl::matmul::launch_ref::<R, E>(
         &Strategy::Simple(
-            SyncReadingStrategy::Cyclic,
+            SyncLoadingStrategy::Cyclic,
             Selection::Inferred(SimpleArgs { multi_rows: true }),
         ),
         &lhs.client,
@@ -197,7 +197,7 @@ fn matmul_double_buffering<R: CubeRuntime, E: MatmulElement>(
 ) -> Result<(), String> {
     cubecl::matmul::launch_ref::<R, E>(
         &Strategy::DoubleBuffering(
-            SyncPartialReadingStrategy::Tilewise,
+            SyncPartialLoadingStrategy::Tilewise,
             Selection::Inferred(DoubleBufferingArgs { specialized: false }),
         ),
         &lhs.client,
@@ -215,7 +215,7 @@ fn matmul_double_buffering_specialized<R: CubeRuntime, E: MatmulElement>(
 ) -> Result<(), String> {
     cubecl::matmul::launch_ref::<R, E>(
         &Strategy::DoubleBuffering(
-            SyncPartialReadingStrategy::Tilewise,
+            SyncPartialLoadingStrategy::Tilewise,
             Selection::Inferred(DoubleBufferingArgs { specialized: true }),
         ),
         &lhs.client,
diff --git a/crates/burn-cuda/Cargo.toml b/crates/burn-cuda/Cargo.toml
@@ -21,7 +21,6 @@ default = ["std", "fusion", "autotune", "burn-cubecl/default", "cubecl/default"]
 doc = ["burn-cubecl/doc"]
 fusion = ["burn-fusion", "burn-cubecl/fusion"]
 std = ["burn-cubecl/std", "cubecl/std"]
-compilation-cache = ["cubecl/compilation-cache"]
 
 [dependencies]
 burn-fusion = { path = "../burn-fusion", version = "0.19.0", optional = true }
diff --git a/crates/burn-fusion/src/stream/multi.rs b/crates/burn-fusion/src/stream/multi.rs
@@ -311,7 +311,12 @@ impl<R: FusionRuntime> MultiStream<R> {
         }
 
         for id in streams_to_sync.drain() {
+            let old = unsafe { StreamId::swap(id) };
+            log::info!("Drain stream {id} for use in current {current}");
             self.resolve_stream(handles, id, nodes);
+            unsafe {
+                StreamId::swap(old);
+            };
         }
     }
 
diff --git a/crates/burn-rocm/Cargo.toml b/crates/burn-rocm/Cargo.toml
@@ -21,7 +21,6 @@ autotune = ["burn-cubecl/autotune"]
 autotune-checks = ["burn-cubecl/autotune-checks"]
 doc = ["burn-cubecl/doc"]
 std = ["burn-cubecl/std", "cubecl/std"]
-compilation-cache = ["cubecl/compilation-cache"]
 
 [dependencies]
 cubecl = { workspace = true, features = ["hip"] }
diff --git a/crates/burn/Cargo.toml b/crates/burn/Cargo.toml
@@ -134,10 +134,6 @@ metal = ["wgpu", "burn-wgpu/metal"]
 wgpu = ["burn-wgpu"]
 cpu = ["burn-cpu"]
 
-compilation-cache = [
-    "burn-cuda?/compilation-cache",
-    "burn-rocm?/compilation-cache",
-]
 
 [dependencies]
 
@@ -162,4 +158,4 @@ burn-tch = { path = "../burn-tch", version = "0.19.0", default-features = false,
 burn-wgpu = { path = "../burn-wgpu", version = "0.19.0", optional = true, default-features = false }
 burn-ir = { path = "../burn-ir", version = "0.19.0", optional = true, default-features = false }
 
-cubecl = { workspace = true, default-features = false, optional = true }
+cubecl = { workspace = true, default-features = false, optional = true }
diff --git a/examples/server/Cargo.toml b/examples/server/Cargo.toml
@@ -11,7 +11,7 @@ workspace = true
 
 [features]
 default = ["webgpu"]
-cuda = ["burn/cuda", "cubecl/compilation-cache"]
+cuda = ["burn/cuda"]
 webgpu = ["burn/webgpu"]
 vulkan = ["burn/vulkan"]
 ndarray = ["burn/ndarray"]
diff --git a/examples/text-classification/examples/ag-news-train.rs b/examples/text-classification/examples/ag-news-train.rs
@@ -18,7 +18,7 @@ type ElemType = burn::tensor::flex32;
 
 pub fn launch<B: AutodiffBackend>(devices: Vec<B::Device>) {
     let config = ExperimentConfig::new(
-        TransformerEncoderConfig::new(512, 2048, 16, 8)
+        TransformerEncoderConfig::new(256, 1024, 8, 4)
             .with_norm_first(true)
             .with_quiet_softmax(true),
         AdamConfig::new().with_weight_decay(Some(WeightDecayConfig::new(5e-5))),
@@ -135,7 +135,7 @@ mod cuda {
 
     pub fn run() {
         let type_id = 0;
-        let num_devices = B::Device::device_count(type_id);
+        let num_devices = 1;
 
         let devices = (0..num_devices)
             .map(|i| CudaDevice::new(i as usize))
@@ -148,10 +148,10 @@ mod cuda {
 #[cfg(feature = "rocm")]
 mod rocm {
     use crate::{ElemType, launch};
-    use burn::backend::{Autodiff, Rocm};
+    use burn::backend::{Autodiff, Rocm, autodiff::checkpoint::strategy::BalancedCheckpointing};
 
     pub fn run() {
-        launch::<Autodiff<Rocm<ElemType, i32>>>(vec![Default::default()]);
+        launch::<Autodiff<Rocm<ElemType, i32>, BalancedCheckpointing>>(vec![Default::default()]);
     }
 }
 
diff --git a/examples/text-classification/src/training.rs b/examples/text-classification/src/training.rs
@@ -37,7 +37,7 @@ pub struct ExperimentConfig {
     pub optimizer: AdamConfig,
     #[config(default = 256)]
     pub max_seq_length: usize,
-    #[config(default = 64)]
+    #[config(default = 32)]
     pub batch_size: usize,
     #[config(default = 5)]
     pub num_epochs: usize,

Original file line number	Diff line number	Diff line change
`@@ -248,7 +248,8 @@ pub fn input_as_linear_view<C: CubePrimitive>(`
`248`	`248`	`#[comptime] pos: u32,`
`249`	`249`	`) -> View<Line<C>, u32> {`
`250`	`250`	`let slice = input_as_slice::<Line<C>>(inputs, pos);`
`251`		`- let layout = LinearLayout::new_Plain(PlainLayout::new(slice.len()));`
	`251`	`+ let line_size = slice.line_size();`
	`252`	`+ let layout = LinearLayout::new_Plain(PlainLayout::new(slice.len(), comptime!(line_size as u8)));`
`252`	`253`	`View::new::<Slice<Line<C>>, u32>(&slice, layout)`
`253`	`254`	`}`
`254`	`255`
Original file line number	Diff line number	Diff line change
`@@ -311,7 +311,12 @@ impl<R: FusionRuntime> MultiStream<R> {`
`311`	`311`	`}`
`312`	`312`
`313`	`313`	`for id in streams_to_sync.drain() {`
	`314`	`+ let old = unsafe { StreamId::swap(id) };`
	`315`	`+ log::info!("Drain stream {id} for use in current {current}");`
`314`	`316`	`self.resolve_stream(handles, id, nodes);`
	`317`	`+ unsafe {`
	`318`	`+ StreamId::swap(old);`
	`319`	`+ };`
`315`	`320`	`}`
`316`	`321`	`}`
`317`	`322`