Update cubecl (#2680)

Cargo.toml

@@ -153,8 +153,8 @@ ahash = { version = "0.8.11", default-features = false }
 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 = "5a3f9ac9f6178c4f76570535bf5e42ef12a19a3d" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "5a3f9ac9f6178c4f76570535bf5e42ef12a19a3d" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "8244dbb4660e373ff1ffb780feb73a5b899e5977" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "8244dbb4660e373ff1ffb780feb73a5b899e5977" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

crates/burn-jit/src/kernel/conv/deform_conv_transpose2d.rs

@@ -509,7 +509,7 @@ fn deform_col2img_kernel<F: Float>(
     offset: &Tensor<F>,
     mask: &Tensor<F>,
     columns: &Tensor<F>,
-    grad_input: &mut Tensor<AtomicU32>,
+    grad_input: &mut Tensor<Atomic<u32>>,
     args: &DeformConv2dCol2ImgArgs,
     #[comptime] use_mask: bool,
 ) {
@@ -589,14 +589,14 @@ fn deform_col2img_kernel<F: Float>(
 }
 
 #[cube]
-fn float_atomic_add(ptr: &mut AtomicU32, value: f32) {
+fn float_atomic_add(ptr: &mut Atomic<u32>, value: f32) {
     if value != 0.0 {
-        let mut v = AtomicU32::load(ptr);
+        let mut v = Atomic::<u32>::load(ptr);
         loop {
             let prev = v;
             let v_float = f32::bitcast_from(v);
             let new = u32::bitcast_from(v_float + value);
-            v = AtomicU32::compare_and_swap(ptr, v, new);
+            v = Atomic::<u32>::compare_and_swap(ptr, v, new);
             if prev == v {
                 break;
             }