Chore: Update CubeCL (#2292)

tracel-ai · Sep 21, 2024 · 20ab5e3 · 20ab5e3
1 parent aa79e36
commit 20ab5e3
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 10 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -151,8 +151,8 @@ tch = "0.15.0"
 portable-atomic-util = { version = "0.2.2", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "447968efed59e8659c255f1d84d82039b533b464" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "447968efed59e8659c255f1d84d82039b533b464" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl" }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common" }

diff --git a/crates/burn-jit/src/fusion/tracing/builder.rs b/crates/burn-jit/src/fusion/tracing/builder.rs
@@ -514,6 +514,11 @@ impl TraceBuilder {
                     Operator::AtomicCompareAndSwap(_op) => {
                         // Nothing to do.
                     }
+                    Operator::Magnitude(op) => mark_unary(
+                        op,
+                        &mut local_tensor_ids_input,
+                        &mut local_tensor_ids_output,
+                    ),
                 },
                 Operation::Procedure(proc) => {
                     match proc {