From 79ba1c32de7ce9181587b640f0e2e4181490d250 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 09:42:29 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Optimize=20autograd=20backp?= =?UTF-8?q?ropagation=20pass?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refactored `accumulate_grad` to take ownership of Tensors, and replaced `.clone()` in `backward` loop with `.take()` and re-insertion to avoid redundant deep cloning of `Tensor` wrappers. Co-authored-by: teerthsharma <78080953+teerthsharma@users.noreply.github.com> --- .jules/bolt.md | 3 +++ crates/aether-core/src/ml/autograd.rs | 36 +++++++++++++-------------- test_autograd.sh | 2 ++ 3 files changed, 23 insertions(+), 18 deletions(-) create mode 100644 .jules/bolt.md create mode 100644 test_autograd.sh diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 0000000..92e4793 --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,3 @@ +## 2026-07-05 - Optimize Autograd Backpropagation +**Learning:** In reverse-mode backpropagation, cloning `Tensor` metadata for intermediate gradients causes redundant heap allocations. Using `Option::take()` to acquire ownership of gradients during backprop and re-insertion avoids borrow checker issues and unnecessary cloning. +**Action:** Pass `Tensor` by value in `accumulate_grad` to utilize empty `None` slots efficiently. diff --git a/crates/aether-core/src/ml/autograd.rs b/crates/aether-core/src/ml/autograd.rs index db09b19..13fea31 100644 --- a/crates/aether-core/src/ml/autograd.rs +++ b/crates/aether-core/src/ml/autograd.rs @@ -200,52 +200,52 @@ impl<'a> Context<'a> { match op { Op::Add { out, lhs, rhs } => { // Solves borrow checker by cloning Option first - let grad_out = grads[out.index].clone(); - if let Some(grad) = grad_out { + if let Some(grad) = grads[out.index].take() { // dL/d(lhs) += dL/dout * 1 - Self::accumulate_grad(&mut grads, lhs.index, &grad); - Self::accumulate_grad(&mut grads, rhs.index, &grad); + Self::accumulate_grad(&mut grads, lhs.index, grad.clone()); + Self::accumulate_grad(&mut grads, rhs.index, grad.clone()); + grads[out.index] = Some(grad); } } Op::Mul { out, lhs, rhs } => { - let grad_out = grads[out.index].clone(); - if let Some(grad) = grad_out { + if let Some(grad) = grads[out.index].take() { let lhs_val = self.heap.get(*lhs).unwrap(); let rhs_val = self.heap.get(*rhs).unwrap(); // dL/dLhs = grad_out * rhs let d_lhs: Tensor = grad.mul(rhs_val); - Self::accumulate_grad(&mut grads, lhs.index, &d_lhs); + Self::accumulate_grad(&mut grads, lhs.index, d_lhs); // dL/dRhs = grad_out * lhs let d_rhs: Tensor = grad.mul(lhs_val); - Self::accumulate_grad(&mut grads, rhs.index, &d_rhs); + Self::accumulate_grad(&mut grads, rhs.index, d_rhs); + grads[out.index] = Some(grad); } } Op::MatMul { out, lhs, rhs } => { - let grad_out = grads[out.index].clone(); - if let Some(grad) = grad_out { + if let Some(grad) = grads[out.index].take() { let lhs_val = self.heap.get(*lhs).unwrap(); let rhs_val = self.heap.get(*rhs).unwrap(); // C = A @ B // dA = dC @ B^T let d_lhs: Tensor = grad.matmul(&rhs_val.transpose()); - Self::accumulate_grad(&mut grads, lhs.index, &d_lhs); + Self::accumulate_grad(&mut grads, lhs.index, d_lhs); // dB = A^T @ dC let d_rhs: Tensor = lhs_val.transpose().matmul(&grad); - Self::accumulate_grad(&mut grads, rhs.index, &d_rhs); + Self::accumulate_grad(&mut grads, rhs.index, d_rhs); + grads[out.index] = Some(grad); } } Op::ReLU { out, input } => { - let grad_out = grads[out.index].clone(); - if let Some(grad) = grad_out { + if let Some(grad) = grads[out.index].take() { let input_val = self.heap.get(*input).unwrap(); // dL/dx = grad_out * (1 if x > 0 else 0) let mask = input_val.map(|x| if x > 0.0 { 1.0 } else { 0.0 }); let d_input: Tensor = grad.mul(&mask); - Self::accumulate_grad(&mut grads, input.index, &d_input); + Self::accumulate_grad(&mut grads, input.index, d_input); + grads[out.index] = Some(grad); } } } @@ -254,18 +254,18 @@ impl<'a> Context<'a> { grads } - fn accumulate_grad(grads: &mut Vec>, idx: usize, grad: &Tensor) { + fn accumulate_grad(grads: &mut Vec>, idx: usize, grad: Tensor) { if idx >= grads.len() { grads.resize(idx + 1 + 256, None); } match &mut grads[idx] { Some(existing) => { - let new = existing.add(grad); + let new = existing.add(&grad); grads[idx] = Some(new); } None => { - grads[idx] = Some(grad.clone()); + grads[idx] = Some(grad); } } } diff --git a/test_autograd.sh b/test_autograd.sh new file mode 100644 index 0000000..36cf175 --- /dev/null +++ b/test_autograd.sh @@ -0,0 +1,2 @@ +cd crates/aether-core +cargo test -p aether-core --offline -- autograd