613 lines
20 KiB
Rust
613 lines
20 KiB
Rust
//! Backward Pass Implementation for PINN Training
|
|
//!
|
|
//! This module provides analytical gradient computation for the Helmholtz PINN.
|
|
//! Instead of using autograd, we derive and implement closed-form gradients.
|
|
//!
|
|
//! ## Loss Function
|
|
//!
|
|
//! L = w_data * L_data + w_pde * L_pde
|
|
//!
|
|
//! where:
|
|
//! - L_data = MSE(u_pred, u_target) = (1/N) Σ ||u_pred - u_target||²
|
|
//! - L_pde = MSE(∂²u/∂x² + k²u, 0) (Helmholtz residual)
|
|
//!
|
|
//! ## Network Architecture
|
|
//!
|
|
//! u_pred = W_L * tanh(W_{L-1} * ... tanh(W_0 * φ(x, B) + b_0) ... + b_{L-1}) + b_L
|
|
//!
|
|
//! where φ(x, B) = [sin(2πBx), cos(2πBx)] are Fourier features
|
|
|
|
use std::f32::consts::PI;
|
|
use rtx_tensor::{Tensor, Device};
|
|
|
|
/// Result type for backward operations
|
|
pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
|
|
|
|
/// Workspace for gradient computation
|
|
///
|
|
/// Stores cached activations from the forward pass needed for backpropagation,
|
|
/// plus gradient accumulators for all trainable parameters.
|
|
#[derive(Debug)]
|
|
pub struct GradientWorkspace {
|
|
// =========================================================================
|
|
// Cached Activations (from forward pass)
|
|
// =========================================================================
|
|
|
|
/// Pre-activation values: z_l = W_l @ h_{l-1} + b_l
|
|
/// Shape: [batch, hidden_dim] for hidden layers, [batch, 2] for output
|
|
pub z: Vec<Tensor>,
|
|
|
|
/// Post-activation values: h_l = tanh(z_l) for hidden, identity for output
|
|
/// Shape: [batch, hidden_dim] for hidden layers, [batch, 2] for output
|
|
/// Note: h[0] is after first layer, NOT the input features
|
|
pub h: Vec<Tensor>,
|
|
|
|
/// Fourier features: φ(x, B) = [sin(2πBx), cos(2πBx)]
|
|
/// Shape: [batch, ff_dim * 2]
|
|
/// This is the input to the first layer
|
|
pub fourier_features: Tensor,
|
|
|
|
/// Raw x values for Fourier gradient computation
|
|
/// Shape: [batch, 1]
|
|
pub x_input: Tensor,
|
|
|
|
// =========================================================================
|
|
// Gradient Accumulators
|
|
// =========================================================================
|
|
|
|
/// Weight gradients: dL/dW_l
|
|
/// Shape: [out_dim, in_dim] for each layer
|
|
pub dW: Vec<Tensor>,
|
|
|
|
/// Bias gradients: dL/db_l
|
|
/// Shape: [out_dim] for each layer
|
|
pub db: Vec<Tensor>,
|
|
|
|
/// Fourier frequency gradient: dL/dB
|
|
/// Shape: [1, ff_dim]
|
|
pub dB: Tensor,
|
|
|
|
// =========================================================================
|
|
// Spatial Derivative Caches (for PDE loss)
|
|
// =========================================================================
|
|
|
|
/// First spatial derivative: ∂u/∂x
|
|
/// Shape: [batch, 2] (real and imaginary parts)
|
|
pub du_dx: Tensor,
|
|
|
|
/// Second spatial derivative: ∂²u/∂x²
|
|
/// Shape: [batch, 2]
|
|
pub d2u_dx2: Tensor,
|
|
|
|
// =========================================================================
|
|
// Backward Pass Intermediate Buffers (pre-allocated to avoid heap allocs)
|
|
// =========================================================================
|
|
|
|
/// h² for tanh derivative: 1 - h²
|
|
/// Shape: [batch, hidden_dim] for each layer (except output)
|
|
pub h_sq: Vec<Tensor>,
|
|
|
|
/// tanh derivative: 1 - h² = sech²(z)
|
|
/// Shape: [batch, hidden_dim] for each layer (except output)
|
|
pub tanh_deriv: Vec<Tensor>,
|
|
|
|
/// dL/dz = dL/dh * tanh'(z) for each layer
|
|
/// Shape: [batch, out_dim] for each layer
|
|
pub dL_dz: Vec<Tensor>,
|
|
|
|
/// dL/dz transposed for weight gradient computation
|
|
/// Shape: [out_dim, batch] for each layer
|
|
pub dL_dz_t: Vec<Tensor>,
|
|
|
|
/// Gradient w.r.t. previous layer output (passed backward)
|
|
/// Shape: [batch, hidden_dim] for hidden, [batch, ff_dim*2] for first layer
|
|
pub dL_dh_prev: Vec<Tensor>,
|
|
|
|
/// MSE gradient output buffer
|
|
/// Shape: [batch, 2]
|
|
pub d_loss: Tensor,
|
|
|
|
// =========================================================================
|
|
// Fourier Backward Intermediate Buffers
|
|
// =========================================================================
|
|
|
|
/// x * scale for Fourier backward
|
|
/// Shape: [batch, 1]
|
|
pub x_scaled: Tensor,
|
|
|
|
/// B @ x_scaled = 2πBx
|
|
/// Shape: [batch, ff_dim]
|
|
pub Bx: Tensor,
|
|
|
|
/// cos(2πBx) for Fourier backward
|
|
/// Shape: [batch, ff_dim]
|
|
pub cos_Bx: Tensor,
|
|
|
|
/// sin(2πBx) for Fourier backward
|
|
/// Shape: [batch, ff_dim]
|
|
pub sin_Bx: Tensor,
|
|
|
|
/// dphi_sin/dB = 2πx * cos(2πBx)
|
|
/// Shape: [batch, ff_dim]
|
|
pub dphi_sin_dB: Tensor,
|
|
|
|
/// dphi_cos/dB = -2πx * sin(2πBx)
|
|
/// Shape: [batch, ff_dim]
|
|
pub dphi_cos_dB: Tensor,
|
|
|
|
/// Slice of dL_dphi for sin part
|
|
/// Shape: [batch, ff_dim]
|
|
pub dL_dphi_sin: Tensor,
|
|
|
|
/// Slice of dL_dphi for cos part
|
|
/// Shape: [batch, ff_dim]
|
|
pub dL_dphi_cos: Tensor,
|
|
|
|
/// dL_dB from sin contribution
|
|
/// Shape: [batch, ff_dim]
|
|
pub dL_dB_from_sin: Tensor,
|
|
|
|
/// dL_dB from cos contribution
|
|
/// Shape: [batch, ff_dim]
|
|
pub dL_dB_from_cos: Tensor,
|
|
|
|
/// Device for tensor allocation
|
|
device: Device,
|
|
|
|
/// Number of layers
|
|
num_layers: usize,
|
|
|
|
/// Batch size
|
|
batch_size: usize,
|
|
|
|
/// Hidden dimension
|
|
hidden_dim: usize,
|
|
|
|
/// Fourier feature dimension
|
|
ff_dim: usize,
|
|
}
|
|
|
|
impl GradientWorkspace {
|
|
/// Create a new gradient workspace
|
|
///
|
|
/// # Arguments
|
|
/// * `batch_size` - Number of samples in a batch
|
|
/// * `ff_dim` - Fourier feature dimension
|
|
/// * `hidden_dim` - Hidden layer dimension
|
|
/// * `num_hidden_layers` - Number of hidden layers (not including output)
|
|
/// * `device` - Device for tensor allocation
|
|
pub fn new(
|
|
batch_size: usize,
|
|
ff_dim: usize,
|
|
hidden_dim: usize,
|
|
num_hidden_layers: usize,
|
|
device: &Device,
|
|
) -> Result<Self> {
|
|
let num_layers = num_hidden_layers + 1; // +1 for output layer
|
|
|
|
// Pre-allocate activation caches
|
|
let mut z = Vec::with_capacity(num_layers);
|
|
let mut h = Vec::with_capacity(num_layers);
|
|
|
|
// Hidden layers
|
|
for _ in 0..num_hidden_layers {
|
|
z.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
|
|
h.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
|
|
}
|
|
// Output layer
|
|
z.push(Tensor::zeros(&[batch_size, 2], device)?);
|
|
h.push(Tensor::zeros(&[batch_size, 2], device)?);
|
|
|
|
// Pre-allocate gradient accumulators
|
|
let mut dW = Vec::with_capacity(num_layers);
|
|
let mut db = Vec::with_capacity(num_layers);
|
|
|
|
// First hidden layer: [hidden_dim, ff_dim * 2]
|
|
dW.push(Tensor::zeros(&[hidden_dim, ff_dim * 2], device)?);
|
|
db.push(Tensor::zeros(&[hidden_dim], device)?);
|
|
|
|
// Middle hidden layers: [hidden_dim, hidden_dim]
|
|
for _ in 1..num_hidden_layers {
|
|
dW.push(Tensor::zeros(&[hidden_dim, hidden_dim], device)?);
|
|
db.push(Tensor::zeros(&[hidden_dim], device)?);
|
|
}
|
|
|
|
// Output layer: [2, hidden_dim]
|
|
dW.push(Tensor::zeros(&[2, hidden_dim], device)?);
|
|
db.push(Tensor::zeros(&[2], device)?);
|
|
|
|
// Fourier frequency gradient
|
|
let dB = Tensor::zeros(&[1, ff_dim], device)?;
|
|
|
|
// Fourier features cache
|
|
let fourier_features = Tensor::zeros(&[batch_size, ff_dim * 2], device)?;
|
|
|
|
// Input cache
|
|
let x_input = Tensor::zeros(&[batch_size, 1], device)?;
|
|
|
|
// Spatial derivatives
|
|
let du_dx = Tensor::zeros(&[batch_size, 2], device)?;
|
|
let d2u_dx2 = Tensor::zeros(&[batch_size, 2], device)?;
|
|
|
|
// =========================================================================
|
|
// Backward Pass Intermediate Buffers
|
|
// =========================================================================
|
|
|
|
// Buffers for tanh derivative computation (only for hidden layers)
|
|
let mut h_sq = Vec::with_capacity(num_hidden_layers);
|
|
let mut tanh_deriv = Vec::with_capacity(num_hidden_layers);
|
|
for _ in 0..num_hidden_layers {
|
|
h_sq.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
|
|
tanh_deriv.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
|
|
}
|
|
|
|
// dL/dz for each layer
|
|
let mut dL_dz = Vec::with_capacity(num_layers);
|
|
let mut dL_dz_t = Vec::with_capacity(num_layers);
|
|
for i in 0..num_layers {
|
|
let out_dim = if i < num_hidden_layers { hidden_dim } else { 2 };
|
|
dL_dz.push(Tensor::zeros(&[batch_size, out_dim], device)?);
|
|
dL_dz_t.push(Tensor::zeros(&[out_dim, batch_size], device)?);
|
|
}
|
|
|
|
// dL/dh_prev for each layer (gradient passed backward)
|
|
let mut dL_dh_prev = Vec::with_capacity(num_layers);
|
|
// First layer: gradient goes to Fourier features [batch, ff_dim*2]
|
|
dL_dh_prev.push(Tensor::zeros(&[batch_size, ff_dim * 2], device)?);
|
|
// Hidden layers: [batch, hidden_dim]
|
|
for _ in 1..num_layers {
|
|
dL_dh_prev.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
|
|
}
|
|
|
|
// MSE gradient buffer
|
|
let d_loss = Tensor::zeros(&[batch_size, 2], device)?;
|
|
|
|
// =========================================================================
|
|
// Fourier Backward Intermediate Buffers
|
|
// =========================================================================
|
|
|
|
let x_scaled = Tensor::zeros(&[batch_size, 1], device)?;
|
|
let Bx = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let cos_Bx = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let sin_Bx = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let dphi_sin_dB = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let dphi_cos_dB = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let dL_dphi_sin = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let dL_dphi_cos = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let dL_dB_from_sin = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
let dL_dB_from_cos = Tensor::zeros(&[batch_size, ff_dim], device)?;
|
|
|
|
Ok(Self {
|
|
z,
|
|
h,
|
|
fourier_features,
|
|
x_input,
|
|
dW,
|
|
db,
|
|
dB,
|
|
du_dx,
|
|
d2u_dx2,
|
|
// Backward intermediates
|
|
h_sq,
|
|
tanh_deriv,
|
|
dL_dz,
|
|
dL_dz_t,
|
|
dL_dh_prev,
|
|
d_loss,
|
|
// Fourier backward intermediates
|
|
x_scaled,
|
|
Bx,
|
|
cos_Bx,
|
|
sin_Bx,
|
|
dphi_sin_dB,
|
|
dphi_cos_dB,
|
|
dL_dphi_sin,
|
|
dL_dphi_cos,
|
|
dL_dB_from_sin,
|
|
dL_dB_from_cos,
|
|
// Metadata
|
|
device: device.clone(),
|
|
num_layers,
|
|
batch_size,
|
|
hidden_dim,
|
|
ff_dim,
|
|
})
|
|
}
|
|
|
|
/// Get the number of layers
|
|
pub fn num_layers(&self) -> usize {
|
|
self.num_layers
|
|
}
|
|
|
|
/// Zero all gradients (call before backward pass)
|
|
pub fn zero_grad(&mut self) -> Result<()> {
|
|
for dw in &mut self.dW {
|
|
*dw = Tensor::zeros(dw.shape(), &self.device)?;
|
|
}
|
|
for db in &mut self.db {
|
|
*db = Tensor::zeros(db.shape(), &self.device)?;
|
|
}
|
|
self.dB = Tensor::zeros(self.dB.shape(), &self.device)?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Compute the gradient of a single layer's backward pass
|
|
///
|
|
/// Given dL/dh (gradient w.r.t. layer output), computes:
|
|
/// - dL/dW (weight gradient)
|
|
/// - dL/db (bias gradient)
|
|
/// - dL/dh_prev (gradient to propagate backward)
|
|
///
|
|
/// For a layer with h = tanh(W @ h_prev + b):
|
|
/// - dL/dz = dL/dh * (1 - h²) (tanh derivative)
|
|
/// - dL/dW = dL/dz.T @ h_prev
|
|
/// - dL/db = sum(dL/dz, axis=0)
|
|
/// - dL/dh_prev = dL/dz @ W
|
|
///
|
|
/// # Arguments
|
|
/// * `dL_dh` - Gradient w.r.t. layer output [batch, out_dim]
|
|
/// * `h` - Cached activation (tanh output) [batch, out_dim]
|
|
/// * `h_prev` - Previous layer activation (input to this layer) [batch, in_dim]
|
|
/// * `W` - Layer weights [out_dim, in_dim]
|
|
/// * `apply_tanh_grad` - Whether to apply tanh derivative (false for output layer)
|
|
///
|
|
/// # Returns
|
|
/// Tuple of (dL/dW, dL/db, dL/dh_prev)
|
|
pub fn layer_backward(
|
|
dL_dh: &Tensor,
|
|
h: &Tensor,
|
|
h_prev: &Tensor,
|
|
W: &Tensor,
|
|
apply_tanh_grad: bool,
|
|
) -> Result<(Tensor, Tensor, Tensor)> {
|
|
// Compute dL/dz
|
|
let dL_dz = if apply_tanh_grad {
|
|
// tanh'(z) = 1 - tanh²(z) = 1 - h²
|
|
let h_sq = h.mul(h)?;
|
|
let one_minus_h_sq = h_sq.mul_scalar(-1.0)?.add_scalar(1.0)?;
|
|
dL_dh.mul(&one_minus_h_sq)?
|
|
} else {
|
|
// No activation (output layer)
|
|
dL_dh.clone()
|
|
};
|
|
|
|
// dL/dW = dL/dz.T @ h_prev = [out_dim, batch] @ [batch, in_dim] = [out_dim, in_dim]
|
|
let dL_dz_t = dL_dz.transpose(0, 1)?;
|
|
let dL_dW = dL_dz_t.matmul(h_prev)?;
|
|
|
|
// dL/db = sum(dL/dz, axis=0) = [out_dim]
|
|
let dL_db = dL_dz.sum(Some(0))?;
|
|
|
|
// dL/dh_prev = dL/dz @ W = [batch, out_dim] @ [out_dim, in_dim] = [batch, in_dim]
|
|
let dL_dh_prev = dL_dz.matmul(W)?;
|
|
|
|
Ok((dL_dW, dL_db, dL_dh_prev))
|
|
}
|
|
|
|
/// Compute gradient of Fourier features w.r.t. B
|
|
///
|
|
/// φ(x, B) = [sin(2πBx), cos(2πBx)]
|
|
/// ∂φ/∂B = [2πx * cos(2πBx), -2πx * sin(2πBx)]
|
|
///
|
|
/// # Arguments
|
|
/// * `dL_dphi` - Gradient w.r.t. Fourier features [batch, ff_dim * 2]
|
|
/// * `x` - Input positions [batch, 1]
|
|
/// * `B` - Fourier frequencies [1, ff_dim]
|
|
///
|
|
/// # Returns
|
|
/// dL/dB with shape [1, ff_dim]
|
|
pub fn fourier_backward(
|
|
dL_dphi: &Tensor,
|
|
x: &Tensor,
|
|
B: &Tensor,
|
|
) -> Result<Tensor> {
|
|
let scale = 2.0 * PI;
|
|
let ff_dim = B.shape()[1] as usize;
|
|
|
|
// Compute 2πBx
|
|
let x_scaled = x.mul_scalar(scale)?; // [batch, 1]
|
|
let Bx = x_scaled.matmul(B)?; // [batch, ff_dim]
|
|
|
|
// ∂sin(2πBx)/∂B = 2πx * cos(2πBx)
|
|
// ∂cos(2πBx)/∂B = -2πx * sin(2πBx)
|
|
let cos_Bx = Bx.cos()?;
|
|
let sin_Bx = Bx.sin()?;
|
|
|
|
// dphi/dB for sin part: 2πx * cos(2πBx)
|
|
let dphi_sin_dB = x_scaled.mul(&cos_Bx)?; // Broadcasting: [batch, 1] * [batch, ff_dim]
|
|
|
|
// dphi/dB for cos part: -2πx * sin(2πBx)
|
|
let dphi_cos_dB = x_scaled.mul(&sin_Bx)?.mul_scalar(-1.0)?;
|
|
|
|
// Split dL_dphi into sin and cos parts
|
|
// dL_dphi has shape [batch, ff_dim * 2] = [batch, sin_part | cos_part]
|
|
// Use slice(dim, start, end) API
|
|
let dL_dphi_sin = dL_dphi.slice(1, 0, ff_dim)?;
|
|
let dL_dphi_cos = dL_dphi.slice(1, ff_dim, ff_dim * 2)?;
|
|
|
|
// Chain rule: dL/dB = dL/dphi_sin * dphi_sin/dB + dL/dphi_cos * dphi_cos/dB
|
|
let dL_dB_from_sin = dL_dphi_sin.mul(&dphi_sin_dB)?;
|
|
let dL_dB_from_cos = dL_dphi_cos.mul(&dphi_cos_dB)?;
|
|
|
|
// Sum over batch dimension to get final gradient
|
|
// sum(dim) where dim=0 for batch dimension
|
|
let dL_dB = dL_dB_from_sin.add(&dL_dB_from_cos)?.sum(Some(0))?;
|
|
|
|
Ok(dL_dB)
|
|
}
|
|
|
|
/// Compute MSE loss gradient
|
|
///
|
|
/// L = (1/N) Σ ||pred - target||²
|
|
/// dL/dpred = (2/N) * (pred - target)
|
|
///
|
|
/// # Arguments
|
|
/// * `pred` - Predictions [batch, 2]
|
|
/// * `target` - Targets [batch, 2]
|
|
///
|
|
/// # Returns
|
|
/// dL/dpred with shape [batch, 2]
|
|
pub fn mse_backward(pred: &Tensor, target: &Tensor) -> Result<Tensor> {
|
|
let batch_size = pred.shape()[0] as f32;
|
|
let diff = pred.sub(target)?;
|
|
let grad = diff.mul_scalar(2.0 / batch_size)?;
|
|
Ok(grad)
|
|
}
|
|
|
|
// =============================================================================
|
|
// GPU-NATIVE ZERO-ALLOCATION BACKWARD PASS
|
|
// =============================================================================
|
|
//
|
|
// These functions use PinnStreamContext GPU ops to avoid tensor allocations.
|
|
// All intermediate results are written to pre-allocated GradientWorkspace buffers.
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use crate::cuda_stream_context::PinnStreamContext;
|
|
|
|
/// Compute MSE loss gradient into pre-allocated buffer (zero-alloc version)
|
|
///
|
|
/// # Arguments
|
|
/// * `pred` - Predictions [batch, 2]
|
|
/// * `target` - Targets [batch, 2]
|
|
/// * `ctx` - CUDA stream context for GPU ops
|
|
/// * `d_loss` - Pre-allocated output buffer [batch, 2]
|
|
#[cfg(feature = "cuda")]
|
|
pub fn mse_backward_inplace(
|
|
pred: &Tensor,
|
|
target: &Tensor,
|
|
ctx: &PinnStreamContext,
|
|
d_loss: &mut Tensor,
|
|
) -> Result<()> {
|
|
let batch_size = pred.shape()[0] as f32;
|
|
let scale = 2.0 / batch_size;
|
|
|
|
// d_loss = pred - target
|
|
ctx.sub_out(pred, target, d_loss)?;
|
|
|
|
// d_loss = d_loss * scale (in-place via copy to same buffer)
|
|
// Note: mul_scalar_out writes to output, we need in-place
|
|
// For now, use existing mul_scalar_ if available
|
|
d_loss.mul_scalar_(scale)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Compute single layer backward pass using pre-allocated buffers (zero-alloc version)
|
|
///
|
|
/// # Arguments
|
|
/// * `dL_dh` - Gradient w.r.t. layer output [batch, out_dim]
|
|
/// * `h` - Cached activation (tanh output) [batch, out_dim]
|
|
/// * `h_prev` - Previous layer activation [batch, in_dim]
|
|
/// * `W` - Layer weights [out_dim, in_dim]
|
|
/// * `layer_idx` - Index of this layer (for accessing workspace buffers)
|
|
/// * `apply_tanh_grad` - Whether to apply tanh derivative
|
|
/// * `ctx` - CUDA stream context for GPU ops
|
|
/// * `ws` - GradientWorkspace with pre-allocated buffers
|
|
#[cfg(feature = "cuda")]
|
|
pub fn layer_backward_inplace(
|
|
dL_dh: &Tensor,
|
|
h: &Tensor,
|
|
h_prev: &Tensor,
|
|
W: &Tensor,
|
|
layer_idx: usize,
|
|
apply_tanh_grad: bool,
|
|
ctx: &PinnStreamContext,
|
|
ws: &mut GradientWorkspace,
|
|
) -> Result<()> {
|
|
// Compute dL/dz into ws.dL_dz[layer_idx]
|
|
if apply_tanh_grad {
|
|
// dL/dz = dL/dh * (1 - h²)
|
|
// Use fused tanh_deriv_mul kernel: out = (1 - h*h) * dL_dh
|
|
ctx.tanh_deriv_mul_out(h, dL_dh, &mut ws.dL_dz[layer_idx])?;
|
|
} else {
|
|
// No activation - copy dL_dh to dL_dz
|
|
// For output layer, dL_dz = dL_dh directly
|
|
// We can use add with zero, or just assign (need copy)
|
|
// Use a simple copy via mul_scalar(1.0)
|
|
ctx.mul_scalar_out(dL_dh, 1.0, &mut ws.dL_dz[layer_idx])?;
|
|
}
|
|
|
|
// dL/dW = dL/dz.T @ h_prev
|
|
// First transpose dL_dz: [batch, out] -> [out, batch]
|
|
// Then matmul: [out, batch] @ [batch, in] -> [out, in]
|
|
//
|
|
// For matmul, we use the Tensor API since we have matmul_out
|
|
let dL_dz_t = ws.dL_dz[layer_idx].transpose(0, 1)?;
|
|
dL_dz_t.matmul_out(h_prev, &mut ws.dW[layer_idx])?;
|
|
|
|
// dL/db = sum(dL/dz, axis=0)
|
|
ctx.sum_rows_out(&ws.dL_dz[layer_idx], &mut ws.db[layer_idx])?;
|
|
|
|
// dL/dh_prev = dL/dz @ W
|
|
ws.dL_dz[layer_idx].matmul_out(W, &mut ws.dL_dh_prev[layer_idx])?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Compute Fourier backward gradient using fused GPU kernel (zero-alloc version)
|
|
///
|
|
/// This replaces the allocating `fourier_backward()` with a single fused CUDA kernel
|
|
/// that eliminates ~10 D2H copies per training step.
|
|
///
|
|
/// # Arguments
|
|
/// * `dL_dphi` - Gradient from layer 0 [batch, 2*ff_dim]
|
|
/// * `x` - Input coordinates [batch, 1]
|
|
/// * `B` - Fourier frequencies [1, ff_dim]
|
|
/// * `ctx` - CUDA stream context for GPU ops
|
|
/// * `dL_dB` - Pre-allocated output buffer [1, ff_dim]
|
|
#[cfg(feature = "cuda")]
|
|
pub fn fourier_backward_inplace(
|
|
dL_dphi: &Tensor,
|
|
x: &Tensor,
|
|
B: &Tensor,
|
|
ctx: &PinnStreamContext,
|
|
dL_dB: &mut Tensor,
|
|
) -> Result<()> {
|
|
use std::f32::consts::PI;
|
|
|
|
// Zero the output buffer before atomic accumulation
|
|
ctx.memset_zero_out(dL_dB)?;
|
|
|
|
// Launch fused kernel: computes entire Fourier gradient in one pass
|
|
ctx.fourier_grad(x, dL_dphi, B, dL_dB, 2.0 * PI)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_gradient_workspace_creation() {
|
|
let device = Device::Cpu;
|
|
let ws = GradientWorkspace::new(32, 16, 64, 3, &device).unwrap();
|
|
|
|
// Check shapes
|
|
assert_eq!(ws.z.len(), 4); // 3 hidden + 1 output
|
|
assert_eq!(ws.h.len(), 4);
|
|
assert_eq!(ws.dW.len(), 4);
|
|
assert_eq!(ws.db.len(), 4);
|
|
|
|
// Check dimensions
|
|
assert_eq!(ws.dW[0].shape(), &[64, 32]); // hidden_dim x ff_dim*2
|
|
assert_eq!(ws.dW[1].shape(), &[64, 64]); // hidden x hidden
|
|
assert_eq!(ws.dW[3].shape(), &[2, 64]); // output x hidden
|
|
}
|
|
|
|
#[test]
|
|
fn test_mse_backward() {
|
|
let device = Device::Cpu;
|
|
let pred = Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap();
|
|
let target = Tensor::from_vec(vec![0.0f32, 0.0, 0.0, 0.0], &[2, 2], &device).unwrap();
|
|
|
|
let grad = mse_backward(&pred, &target).unwrap();
|
|
|
|
// dL/dpred = (2/2) * (pred - 0) = pred
|
|
let grad_cpu = grad.to_cpu().unwrap();
|
|
assert!((grad_cpu[0] - 1.0).abs() < 1e-5);
|
|
assert!((grad_cpu[1] - 2.0).abs() < 1e-5);
|
|
}
|
|
}
|