Fix cubecl std usage (tracel-ai#3306)

Author: laggui

Cargo.lock

@@ -109,7 +109,6 @@ chrono = "0.4.39"
 os_info = "3.10.0"
 wsl = "0.1.0"
 
-# TODO: update to official 2.0 release once MSRV is bumped to 1.85
 bincode = { version = "2.0.1", features = [
     "alloc",
     "serde",
@@ -158,17 +157,14 @@ 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 = "c5817b4e6bc355d9f231a411927c72c637699a09" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c5817b4e6bc355d9f231a411927c72c637699a09" }
-cubecl-std = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c5817b4e6bc355d9f231a411927c72c637699a09" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "4dc928dbd34c59884ba774893df4331462a71386" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "4dc928dbd34c59884ba774893df4331462a71386" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }
-# cubecl-std = { path = "../cubecl/crates/cubecl-std", default-features = false }
 ### For the release. ###
 # cubecl = { version = "0.5.0", default-features = false }
 # cubecl-common = { version = "0.5.0", default-features = false }
-# cubecl-std = { version = "0.5.0", default-features = false }
 
 ### For xtask crate ###
 tracel-xtask = { version = "=1.1.9" }

Cargo.toml

@@ -29,7 +29,6 @@ burn-tensor = { path = "../burn-tensor", version = "0.18.0", default-features =
     "cubecl",
 ] }
 cubecl = { workspace = true, features = ["matmul", "convolution", "reduce"] }
-cubecl-std = { workspace = true }
 
 half = { workspace = true }
 serde = { workspace = true }

crates/burn-cubecl-fusion/Cargo.toml

@@ -22,8 +22,8 @@ use cubecl::matmul::kernels::matmul::double_buffering::CyclicDoubleBufferingAlgo
 use cubecl::matmul::kernels::matmul::launch_kernel_virtual;
 use cubecl::matmul::kernels::matmul::simple::SimpleAlgorithm;
 use cubecl::matmul::kernels::{MatmulAvailabilityError, MatmulSetupError};
+use cubecl::std::tensor::{MatrixBatchLayout, matrix_batch_layout};
 use cubecl::{client::ComputeClient, prelude::*};
-use cubecl_std::tensor::{MatrixBatchLayout, matrix_batch_layout};
 use half::{bf16, f16};
 use serde::{Deserialize, Serialize};
 

crates/burn-cubecl-fusion/src/matmul/optimization.rs

@@ -60,8 +60,8 @@ cubecl = { workspace = true, features = [
     "convolution",
     "reduce",
     "random",
+    "stdlib",
 ] }
-cubecl-std = { workspace = true }
 
 bytemuck = { workspace = true }
 derive-new = { workspace = true }

crates/burn-cubecl/Cargo.toml

@@ -1,9 +1,9 @@
 use burn_tensor::ops::{ConvOptions, conv::calculate_conv_output_sizes};
+use cubecl::std::{CubeOption, CubeOptionExpand, FastDivmod};
 use cubecl::{
     calculate_cube_count_elemwise, convolution::ConvLaunchError, prelude::*,
     std::tensor::StridedLayout, tensor_line_size_parallel,
 };
-use cubecl_std::{CubeOption, CubeOptionExpand, FastDivmod};
 
 use crate::{
     CubeElement, CubeRuntime,

crates/burn-cubecl/src/kernel/conv/direct.rs

@@ -3,14 +3,14 @@ use burn_tensor::{
     ops::{ConvOptions, conv::calculate_conv_output_sizes},
 };
 use core::iter;
+use cubecl::std::{FastDivmod, FastDivmodArgs};
 use cubecl::{
     calculate_cube_count_elemwise,
     convolution::ConvLaunchError,
     intrinsic,
     prelude::*,
     std::tensor::{TensorHandle, into_contiguous_pitched},
 };
-use cubecl_std::{FastDivmod, FastDivmodArgs};
 
 use crate::{
     CubeElement, CubeRuntime, FloatElement,

crates/burn-cubecl/src/kernel/conv/im2col.rs

@@ -1,6 +1,6 @@
 use crate::{CubeRuntime, element::CubeElement, tensor::CubeTensor};
+use cubecl::std::{FastDivmod, FastDivmodArgs};
 use cubecl::{calculate_cube_count_elemwise, prelude::*};
-use cubecl_std::{FastDivmod, FastDivmodArgs};
 use std::ops::Range;
 
 #[cube(launch_unchecked)]

crates/burn-cubecl/src/kernel/index/slice_assign.rs

@@ -1,5 +1,5 @@
+use cubecl::std::FastDivmod;
 use cubecl::{calculate_cube_count_elemwise, prelude::*, std::tensor::StridedLayout};
-use cubecl_std::FastDivmod;
 
 use crate::{
     CubeRuntime, FloatElement,

crates/burn-cubecl/src/kernel/interpolate/bicubic.rs

@@ -1,5 +1,5 @@
+use cubecl::std::FastDivmod;
 use cubecl::{calculate_cube_count_elemwise, prelude::*, std::tensor::StridedLayout};
-use cubecl_std::FastDivmod;
 
 use crate::{
     CubeRuntime, FloatElement,