From 20ab5e31d7172bcca09e5a6169ac9f614aae63e3 Mon Sep 17 00:00:00 2001
From: Nathaniel Simard <nathaniel.simard.42@gmail.com>
Date: Sat, 21 Sep 2024 13:28:07 -0400
Subject: [PATCH] Chore: Update CubeCL (#2292)

---
 Cargo.lock                                    | 16 ++++++++--------
 Cargo.toml                                    |  4 ++--
 crates/burn-jit/src/fusion/tracing/builder.rs |  5 +++++
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 75794f5926..3c0056a397 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1397,7 +1397,7 @@ dependencies = [
 [[package]]
 name = "cubecl"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "cubecl-core",
  "cubecl-cuda",
@@ -1408,7 +1408,7 @@ dependencies = [
 [[package]]
 name = "cubecl-common"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "derive-new",
  "getrandom",
@@ -1423,7 +1423,7 @@ dependencies = [
 [[package]]
 name = "cubecl-core"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1440,7 +1440,7 @@ dependencies = [
 [[package]]
 name = "cubecl-cuda"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1455,7 +1455,7 @@ dependencies = [
 [[package]]
 name = "cubecl-linalg"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "bytemuck",
  "cubecl-core",
@@ -1466,7 +1466,7 @@ dependencies = [
 [[package]]
 name = "cubecl-macros"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "cubecl-common",
  "darling",
@@ -1481,7 +1481,7 @@ dependencies = [
 [[package]]
 name = "cubecl-runtime"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "async-channel",
  "cfg_aliases 0.2.1",
@@ -1501,7 +1501,7 @@ dependencies = [
 [[package]]
 name = "cubecl-wgpu"
 version = "0.2.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b#0c494b6d6ec3cc28635e0c6d7a4935008c3c8b8b"
+source = "git+https://github.com/tracel-ai/cubecl?rev=447968efed59e8659c255f1d84d82039b533b464#447968efed59e8659c255f1d84d82039b533b464"
 dependencies = [
  "async-channel",
  "bytemuck",
diff --git a/Cargo.toml b/Cargo.toml
index b06407b592..a48b9aba8b 100644
--- 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
index 24fd084441..dc416d9ec9 100644
--- 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 {