2174 lines
82 KiB
Rust
2174 lines
82 KiB
Rust
//! Physics-Informed Neural Network (PINN) for 1D MRE Helmholtz Equation - Library Interface
|
||
//!
|
||
//! This module exposes the core types and functions for benchmarking.
|
||
|
||
pub mod uber_kernel;
|
||
pub mod cuda_stream_context;
|
||
pub mod pinn_graph;
|
||
pub mod backward;
|
||
pub mod gpu_adam;
|
||
#[cfg(feature = "cuda")]
|
||
pub mod unsafe_graph;
|
||
|
||
// Re-exports for convenience
|
||
#[cfg(feature = "cuda")]
|
||
pub use cuda_stream_context::PinnStreamContext;
|
||
#[cfg(feature = "cuda")]
|
||
pub use cuda_stream_context::CachedGpuPtrs;
|
||
#[cfg(feature = "cuda")]
|
||
pub use pinn_graph::PinnGraph;
|
||
#[cfg(feature = "cuda")]
|
||
pub use unsafe_graph::UnsafeGraph;
|
||
pub use backward::{GradientWorkspace, layer_backward, fourier_backward, mse_backward};
|
||
#[cfg(feature = "cuda")]
|
||
pub use backward::{layer_backward_inplace, mse_backward_inplace, fourier_backward_inplace};
|
||
pub use gpu_adam::GpuAdam;
|
||
|
||
use anyhow::Result;
|
||
use num_complex::Complex64;
|
||
use std::f64::consts::PI;
|
||
use std::collections::HashMap;
|
||
|
||
// RustyTorch++ imports
|
||
use rtx_tensor::{Tensor, Device};
|
||
use rtx_autograd::clear_tape;
|
||
use rtx_nn::Linear;
|
||
#[cfg(feature = "cuda")]
|
||
use std::sync::Arc;
|
||
#[cfg(feature = "cuda")]
|
||
use cudarc::driver::safe::CudaStream;
|
||
use rtx_nn::layers::Module;
|
||
use rtx_transformers::optimizers::{AdamOptimizer, AdamConfig};
|
||
use rtx_transformers::schedulers::{ReduceLROnPlateauScheduler, PlateauMode, ThresholdMode};
|
||
|
||
// =============================================================================
|
||
// CONFIGURATION
|
||
// =============================================================================
|
||
|
||
/// Configuration for physics, network, and training parameters
|
||
#[derive(Debug, Clone)]
|
||
pub struct Config {
|
||
// Physics parameters
|
||
pub rho: f64,
|
||
pub freq: f64,
|
||
pub l: f64,
|
||
pub u0: f64,
|
||
pub g_prime_true: f64,
|
||
pub g_double_true: f64,
|
||
|
||
// Grid parameters
|
||
pub n_data: usize,
|
||
pub n_pde: usize,
|
||
|
||
// Network architecture (LFFN-MLP)
|
||
pub u_layers: usize,
|
||
pub u_hidden: usize,
|
||
pub u_ff_dim: usize,
|
||
pub u_ff_scale: f64,
|
||
|
||
// Training parameters
|
||
pub lr: f64,
|
||
pub epochs: usize,
|
||
pub scheduler_patience: usize,
|
||
pub scheduler_factor: f64,
|
||
|
||
// Loss weights
|
||
pub data_weight: f64,
|
||
pub pde_weight: f64,
|
||
|
||
// Logging
|
||
pub print_every: usize,
|
||
}
|
||
|
||
impl Default for Config {
|
||
fn default() -> Self {
|
||
Self {
|
||
rho: 1040.0,
|
||
freq: 50.0,
|
||
l: 0.1,
|
||
u0: 1e-6,
|
||
g_prime_true: 3000.0,
|
||
g_double_true: 1500.0,
|
||
n_data: 200,
|
||
n_pde: 200,
|
||
u_layers: 4,
|
||
u_hidden: 64,
|
||
u_ff_dim: 64,
|
||
u_ff_scale: 10.0,
|
||
lr: 1e-3,
|
||
epochs: 50_000,
|
||
scheduler_patience: 500,
|
||
scheduler_factor: 0.5,
|
||
data_weight: 1.0,
|
||
pde_weight: 1e-6,
|
||
print_every: 500,
|
||
}
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// SYNTHETIC DATA GENERATION
|
||
// =============================================================================
|
||
|
||
/// Calculate the complex wave number k = sqrt(ρω²/G*)
|
||
pub fn calculate_k(cfg: &Config) -> Complex64 {
|
||
let omega = 2.0 * PI * cfg.freq;
|
||
let g_complex = Complex64::new(cfg.g_prime_true, cfg.g_double_true);
|
||
let k_sq = cfg.rho * omega.powi(2) / g_complex;
|
||
let mut k = k_sq.sqrt();
|
||
if k.im < 0.0 {
|
||
k = -k;
|
||
}
|
||
k
|
||
}
|
||
|
||
/// Generate synthetic displacement data: u(x) = U0 * exp(i*k*x)
|
||
pub fn synthesize_displacement(cfg: &Config) -> (Vec<f64>, Vec<f64>, Vec<f64>, Complex64) {
|
||
let k = calculate_k(cfg);
|
||
let x: Vec<f64> = (0..cfg.n_data)
|
||
.map(|i| i as f64 * cfg.l / (cfg.n_data - 1) as f64)
|
||
.collect();
|
||
let u_complex: Vec<Complex64> = x.iter()
|
||
.map(|&xi| cfg.u0 * (Complex64::i() * k * xi).exp())
|
||
.collect();
|
||
let u_real: Vec<f64> = u_complex.iter().map(|u| u.re).collect();
|
||
let u_imag: Vec<f64> = u_complex.iter().map(|u| u.im).collect();
|
||
(x, u_real, u_imag, k)
|
||
}
|
||
|
||
/// Generate exact analytical derivatives for validation
|
||
pub fn synthesize_exact_derivatives(
|
||
cfg: &Config,
|
||
k: Complex64,
|
||
) -> (Vec<[f64; 2]>, Vec<[f64; 2]>) {
|
||
let (_x, u_r, u_i, _) = synthesize_displacement(cfg);
|
||
let u_complex: Vec<Complex64> = u_r.iter().zip(u_i.iter())
|
||
.map(|(&r, &i)| Complex64::new(r, i))
|
||
.collect();
|
||
let dudx: Vec<[f64; 2]> = u_complex.iter()
|
||
.map(|&u| {
|
||
let d = Complex64::i() * k * u;
|
||
[d.re, d.im]
|
||
})
|
||
.collect();
|
||
let d2udx2: Vec<[f64; 2]> = u_complex.iter()
|
||
.map(|&u| {
|
||
let d = -(k * k) * u;
|
||
[d.re, d.im]
|
||
})
|
||
.collect();
|
||
(dudx, d2udx2)
|
||
}
|
||
|
||
// =============================================================================
|
||
// FORWARD WORKSPACE - Pre-allocated buffers for zero-allocation forward pass
|
||
// =============================================================================
|
||
|
||
/// Pre-allocated buffers for forward pass (zero allocation after initialization)
|
||
///
|
||
/// This workspace holds all intermediate tensors needed during the forward pass,
|
||
/// allowing the PINN to reuse memory across iterations instead of allocating new
|
||
/// tensors on every forward call.
|
||
///
|
||
/// Phase 4 optimizations: Uses fused Fourier features kernel that combines
|
||
/// matmul + scale + sin/cos + concat into a single CUDA kernel launch.
|
||
#[derive(Debug)]
|
||
pub struct ForwardWorkspace {
|
||
/// Fourier features buffer: [batch, ff_dim * 2] = [sin(x @ B * scale), cos(x @ B * scale)]
|
||
/// This is the output of the fused Fourier features kernel
|
||
pub features: Tensor,
|
||
/// Hidden layer buffers: [batch, hidden_dim] × num_layers
|
||
pub hidden_layers: Vec<Tensor>,
|
||
/// Final output buffer: [batch, 2] for u_real and u_imag
|
||
pub output: Tensor,
|
||
}
|
||
|
||
impl ForwardWorkspace {
|
||
/// Create a new workspace with pre-allocated buffers
|
||
pub fn new(batch_size: usize, cfg: &Config, device: &Device) -> Result<Self> {
|
||
// Fused Fourier features output: [batch, ff_dim * 2]
|
||
// Contains [sin(x @ B * scale), cos(x @ B * scale)] concatenated
|
||
let features = Tensor::zeros(&[batch_size, cfg.u_ff_dim * 2], device)?;
|
||
|
||
// Hidden layer buffers
|
||
let mut hidden_layers = Vec::with_capacity(cfg.u_layers + 1);
|
||
// First hidden layer output: [batch, hidden_dim]
|
||
hidden_layers.push(Tensor::zeros(&[batch_size, cfg.u_hidden], device)?);
|
||
// Subsequent hidden layers: [batch, hidden_dim]
|
||
for _ in 1..cfg.u_layers {
|
||
hidden_layers.push(Tensor::zeros(&[batch_size, cfg.u_hidden], device)?);
|
||
}
|
||
|
||
// Final output: [batch, 2]
|
||
let output = Tensor::zeros(&[batch_size, 2], device)?;
|
||
|
||
Ok(Self {
|
||
features,
|
||
hidden_layers,
|
||
output,
|
||
})
|
||
}
|
||
|
||
/// Create a new workspace with all tensors allocated on a specific CUDA stream.
|
||
///
|
||
/// This method enables CUDA graph capture by ensuring all workspace tensors
|
||
/// are allocated on the same stream. CUDA graphs require all operations to
|
||
/// use the same stream to avoid cross-stream dependency errors.
|
||
///
|
||
/// # Arguments
|
||
/// * `batch_size` - Number of samples per batch
|
||
/// * `cfg` - Network configuration
|
||
/// * `stream` - The CUDA stream to allocate tensors on
|
||
///
|
||
/// # Example
|
||
/// ```ignore
|
||
/// let ctx = PinnStreamContext::new(0)?;
|
||
/// let ws = ForwardWorkspace::new_on_stream(200, &cfg, ctx.stream_for_alloc())?;
|
||
///
|
||
/// // Now graph capture will work without cross-stream errors
|
||
/// graph.capture(&[200, 1], || {
|
||
/// solver.forward_on_stream(&x, &mut ws, &ctx)
|
||
/// })?;
|
||
/// ```
|
||
#[cfg(feature = "cuda")]
|
||
pub fn new_on_stream(
|
||
batch_size: usize,
|
||
cfg: &Config,
|
||
stream: &Arc<CudaStream>,
|
||
) -> Result<Self> {
|
||
// Fused Fourier features output: [batch, ff_dim * 2]
|
||
// Contains [sin(x @ B * scale), cos(x @ B * scale)] concatenated
|
||
let features = Tensor::zeros_on_stream(&[batch_size, cfg.u_ff_dim * 2], stream)?;
|
||
|
||
// Hidden layer buffers
|
||
let mut hidden_layers = Vec::with_capacity(cfg.u_layers + 1);
|
||
// First hidden layer output: [batch, hidden_dim]
|
||
hidden_layers.push(Tensor::zeros_on_stream(&[batch_size, cfg.u_hidden], stream)?);
|
||
// Subsequent hidden layers: [batch, hidden_dim]
|
||
for _ in 1..cfg.u_layers {
|
||
hidden_layers.push(Tensor::zeros_on_stream(&[batch_size, cfg.u_hidden], stream)?);
|
||
}
|
||
|
||
// Final output: [batch, 2]
|
||
let output = Tensor::zeros_on_stream(&[batch_size, 2], stream)?;
|
||
|
||
Ok(Self {
|
||
features,
|
||
hidden_layers,
|
||
output,
|
||
})
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// CUDA GRAPH CAPTURE FOR FORWARD PASS
|
||
// =============================================================================
|
||
|
||
/// Captured CUDA graph for forward pass acceleration.
|
||
///
|
||
/// CUDA graphs eliminate kernel launch overhead by recording a sequence of
|
||
/// GPU operations and replaying them with a single API call. For the PINN
|
||
/// forward pass (11 kernel launches), this can provide significant speedup
|
||
/// on small batch sizes where launch overhead dominates.
|
||
///
|
||
/// # Usage
|
||
/// ```ignore
|
||
/// let mut graph = CapturedForwardGraph::new();
|
||
///
|
||
/// // First call captures the graph
|
||
/// solver.forward_with_graph(&x, &mut ws, &mut graph)?;
|
||
///
|
||
/// // Subsequent calls replay the captured graph (fast!)
|
||
/// for _ in 0..1000 {
|
||
/// solver.forward_with_graph(&x, &mut ws, &mut graph)?;
|
||
/// }
|
||
/// ```
|
||
#[cfg(feature = "cuda")]
|
||
pub struct CapturedForwardGraph {
|
||
/// The captured graph ID (None if not yet captured)
|
||
graph_id: Option<u64>,
|
||
/// The graph manager reference
|
||
graph_manager: Option<std::sync::Arc<rtx_runtime::CudaGraphManager>>,
|
||
/// Input shape for validation (graphs are shape-specific)
|
||
input_shape: Vec<usize>,
|
||
/// Whether the graph has been captured
|
||
is_captured: bool,
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
impl CapturedForwardGraph {
|
||
/// Create a new (uncaptured) graph wrapper
|
||
pub fn new() -> Self {
|
||
Self {
|
||
graph_id: None,
|
||
graph_manager: None,
|
||
input_shape: Vec::new(),
|
||
is_captured: false,
|
||
}
|
||
}
|
||
|
||
/// Check if the graph has been captured
|
||
pub fn is_captured(&self) -> bool {
|
||
self.is_captured
|
||
}
|
||
|
||
/// Invalidate the captured graph (e.g., after parameter updates)
|
||
///
|
||
/// Call this when model parameters change, as the captured graph
|
||
/// references the old parameter values.
|
||
pub fn invalidate(&mut self) {
|
||
if let (Some(gm), Some(gid)) = (&self.graph_manager, self.graph_id) {
|
||
let _ = gm.destroy_graph(gid);
|
||
}
|
||
self.graph_id = None;
|
||
self.is_captured = false;
|
||
self.input_shape.clear();
|
||
}
|
||
|
||
/// Get the number of times this graph has been launched
|
||
pub fn launch_count(&self) -> u64 {
|
||
if let (Some(gm), Some(gid)) = (&self.graph_manager, self.graph_id) {
|
||
gm.graph_info(gid).map(|i| i.launch_count).unwrap_or(0)
|
||
} else {
|
||
0
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
impl Default for CapturedForwardGraph {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
impl Drop for CapturedForwardGraph {
|
||
fn drop(&mut self) {
|
||
self.invalidate();
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// GPU-RESIDENT LOSS ACCUMULATOR (Zero-Sync Training)
|
||
// =============================================================================
|
||
|
||
/// GPU-resident loss accumulator for deferred CPU synchronization.
|
||
///
|
||
/// This struct keeps all loss tensors on the GPU and only transfers to CPU
|
||
/// periodically (every `sync_interval` steps) or on demand. This eliminates
|
||
/// the ~35-55µs overhead from 3-4 blocking `to_cpu()` calls per training step.
|
||
///
|
||
/// # Performance Impact
|
||
/// - Before: 4 GPU-CPU syncs per step (~35-55µs overhead)
|
||
/// - After: ~0.01 syncs per step on average (sync every 100 steps)
|
||
///
|
||
/// # Usage
|
||
/// ```ignore
|
||
/// let mut accumulator = GpuLossAccumulator::new(&device, 100)?;
|
||
/// for epoch in 1..=epochs {
|
||
/// solver.training_step_zero_sync(&mut ws, &mut accumulator)?;
|
||
/// }
|
||
/// let (data_loss, pde_loss, total_loss) = accumulator.force_sync()?.unwrap();
|
||
/// ```
|
||
#[derive(Debug)]
|
||
pub struct GpuLossAccumulator {
|
||
/// Accumulated total loss (stays on GPU)
|
||
total_loss: Tensor,
|
||
/// Accumulated data loss (stays on GPU)
|
||
data_loss: Tensor,
|
||
/// Accumulated PDE loss real part (stays on GPU)
|
||
pde_loss_re: Tensor,
|
||
/// Accumulated PDE loss imaginary part (stays on GPU)
|
||
pde_loss_im: Tensor,
|
||
/// Number of steps accumulated since last sync
|
||
steps: usize,
|
||
/// Sync to CPU every N steps (0 = never auto-sync)
|
||
sync_interval: usize,
|
||
/// Last synced loss values (data_loss, pde_loss, total_loss)
|
||
last_synced: Option<(f64, f64, f64)>,
|
||
/// Device for tensor allocation
|
||
device: Device,
|
||
}
|
||
|
||
impl GpuLossAccumulator {
|
||
/// Create a new GPU-resident loss accumulator.
|
||
///
|
||
/// # Arguments
|
||
/// * `device` - Device to allocate tensors on (should be CUDA)
|
||
/// * `sync_interval` - Sync to CPU every N steps (use 100 for scheduler with patience=500)
|
||
pub fn new(device: &Device, sync_interval: usize) -> Result<Self> {
|
||
Ok(Self {
|
||
total_loss: Tensor::zeros(&[1], device)?,
|
||
data_loss: Tensor::zeros(&[1], device)?,
|
||
pde_loss_re: Tensor::zeros(&[1], device)?,
|
||
pde_loss_im: Tensor::zeros(&[1], device)?,
|
||
steps: 0,
|
||
sync_interval,
|
||
last_synced: None,
|
||
device: device.clone(),
|
||
})
|
||
}
|
||
|
||
/// Accumulate losses WITHOUT any CPU transfer.
|
||
///
|
||
/// All operations stay on GPU - no blocking synchronization.
|
||
pub fn accumulate(
|
||
&mut self,
|
||
data_loss: &Tensor,
|
||
pde_re: &Tensor,
|
||
pde_im: &Tensor,
|
||
data_weight: f32,
|
||
pde_weight: f32,
|
||
) -> Result<()> {
|
||
// All operations stay on GPU (no to_cpu calls!)
|
||
self.data_loss = self.data_loss.add(data_loss)?;
|
||
self.pde_loss_re = self.pde_loss_re.add(pde_re)?;
|
||
self.pde_loss_im = self.pde_loss_im.add(pde_im)?;
|
||
|
||
// Compute weighted total on GPU
|
||
let weighted_data = data_loss.mul_scalar(data_weight)?;
|
||
let pde_sum = pde_re.add(pde_im)?;
|
||
let weighted_pde = pde_sum.mul_scalar(pde_weight)?;
|
||
let step_total = weighted_data.add(&weighted_pde)?;
|
||
self.total_loss = self.total_loss.add(&step_total)?;
|
||
|
||
self.steps += 1;
|
||
Ok(())
|
||
}
|
||
|
||
/// Check if sync is needed and perform it if so.
|
||
///
|
||
/// Returns `Some((data_loss, pde_loss, total_loss))` if synced, `None` otherwise.
|
||
/// Uses `last_synced` value if no sync was performed.
|
||
pub fn maybe_sync(&mut self) -> Result<Option<(f64, f64, f64)>> {
|
||
if self.sync_interval > 0 && self.steps > 0 && self.steps % self.sync_interval == 0 {
|
||
self.force_sync()
|
||
} else {
|
||
Ok(self.last_synced)
|
||
}
|
||
}
|
||
|
||
/// Force synchronization and return averaged losses.
|
||
///
|
||
/// This is the ONLY place where GPU-CPU transfer happens.
|
||
/// Call this at the end of training or when you need the actual loss values.
|
||
pub fn force_sync(&mut self) -> Result<Option<(f64, f64, f64)>> {
|
||
if self.steps == 0 {
|
||
return Ok(None);
|
||
}
|
||
|
||
let scale = 1.0 / self.steps as f32;
|
||
|
||
// These to_cpu() calls are the ONLY blocking syncs
|
||
let avg_data = self.data_loss.mul_scalar(scale)?.to_cpu()?[0] as f64;
|
||
let avg_pde_re = self.pde_loss_re.mul_scalar(scale)?.to_cpu()?[0] as f64;
|
||
let avg_pde_im = self.pde_loss_im.mul_scalar(scale)?.to_cpu()?[0] as f64;
|
||
let avg_total = self.total_loss.mul_scalar(scale)?.to_cpu()?[0] as f64;
|
||
|
||
// Reset accumulators (reuse existing tensors to avoid allocation)
|
||
self.data_loss = Tensor::zeros(&[1], &self.device)?;
|
||
self.pde_loss_re = Tensor::zeros(&[1], &self.device)?;
|
||
self.pde_loss_im = Tensor::zeros(&[1], &self.device)?;
|
||
self.total_loss = Tensor::zeros(&[1], &self.device)?;
|
||
self.steps = 0;
|
||
|
||
let pde_loss = avg_pde_re + avg_pde_im;
|
||
self.last_synced = Some((avg_data, pde_loss, avg_total));
|
||
Ok(self.last_synced)
|
||
}
|
||
|
||
/// Get the last synced values without performing a sync.
|
||
pub fn get_last_synced(&self) -> Option<(f64, f64, f64)> {
|
||
self.last_synced
|
||
}
|
||
|
||
/// Get the number of steps accumulated since last sync.
|
||
pub fn steps_since_sync(&self) -> usize {
|
||
self.steps
|
||
}
|
||
|
||
/// Reset the accumulator without syncing.
|
||
pub fn reset(&mut self) -> Result<()> {
|
||
self.data_loss = Tensor::zeros(&[1], &self.device)?;
|
||
self.pde_loss_re = Tensor::zeros(&[1], &self.device)?;
|
||
self.pde_loss_im = Tensor::zeros(&[1], &self.device)?;
|
||
self.total_loss = Tensor::zeros(&[1], &self.device)?;
|
||
self.steps = 0;
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// LFFN-MLP NETWORK
|
||
// =============================================================================
|
||
|
||
/// Learnable Fourier Feature Network with MLP
|
||
#[derive(Debug)]
|
||
pub struct LffnUNet1D {
|
||
b_learnable: Tensor,
|
||
layers: Vec<Linear>,
|
||
ff_dim: usize,
|
||
device: Device,
|
||
}
|
||
|
||
impl LffnUNet1D {
|
||
/// Get reference to the learnable Fourier feature B matrix
|
||
pub fn b_learnable(&self) -> &Tensor {
|
||
&self.b_learnable
|
||
}
|
||
|
||
/// Get mutable reference to the learnable Fourier feature B matrix
|
||
pub fn b_learnable_mut(&mut self) -> &mut Tensor {
|
||
&mut self.b_learnable
|
||
}
|
||
|
||
/// Get reference to the layer weights (for raw FFI graph capture)
|
||
pub fn layers(&self) -> &[Linear] {
|
||
&self.layers
|
||
}
|
||
|
||
/// Get mutable reference to the layers
|
||
pub fn layers_mut(&mut self) -> &mut [Linear] {
|
||
&mut self.layers
|
||
}
|
||
|
||
/// Extract weights for uber-kernel
|
||
#[cfg(feature = "cuda")]
|
||
pub fn to_uber_kernel_weights(&self, cfg: &Config) -> uber_kernel::UberKernelWeights {
|
||
// B weights need to be reshaped from [1, ff_dim] to [ff_dim]
|
||
let b = self.b_learnable.reshape(&[cfg.u_ff_dim]).expect("Failed to reshape B");
|
||
|
||
uber_kernel::UberKernelWeights {
|
||
b,
|
||
w0: self.layers[0].weight().clone(),
|
||
b0: self.layers[0].bias().expect("Missing bias for layer 0").clone(),
|
||
w1: self.layers[1].weight().clone(),
|
||
b1: self.layers[1].bias().expect("Missing bias for layer 1").clone(),
|
||
w2: self.layers[2].weight().clone(),
|
||
b2: self.layers[2].bias().expect("Missing bias for layer 2").clone(),
|
||
w3: self.layers[3].weight().clone(), // Output layer
|
||
b3: self.layers[3].bias().expect("Missing bias for output layer").clone(),
|
||
ff_dim: cfg.u_ff_dim,
|
||
hidden_dim: cfg.u_hidden,
|
||
}
|
||
}
|
||
|
||
pub fn new(cfg: &Config, device: &Device) -> Result<Self> {
|
||
let b_data: Vec<f32> = (0..cfg.u_ff_dim)
|
||
.map(|_| fastrand::f32() * 2.0 - 1.0)
|
||
.map(|x| x * cfg.u_ff_scale as f32)
|
||
.collect();
|
||
let b_learnable = Tensor::from_slice(&b_data, &[1, cfg.u_ff_dim], device)?;
|
||
let mut layers = Vec::new();
|
||
let mut dim = cfg.u_ff_dim * 2;
|
||
for _ in 0..cfg.u_layers {
|
||
layers.push(Linear::new(dim, cfg.u_hidden, true, device)?);
|
||
dim = cfg.u_hidden;
|
||
}
|
||
layers.push(Linear::new(dim, 2, true, device)?);
|
||
Ok(Self {
|
||
b_learnable,
|
||
layers,
|
||
ff_dim: cfg.u_ff_dim,
|
||
device: device.clone(),
|
||
})
|
||
}
|
||
|
||
/// Create a new network with all tensors allocated on a specific CUDA stream.
|
||
///
|
||
/// This is required for CUDA graph capture - all tensors used during the captured
|
||
/// forward pass must be allocated on the same stream as the graph capture.
|
||
///
|
||
/// # Arguments
|
||
/// * `cfg` - Network configuration
|
||
/// * `stream` - CUDA stream for tensor allocation
|
||
#[cfg(feature = "cuda")]
|
||
pub fn new_on_stream(cfg: &Config, stream: &std::sync::Arc<cudarc::driver::safe::CudaStream>) -> Result<Self> {
|
||
let b_data: Vec<f32> = (0..cfg.u_ff_dim)
|
||
.map(|_| fastrand::f32() * 2.0 - 1.0)
|
||
.map(|x| x * cfg.u_ff_scale as f32)
|
||
.collect();
|
||
let b_learnable = Tensor::from_vec_on_stream(b_data, &[1, cfg.u_ff_dim], stream)?;
|
||
|
||
let mut layers = Vec::new();
|
||
let mut dim = cfg.u_ff_dim * 2;
|
||
for _ in 0..cfg.u_layers {
|
||
layers.push(Linear::new_on_stream(dim, cfg.u_hidden, true, stream)?);
|
||
dim = cfg.u_hidden;
|
||
}
|
||
layers.push(Linear::new_on_stream(dim, 2, true, stream)?);
|
||
|
||
let device = Device::cuda(0)?;
|
||
|
||
Ok(Self {
|
||
b_learnable,
|
||
layers,
|
||
ff_dim: cfg.u_ff_dim,
|
||
device,
|
||
})
|
||
}
|
||
|
||
pub fn forward(&self, x_norm: &Tensor) -> Result<Tensor> {
|
||
// Fourier features computation with in-place operations
|
||
// matmul creates new tensor, then we modify it in-place
|
||
let mut y = x_norm.matmul(&self.b_learnable)?;
|
||
y.mul_scalar_(2.0 * PI as f32)?; // In-place: saves 1 allocation
|
||
|
||
// Need two copies for sin/cos, but apply in-place
|
||
let mut sin_feat = y.clone();
|
||
sin_feat.sin_()?; // In-place: saves 1 allocation
|
||
|
||
// Reuse y for cos (no need to clone again)
|
||
y.cos_()?; // In-place: saves 1 allocation
|
||
let cos_feat = y;
|
||
|
||
let feat = Tensor::cat(&[sin_feat, cos_feat], 1)?;
|
||
|
||
// Process through MLP layers with in-place activations
|
||
let mut x = feat;
|
||
for (i, layer) in self.layers.iter().enumerate() {
|
||
x = layer.forward(&x)?;
|
||
if i < self.layers.len() - 1 {
|
||
x.tanh_()?; // In-place: saves N allocations (one per hidden layer)
|
||
}
|
||
}
|
||
Ok(x)
|
||
}
|
||
|
||
/// Optimized forward pass using pre-allocated workspace buffers
|
||
///
|
||
/// Phase 5 optimizations - FUSED CUDA kernels:
|
||
/// 1. `fused_fourier_features_out` - combines matmul + scale + sin/cos + concat
|
||
/// into a SINGLE CUDA kernel launch with 2D parallelization (replaces 5+ separate kernel launches)
|
||
/// 2. `forward_fused_tanh_out` - Linear layers with FUSED bias+tanh (2 kernels instead of 3)
|
||
/// 3. `forward_out` - Final layer without activation
|
||
///
|
||
/// This dramatically reduces GPU overhead by minimizing kernel launches and
|
||
/// synchronization barriers. After initialization, performs ZERO heap allocations.
|
||
///
|
||
/// Kernel launch count comparison (4 hidden layers):
|
||
/// - Old: 1 (fourier) + 4*(matmul+bias+tanh) + (matmul+bias) = 15 kernel launches
|
||
/// - New: 1 (fourier parallel) + 4*(matmul+fused_bias_tanh) + (matmul+bias) = 11 kernel launches
|
||
/// - Reduction: ~27% fewer kernel launches
|
||
pub fn forward_with_workspace(&self, x_norm: &Tensor, ws: &mut ForwardWorkspace) -> Result<()> {
|
||
// Step 1: FUSED Fourier features computation (ONE kernel with 2D parallelization)
|
||
// Each thread handles one (batch_idx, ff_idx) pair instead of looping
|
||
x_norm.fused_fourier_features_out(&self.b_learnable, 2.0 * PI as f32, &mut ws.features)?;
|
||
|
||
// Step 2: Process through MLP layers with FUSED bias+tanh kernels
|
||
// First hidden layer takes features as input
|
||
self.layers[0].forward_fused_tanh_out(&ws.features, &mut ws.hidden_layers[0])?;
|
||
|
||
// Middle hidden layers - use split_at_mut to avoid borrow checker issues
|
||
for i in 1..self.layers.len() - 1 {
|
||
let (left, right) = ws.hidden_layers.split_at_mut(i);
|
||
let input = &left[i - 1];
|
||
let output = &mut right[0];
|
||
self.layers[i].forward_fused_tanh_out(input, output)?;
|
||
}
|
||
|
||
// Final layer: write to output buffer (no activation)
|
||
let last_hidden_idx = self.layers.len() - 2;
|
||
let last_layer_idx = self.layers.len() - 1;
|
||
self.layers[last_layer_idx].forward_out(&ws.hidden_layers[last_hidden_idx], &mut ws.output)?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
pub fn parameters(&self) -> Vec<&Tensor> {
|
||
let mut params = vec![&self.b_learnable];
|
||
for layer in &self.layers {
|
||
params.extend(layer.parameters());
|
||
}
|
||
params
|
||
}
|
||
|
||
/// Forward pass that caches intermediate activations for backprop.
|
||
///
|
||
/// Pipeline: x -> Fourier -> h0 -> h1 -> h2 -> h3 -> output
|
||
/// Each h_i is saved in workspace.h[i] for backward pass.
|
||
///
|
||
/// This method is used by `train_step_data_only()` to enable analytical
|
||
/// gradient computation without autograd.
|
||
pub fn forward_training(
|
||
&self,
|
||
x: &Tensor,
|
||
workspace: &mut GradientWorkspace,
|
||
) -> Result<()> {
|
||
use std::f64::consts::PI;
|
||
|
||
// 1. Fourier Layer: x -> [sin(2πBx), cos(2πBx)]
|
||
// Save as workspace.fourier_features (input to layer 0)
|
||
x.fused_fourier_features_out(&self.b_learnable, 2.0 * PI as f32, &mut workspace.fourier_features)?;
|
||
|
||
// Copy x for Fourier gradient computation later
|
||
workspace.x_input = x.clone();
|
||
|
||
// 2. Hidden Layers with activation caching
|
||
// Layer 0: features -> h[0] (with tanh)
|
||
self.layers[0].forward_fused_tanh_out(&workspace.fourier_features, &mut workspace.h[0])?;
|
||
|
||
// Layers 1..N-1: h[i-1] -> h[i] (with tanh)
|
||
for i in 1..self.layers.len() - 1 {
|
||
let (left, right) = workspace.h.split_at_mut(i);
|
||
self.layers[i].forward_fused_tanh_out(&left[i - 1], &mut right[0])?;
|
||
}
|
||
|
||
// Output layer: h[N-2] -> h[N-1] (NO tanh)
|
||
// Use split_at_mut to avoid borrow checker issues
|
||
let last_hidden = self.layers.len() - 2;
|
||
let last_layer = self.layers.len() - 1;
|
||
let (left, right) = workspace.h.split_at_mut(last_layer);
|
||
self.layers[last_layer].forward_out(&left[last_hidden], &mut right[0])?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Forward pass using unified stream context for CUDA graph capture.
|
||
///
|
||
/// This method executes ALL operations on a single unified stream, enabling
|
||
/// CUDA graph capture. Unlike `forward_with_workspace`, which uses multiple
|
||
/// internal streams (breaking graph capture), this method uses the stream
|
||
/// from `PinnStreamContext` for all operations.
|
||
///
|
||
/// ## Why This Works
|
||
///
|
||
/// The `PinnStreamContext` provides:
|
||
/// - A single CudaContext and CudaStream for all operations
|
||
/// - A cuBLAS handle bound to the unified stream
|
||
/// - Pre-loaded PTX kernels that launch on the unified stream
|
||
///
|
||
/// ## Stable Pointer Architecture
|
||
///
|
||
/// CUDA Graphs capture POINTERS, not VALUES. With in-place optimizer updates:
|
||
/// - Graph captured once at first iteration
|
||
/// - Replayed ~50,000 times during training
|
||
/// - No recapture needed (memory addresses stay stable)
|
||
///
|
||
/// ## Kernel Launch Sequence
|
||
///
|
||
/// 1 Fourier + 4 (matmul+bias_tanh) + 2 output = 11 launches
|
||
/// With graph capture: 1 launch (graph replay)
|
||
///
|
||
/// # Arguments
|
||
/// * `x_norm` - Normalized input tensor [batch, 1]
|
||
/// * `ws` - Pre-allocated workspace buffers
|
||
/// * `ctx` - Unified stream context (all ops use this stream)
|
||
#[cfg(feature = "cuda")]
|
||
pub fn forward_on_stream(
|
||
&self,
|
||
x_norm: &Tensor,
|
||
ws: &mut ForwardWorkspace,
|
||
ctx: &PinnStreamContext,
|
||
) -> Result<()> {
|
||
// Step 1: Fused Fourier features (1 kernel launch)
|
||
ctx.fourier_features_out(x_norm, &self.b_learnable, 2.0 * PI as f32, &mut ws.features)?;
|
||
|
||
// Step 2: Hidden layers with fused bias+tanh (2 kernels each: matmul + fused_bias_tanh)
|
||
// First hidden layer: features -> hidden[0]
|
||
ctx.matmul_out(&ws.features, self.layers[0].weight_t(), &mut ws.hidden_layers[0])?;
|
||
if let Some(bias) = self.layers[0].bias() {
|
||
ctx.bias_add_tanh_(&mut ws.hidden_layers[0], bias)?;
|
||
}
|
||
|
||
// Middle hidden layers: hidden[i-1] -> hidden[i]
|
||
for i in 1..self.layers.len() - 1 {
|
||
let (left, right) = ws.hidden_layers.split_at_mut(i);
|
||
let input = &left[i - 1];
|
||
let output = &mut right[0];
|
||
|
||
ctx.matmul_out(input, self.layers[i].weight_t(), output)?;
|
||
if let Some(bias) = self.layers[i].bias() {
|
||
ctx.bias_add_tanh_(output, bias)?;
|
||
}
|
||
}
|
||
|
||
// Step 3: Output layer (2 kernels: matmul + bias_add)
|
||
let last_hidden_idx = self.layers.len() - 2;
|
||
let last_layer_idx = self.layers.len() - 1;
|
||
ctx.matmul_out(
|
||
&ws.hidden_layers[last_hidden_idx],
|
||
self.layers[last_layer_idx].weight_t(),
|
||
&mut ws.output,
|
||
)?;
|
||
if let Some(bias) = self.layers[last_layer_idx].bias() {
|
||
ctx.add_bias_(&mut ws.output, bias)?;
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Forward pass using CUDA graph capture for minimal overhead.
|
||
///
|
||
/// On the first call, this attempts to capture the forward pass operations into a CUDA graph.
|
||
/// On subsequent calls, it replays the captured graph with a single API call,
|
||
/// eliminating kernel launch overhead.
|
||
///
|
||
/// **Note**: CUDA graph capture requires all operations to use the same CUDA stream.
|
||
/// Currently, rtx-tensor uses multiple cached CUDA contexts internally, which may
|
||
/// cause capture to fail. In that case, this method falls back to the regular
|
||
/// `forward_with_workspace` method.
|
||
///
|
||
/// # Arguments
|
||
/// * `x_norm` - Normalized input tensor (must have same shape as during capture)
|
||
/// * `ws` - Pre-allocated workspace
|
||
/// * `graph` - Graph capture state (reused across calls)
|
||
///
|
||
/// # Notes
|
||
/// - The graph is invalidated if input shape changes
|
||
/// - Call `graph.invalidate()` after parameter updates
|
||
/// - Falls back to regular forward if capture fails
|
||
#[cfg(feature = "cuda")]
|
||
pub fn forward_with_graph(
|
||
&self,
|
||
x_norm: &Tensor,
|
||
ws: &mut ForwardWorkspace,
|
||
graph: &mut CapturedForwardGraph,
|
||
) -> Result<()> {
|
||
// For now, CUDA graph capture is disabled due to stream synchronization issues.
|
||
// rtx-tensor uses multiple cached CUDA contexts internally, each with its own stream.
|
||
// CUDA graph capture requires all operations to use the SAME stream.
|
||
//
|
||
// TODO: Phase 1.5 - Unify rtx-tensor to use a single shared CUDA context/stream
|
||
// This will enable CUDA graph capture across all tensor operations.
|
||
//
|
||
// For now, fall back to the optimized workspace-based forward pass.
|
||
self.forward_with_workspace(x_norm, ws)
|
||
}
|
||
|
||
/// [EXPERIMENTAL] Forward pass with CUDA graph capture.
|
||
///
|
||
/// This is an experimental implementation that attempts stream capture.
|
||
/// Due to rtx-tensor's multi-context architecture, capture may fail.
|
||
/// Use `forward_with_workspace` for production code.
|
||
#[cfg(feature = "cuda")]
|
||
#[allow(dead_code)]
|
||
fn forward_with_graph_experimental(
|
||
&self,
|
||
x_norm: &Tensor,
|
||
ws: &mut ForwardWorkspace,
|
||
graph: &mut CapturedForwardGraph,
|
||
) -> Result<()> {
|
||
use rtx_runtime::{CudaBackend, CudaGraphManager, DeviceId};
|
||
|
||
// Validate input shape matches captured graph
|
||
let current_shape: Vec<usize> = x_norm.shape().iter().map(|&x| x as usize).collect();
|
||
if graph.is_captured && graph.input_shape != current_shape {
|
||
// Shape changed - invalidate and re-capture
|
||
graph.invalidate();
|
||
}
|
||
|
||
if !graph.is_captured {
|
||
// First call: attempt to capture the forward pass into a CUDA graph
|
||
let device_id = match self.device {
|
||
Device::Cuda(id) => id,
|
||
_ => return self.forward_with_workspace(x_norm, ws), // Fall back on CPU
|
||
};
|
||
|
||
// Try to create backend and graph manager
|
||
let backend = match CudaBackend::new(DeviceId(device_id as u32)) {
|
||
Ok(b) => std::sync::Arc::new(b),
|
||
Err(_) => return self.forward_with_workspace(x_norm, ws), // Fall back
|
||
};
|
||
let graph_manager = match CudaGraphManager::new(backend.clone()) {
|
||
Ok(gm) => std::sync::Arc::new(gm),
|
||
Err(_) => return self.forward_with_workspace(x_norm, ws), // Fall back
|
||
};
|
||
|
||
// Get the default stream
|
||
let stream = backend.default_stream();
|
||
|
||
// Try to begin capture - may fail due to multi-stream issues
|
||
let graph_id = match graph_manager.begin_capture(&stream) {
|
||
Ok(id) => id,
|
||
Err(e) => {
|
||
eprintln!("CUDA graph capture not supported: {:?}", e);
|
||
return self.forward_with_workspace(x_norm, ws); // Fall back
|
||
}
|
||
};
|
||
|
||
// Execute the forward pass (operations get recorded)
|
||
self.forward_with_workspace(x_norm, ws)?;
|
||
|
||
// End capture - may fail if operations used different streams
|
||
match graph_manager.end_capture(&stream) {
|
||
Ok(captured_id) => {
|
||
if graph_id != captured_id {
|
||
eprintln!("Graph ID mismatch, falling back to regular forward");
|
||
return Ok(()); // Forward pass already executed
|
||
}
|
||
// Store capture state
|
||
graph.graph_id = Some(graph_id);
|
||
graph.graph_manager = Some(graph_manager);
|
||
graph.input_shape = current_shape;
|
||
graph.is_captured = true;
|
||
}
|
||
Err(e) => {
|
||
eprintln!("CUDA graph capture failed: {:?}, using regular forward", e);
|
||
// Forward pass already executed during capture attempt
|
||
}
|
||
}
|
||
Ok(())
|
||
} else {
|
||
// Subsequent calls: replay the captured graph
|
||
if let (Some(gm), Some(gid)) = (&graph.graph_manager, graph.graph_id) {
|
||
gm.launch(gid)
|
||
.map_err(|e| anyhow::anyhow!("Failed to launch graph: {:?}", e))?;
|
||
}
|
||
Ok(())
|
||
}
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// PINN SOLVER
|
||
// =============================================================================
|
||
|
||
/// PINN Solver for 1D MRE Helmholtz equation
|
||
pub struct Mre1DPinnSolver {
|
||
pub cfg: Config,
|
||
pub omega: f64,
|
||
pub k_true: Complex64,
|
||
pub u_scale: f64,
|
||
pub l: f64,
|
||
pub u_net: LffnUNet1D,
|
||
pub x_data: Tensor,
|
||
pub u_data_target: Tensor,
|
||
g_prime: f64,
|
||
g_double: f64,
|
||
optimizer: AdamOptimizer,
|
||
scheduler: ReduceLROnPlateauScheduler,
|
||
pub device: Device,
|
||
|
||
// === PHASE 1 FIX: Cached PDE tensors (pre-computed once) ===
|
||
// These tensors are computed once at construction and reused every epoch,
|
||
// eliminating ~89% overhead from repeated allocations.
|
||
cached_u_r_t: Tensor, // u_real as tensor [n_data, 1]
|
||
cached_u_i_t: Tensor, // u_imag as tensor [n_data, 1]
|
||
cached_d2u_r_t: Tensor, // d²u/dx² real part [n_data, 1]
|
||
cached_d2u_i_t: Tensor, // d²u/dx² imag part [n_data, 1]
|
||
cached_k_sq_re: f32, // k² real part
|
||
cached_k_sq_im: f32, // k² imag part
|
||
|
||
// === PHASE 8: Zero-Sync Training Fields ===
|
||
// Gradient workspace (allocated once at construction, reused every step)
|
||
grad_workspace: GradientWorkspace,
|
||
// GPU-native Adam optimizer (all state on GPU, no CPU-GPU sync)
|
||
gpu_adam: GpuAdam,
|
||
// CUDA stream context for zero-allocation ops (Option because CPU doesn't have it)
|
||
#[cfg(feature = "cuda")]
|
||
stream_ctx: PinnStreamContext,
|
||
}
|
||
|
||
impl Mre1DPinnSolver {
|
||
pub fn new(cfg: Config) -> Result<Self> {
|
||
#[cfg(feature = "cuda")]
|
||
let device = Device::Cuda(0);
|
||
#[cfg(not(feature = "cuda"))]
|
||
let device = Device::Cpu;
|
||
|
||
let omega = 2.0 * PI * cfg.freq;
|
||
let l = cfg.l;
|
||
let (x, u_real, u_imag, k_true) = synthesize_displacement(&cfg);
|
||
let u_mag_max = u_real.iter().zip(u_imag.iter())
|
||
.map(|(&r, &i)| (r * r + i * i).sqrt())
|
||
.fold(0.0f64, |a, b| a.max(b));
|
||
let u_scale = u_mag_max + 1e-16;
|
||
let x_norm: Vec<f32> = x.iter().map(|&xi| (xi / l) as f32).collect();
|
||
let u_data_norm: Vec<f32> = u_real.iter().zip(u_imag.iter())
|
||
.flat_map(|(&r, &i)| vec![(r / u_scale) as f32, (i / u_scale) as f32])
|
||
.collect();
|
||
let x_data = Tensor::from_slice(&x_norm, &[cfg.n_data, 1], &device)?;
|
||
let u_data_target = Tensor::from_slice(&u_data_norm, &[cfg.n_data, 2], &device)?;
|
||
let u_net = LffnUNet1D::new(&cfg, &device)?;
|
||
let adam_config = AdamConfig {
|
||
learning_rate: cfg.lr,
|
||
beta1: 0.9,
|
||
beta2: 0.999,
|
||
epsilon: 1e-8,
|
||
weight_decay: 0.0,
|
||
amsgrad: false,
|
||
};
|
||
let optimizer = AdamOptimizer::new(adam_config)?;
|
||
let scheduler = ReduceLROnPlateauScheduler::builder(cfg.lr)
|
||
.mode(PlateauMode::Min)
|
||
.factor(cfg.scheduler_factor)
|
||
.patience(cfg.scheduler_patience)
|
||
.threshold(1e-4)
|
||
.threshold_mode(ThresholdMode::Rel)
|
||
.min_lr(1e-8)
|
||
.verbose(false)
|
||
.build()?;
|
||
|
||
// === PHASE 1 FIX: Pre-compute and cache PDE tensors ONCE ===
|
||
// This eliminates ~89% overhead from repeated allocations in training loop
|
||
let (_, d2udx2) = synthesize_exact_derivatives(&cfg, k_true);
|
||
let cached_u_r_t = Tensor::from_slice(
|
||
&u_real.iter().map(|&x| x as f32).collect::<Vec<_>>(),
|
||
&[cfg.n_data, 1],
|
||
&device
|
||
)?;
|
||
let cached_u_i_t = Tensor::from_slice(
|
||
&u_imag.iter().map(|&x| x as f32).collect::<Vec<_>>(),
|
||
&[cfg.n_data, 1],
|
||
&device
|
||
)?;
|
||
let cached_d2u_r_t = Tensor::from_slice(
|
||
&d2udx2.iter().map(|x| x[0] as f32).collect::<Vec<_>>(),
|
||
&[cfg.n_data, 1],
|
||
&device
|
||
)?;
|
||
let cached_d2u_i_t = Tensor::from_slice(
|
||
&d2udx2.iter().map(|x| x[1] as f32).collect::<Vec<_>>(),
|
||
&[cfg.n_data, 1],
|
||
&device
|
||
)?;
|
||
let k_sq = k_true * k_true;
|
||
let cached_k_sq_re = k_sq.re as f32;
|
||
let cached_k_sq_im = k_sq.im as f32;
|
||
|
||
// === PHASE 8: Initialize zero-sync training infrastructure ===
|
||
// GradientWorkspace: pre-allocate all buffers for backward pass
|
||
let grad_workspace = GradientWorkspace::new(
|
||
cfg.n_data,
|
||
cfg.u_ff_dim,
|
||
cfg.u_hidden,
|
||
cfg.u_layers, // num_hidden_layers
|
||
&device,
|
||
).map_err(|e| anyhow::anyhow!("Failed to create GradientWorkspace: {}", e))?;
|
||
|
||
// Collect parameter shapes for GpuAdam initialization
|
||
// Order: B, then for each layer: (W, b)
|
||
let mut param_shapes = Vec::new();
|
||
// B matrix: [1, ff_dim]
|
||
param_shapes.push(vec![1, cfg.u_ff_dim]);
|
||
// Layer weights and biases
|
||
let mut in_dim = cfg.u_ff_dim * 2; // First layer input is Fourier features
|
||
for i in 0..=cfg.u_layers {
|
||
let out_dim = if i == cfg.u_layers { 2 } else { cfg.u_hidden };
|
||
// Weight: [out_dim, in_dim]
|
||
param_shapes.push(vec![out_dim, in_dim]);
|
||
// Bias: [out_dim]
|
||
param_shapes.push(vec![out_dim]);
|
||
in_dim = out_dim;
|
||
}
|
||
|
||
let gpu_adam = GpuAdam::new(
|
||
¶m_shapes,
|
||
&device,
|
||
cfg.lr as f32,
|
||
0.9, // beta1
|
||
0.999, // beta2
|
||
1e-8, // eps
|
||
).map_err(|e| anyhow::anyhow!("Failed to create GpuAdam: {}", e))?;
|
||
|
||
// Create CUDA stream context for zero-allocation operations
|
||
#[cfg(feature = "cuda")]
|
||
let stream_ctx = PinnStreamContext::new(0)
|
||
.map_err(|e| anyhow::anyhow!("Failed to create PinnStreamContext: {}", e))?;
|
||
|
||
Ok(Self {
|
||
omega,
|
||
k_true,
|
||
u_scale,
|
||
l,
|
||
u_net,
|
||
x_data,
|
||
u_data_target,
|
||
g_prime: cfg.g_prime_true,
|
||
g_double: cfg.g_double_true,
|
||
optimizer,
|
||
scheduler,
|
||
device,
|
||
cfg,
|
||
// Cached tensors
|
||
cached_u_r_t,
|
||
cached_u_i_t,
|
||
cached_d2u_r_t,
|
||
cached_d2u_i_t,
|
||
cached_k_sq_re,
|
||
cached_k_sq_im,
|
||
// Zero-sync training
|
||
grad_workspace,
|
||
gpu_adam,
|
||
#[cfg(feature = "cuda")]
|
||
stream_ctx,
|
||
})
|
||
}
|
||
|
||
/// Compute MSE loss between two tensors
|
||
pub fn mse_loss(pred: &Tensor, target: &Tensor) -> Result<f64> {
|
||
let diff = pred.sub(target)?;
|
||
let sq = diff.mul(&diff)?;
|
||
let mean = sq.mean(&[], false)?;
|
||
let data = mean.to_cpu()?;
|
||
Ok(data[0] as f64)
|
||
}
|
||
|
||
/// Compute data loss
|
||
pub fn compute_data_loss(&self, u_pred: &Tensor) -> Result<f64> {
|
||
Self::mse_loss(u_pred, &self.u_data_target)
|
||
}
|
||
|
||
// =========================================================================
|
||
// GPU-ONLY LOSS FUNCTIONS (Zero-Sync)
|
||
// =========================================================================
|
||
// These functions return Tensors (on GPU) instead of f64 scalars,
|
||
// avoiding the blocking to_cpu() call that causes GPU-CPU synchronization.
|
||
|
||
/// Compute MSE loss - GPU-resident version (no sync).
|
||
///
|
||
/// Returns a scalar Tensor on GPU instead of f64.
|
||
/// Use this with GpuLossAccumulator for zero-sync training.
|
||
pub fn mse_loss_gpu(pred: &Tensor, target: &Tensor) -> Result<Tensor> {
|
||
let diff = pred.sub(target)?;
|
||
let sq = diff.mul(&diff)?;
|
||
Ok(sq.mean(&[], false)?) // Returns scalar tensor on GPU - NO to_cpu()!
|
||
}
|
||
|
||
/// Compute data loss - GPU-resident version (no sync).
|
||
pub fn compute_data_loss_gpu(&self, u_pred: &Tensor) -> Result<Tensor> {
|
||
Self::mse_loss_gpu(u_pred, &self.u_data_target)
|
||
}
|
||
|
||
/// Compute PDE residual - GPU-resident version (no sync).
|
||
///
|
||
/// Returns (mse_re_tensor, mse_im_tensor) both on GPU.
|
||
/// Avoids 2 blocking to_cpu() calls per training step.
|
||
pub fn compute_pde_residual_gpu(&self) -> Result<(Tensor, Tensor)> {
|
||
// Use cached tensors instead of recreating them
|
||
let ku_re = self.cached_u_r_t.mul_scalar(self.cached_k_sq_re)?
|
||
.sub(&self.cached_u_i_t.mul_scalar(self.cached_k_sq_im)?)?;
|
||
let ku_im = self.cached_u_i_t.mul_scalar(self.cached_k_sq_re)?
|
||
.add(&self.cached_u_r_t.mul_scalar(self.cached_k_sq_im)?)?;
|
||
let res_re = self.cached_d2u_r_t.add(&ku_re)?;
|
||
let res_im = self.cached_d2u_i_t.add(&ku_im)?;
|
||
let mse_re = res_re.mul(&res_re)?.mean(&[], false)?;
|
||
let mse_im = res_im.mul(&res_im)?.mean(&[], false)?;
|
||
// Return GPU tensors - NO to_cpu()!
|
||
Ok((mse_re, mse_im))
|
||
}
|
||
|
||
|
||
/// Compute PDE residual using analytical derivatives (CPU-based)
|
||
pub fn compute_pde_residual_analytical(&self) -> (f64, f64) {
|
||
let (_dudx, d2udx2) = synthesize_exact_derivatives(&self.cfg, self.k_true);
|
||
let (_, u_r, u_i, _) = synthesize_displacement(&self.cfg);
|
||
let k_sq = self.k_true * self.k_true;
|
||
let mut sum_rr = 0.0;
|
||
let mut sum_ri = 0.0;
|
||
for i in 0..self.cfg.n_data {
|
||
let u = Complex64::new(u_r[i], u_i[i]);
|
||
let d2u = Complex64::new(d2udx2[i][0], d2udx2[i][1]);
|
||
let residual = d2u + k_sq * u;
|
||
sum_rr += residual.re.powi(2);
|
||
sum_ri += residual.im.powi(2);
|
||
}
|
||
(sum_rr / self.cfg.n_data as f64, sum_ri / self.cfg.n_data as f64)
|
||
}
|
||
|
||
/// Compute PDE residual using tensor operations (GPU-accelerable)
|
||
/// WARNING: This method allocates ~6 tensors per call - use compute_pde_residual_cached() instead!
|
||
pub fn compute_pde_residual_tensor(&self) -> Result<(f64, f64)> {
|
||
let (_, d2udx2) = synthesize_exact_derivatives(&self.cfg, self.k_true);
|
||
let (_, u_r, u_i, _) = synthesize_displacement(&self.cfg);
|
||
let u_r_t = Tensor::from_slice(
|
||
&u_r.iter().map(|&x| x as f32).collect::<Vec<_>>(),
|
||
&[self.cfg.n_data, 1],
|
||
&self.device
|
||
)?;
|
||
let u_i_t = Tensor::from_slice(
|
||
&u_i.iter().map(|&x| x as f32).collect::<Vec<_>>(),
|
||
&[self.cfg.n_data, 1],
|
||
&self.device
|
||
)?;
|
||
let d2u_r_t = Tensor::from_slice(
|
||
&d2udx2.iter().map(|x| x[0] as f32).collect::<Vec<_>>(),
|
||
&[self.cfg.n_data, 1],
|
||
&self.device
|
||
)?;
|
||
let d2u_i_t = Tensor::from_slice(
|
||
&d2udx2.iter().map(|x| x[1] as f32).collect::<Vec<_>>(),
|
||
&[self.cfg.n_data, 1],
|
||
&self.device
|
||
)?;
|
||
let k_sq = self.k_true * self.k_true;
|
||
let k_sq_re = k_sq.re as f32;
|
||
let k_sq_im = k_sq.im as f32;
|
||
let ku_re = u_r_t.mul_scalar(k_sq_re)?.sub(&u_i_t.mul_scalar(k_sq_im)?)?;
|
||
let ku_im = u_i_t.mul_scalar(k_sq_re)?.add(&u_r_t.mul_scalar(k_sq_im)?)?;
|
||
let res_re = d2u_r_t.add(&ku_re)?;
|
||
let res_im = d2u_i_t.add(&ku_im)?;
|
||
let mse_re = res_re.mul(&res_re)?.mean(&[], false)?;
|
||
let mse_im = res_im.mul(&res_im)?.mean(&[], false)?;
|
||
Ok((mse_re.to_cpu()?[0] as f64, mse_im.to_cpu()?[0] as f64))
|
||
}
|
||
|
||
/// Compute PDE residual using CACHED tensors (ZERO allocation version)
|
||
///
|
||
/// This is the optimized version that uses pre-computed tensors stored in the solver.
|
||
/// Eliminates ~89% overhead from the training loop by avoiding tensor creation.
|
||
///
|
||
/// Performance: ~2.87ms saved per call (at 200 points)
|
||
pub fn compute_pde_residual_cached(&self) -> Result<(f64, f64)> {
|
||
// Use cached tensors instead of recreating them
|
||
let ku_re = self.cached_u_r_t.mul_scalar(self.cached_k_sq_re)?
|
||
.sub(&self.cached_u_i_t.mul_scalar(self.cached_k_sq_im)?)?;
|
||
let ku_im = self.cached_u_i_t.mul_scalar(self.cached_k_sq_re)?
|
||
.add(&self.cached_u_r_t.mul_scalar(self.cached_k_sq_im)?)?;
|
||
let res_re = self.cached_d2u_r_t.add(&ku_re)?;
|
||
let res_im = self.cached_d2u_i_t.add(&ku_im)?;
|
||
let mse_re = res_re.mul(&res_re)?.mean(&[], false)?;
|
||
let mse_im = res_im.mul(&res_im)?.mean(&[], false)?;
|
||
Ok((mse_re.to_cpu()?[0] as f64, mse_im.to_cpu()?[0] as f64))
|
||
}
|
||
|
||
/// Run a single training step (for benchmarking)
|
||
/// NOTE: Uses compute_pde_residual_tensor() which allocates ~6 tensors per call
|
||
pub fn training_step(&mut self) -> Result<(f64, f64, f64)> {
|
||
clear_tape();
|
||
let u_pred = self.u_net.forward(&self.x_data)?;
|
||
let loss_data = self.compute_data_loss(&u_pred)?;
|
||
let (mse_rr, mse_ri) = self.compute_pde_residual_tensor()?;
|
||
let loss_pde = mse_rr + mse_ri;
|
||
let total_loss = self.cfg.data_weight * loss_data + self.cfg.pde_weight * loss_pde;
|
||
let _new_lr = self.scheduler.step_metric(loss_data);
|
||
Ok((loss_data, loss_pde, total_loss))
|
||
}
|
||
|
||
/// Run a single training step using CACHED PDE tensors (optimized)
|
||
///
|
||
/// This is the recommended method for training loops. Uses pre-computed tensors
|
||
/// to eliminate allocation overhead from compute_pde_residual_tensor().
|
||
///
|
||
/// Performance improvement: ~5-10x faster training loop
|
||
pub fn training_step_cached(&mut self) -> Result<(f64, f64, f64)> {
|
||
clear_tape();
|
||
let u_pred = self.u_net.forward(&self.x_data)?;
|
||
let loss_data = self.compute_data_loss(&u_pred)?;
|
||
let (mse_rr, mse_ri) = self.compute_pde_residual_cached()?; // Use cached version!
|
||
let loss_pde = mse_rr + mse_ri;
|
||
let total_loss = self.cfg.data_weight * loss_data + self.cfg.pde_weight * loss_pde;
|
||
let _new_lr = self.scheduler.step_metric(loss_data);
|
||
Ok((loss_data, loss_pde, total_loss))
|
||
}
|
||
|
||
/// Create a workspace for optimized forward passes
|
||
pub fn create_workspace(&self) -> Result<ForwardWorkspace> {
|
||
ForwardWorkspace::new(self.cfg.n_data, &self.cfg, &self.device)
|
||
}
|
||
|
||
/// Create a unified stream context for CUDA graph capture.
|
||
///
|
||
/// The returned `PinnStreamContext` provides:
|
||
/// - A single CudaContext and CudaStream for all operations
|
||
/// - A cuBLAS handle bound to the unified stream
|
||
/// - Pre-loaded PTX kernels (no JIT during graph capture)
|
||
///
|
||
/// Use with `forward_on_stream` and `PinnGraph` for maximum performance.
|
||
#[cfg(feature = "cuda")]
|
||
pub fn create_stream_context(&self) -> Result<PinnStreamContext> {
|
||
let device_id = match self.device {
|
||
Device::Cuda(id) => id,
|
||
_ => return Err(anyhow::anyhow!("CUDA device required for stream context")),
|
||
};
|
||
PinnStreamContext::new(device_id).map_err(|e| anyhow::anyhow!("{}", e))
|
||
}
|
||
|
||
/// Create a graph-ready workspace and stream context pair.
|
||
///
|
||
/// This method creates:
|
||
/// 1. A `PinnStreamContext` with unified stream for all operations
|
||
/// 2. A `ForwardWorkspace` with all tensors allocated on that stream
|
||
///
|
||
/// By allocating workspace tensors on the capture stream, CUDA graph
|
||
/// capture will work without cross-stream dependency errors.
|
||
///
|
||
/// # Returns
|
||
/// A tuple of (ForwardWorkspace, PinnStreamContext) ready for graph capture.
|
||
///
|
||
/// # Example
|
||
/// ```ignore
|
||
/// let (mut ws, ctx) = solver.create_graph_ready_workspace()?;
|
||
/// let mut graph = PinnGraph::new(Arc::new(ctx));
|
||
///
|
||
/// // Graph capture will succeed - all tensors on same stream
|
||
/// graph.capture(&[200, 1], || {
|
||
/// solver.u_net.forward_on_stream(&solver.x_data, &mut ws, &ctx)
|
||
/// })?;
|
||
/// ```
|
||
#[cfg(feature = "cuda")]
|
||
pub fn create_graph_ready_workspace(&self) -> Result<(ForwardWorkspace, PinnStreamContext)> {
|
||
let ctx = self.create_stream_context()?;
|
||
let ws = ForwardWorkspace::new_on_stream(
|
||
self.cfg.n_data,
|
||
&self.cfg,
|
||
ctx.stream_for_alloc(),
|
||
).map_err(|e| anyhow::anyhow!("{}", e))?;
|
||
Ok((ws, ctx))
|
||
}
|
||
|
||
/// Create a solver with ALL tensors allocated on a single CUDA stream.
|
||
///
|
||
/// This is the **only** way to achieve CUDA graph capture. All tensors that
|
||
/// are accessed during the captured forward pass (model weights, x_data,
|
||
/// workspace) must be allocated on the same stream.
|
||
///
|
||
/// # Returns
|
||
/// A tuple of (solver, workspace, stream_context) all using the same stream.
|
||
///
|
||
/// # Example
|
||
/// ```ignore
|
||
/// let (solver, mut ws, ctx) = Mre1DPinnSolver::new_for_graph_capture(Config::default())?;
|
||
/// let mut graph = PinnGraph::new(Arc::new(ctx));
|
||
///
|
||
/// // Graph capture will succeed - ALL tensors on same stream
|
||
/// graph.capture(&[solver.cfg.n_data, 1], || {
|
||
/// solver.u_net.forward_on_stream(&solver.x_data, &mut ws, &ctx)
|
||
/// })?;
|
||
/// ```
|
||
#[cfg(feature = "cuda")]
|
||
pub fn new_for_graph_capture(cfg: Config) -> Result<(Self, ForwardWorkspace, PinnStreamContext)> {
|
||
use std::sync::Arc;
|
||
|
||
let omega = 2.0 * PI * cfg.freq;
|
||
let l = cfg.l;
|
||
let (x, u_real, u_imag, k_true) = synthesize_displacement(&cfg);
|
||
let u_mag_max = u_real.iter().zip(u_imag.iter())
|
||
.map(|(&r, &i)| (r * r + i * i).sqrt())
|
||
.fold(0.0f64, |a, b| a.max(b));
|
||
let u_scale = u_mag_max + 1e-16;
|
||
let x_norm: Vec<f32> = x.iter().map(|&xi| (xi / l) as f32).collect();
|
||
let u_data_norm: Vec<f32> = u_real.iter().zip(u_imag.iter())
|
||
.flat_map(|(&r, &i)| vec![(r / u_scale) as f32, (i / u_scale) as f32])
|
||
.collect();
|
||
|
||
// Create stream context FIRST - all allocations will use this stream
|
||
let ctx = PinnStreamContext::new(0)?;
|
||
let stream = ctx.stream_for_alloc();
|
||
|
||
// Allocate ALL tensors on the capture stream
|
||
let x_data = Tensor::from_vec_on_stream(x_norm.clone(), &[cfg.n_data, 1], stream)?;
|
||
let u_data_target = Tensor::from_vec_on_stream(u_data_norm, &[cfg.n_data, 2], stream)?;
|
||
|
||
// Network with weights on capture stream
|
||
let u_net = LffnUNet1D::new_on_stream(&cfg, stream)?;
|
||
|
||
// Workspace on capture stream
|
||
let ws = ForwardWorkspace::new_on_stream(cfg.n_data, &cfg, stream)?;
|
||
|
||
let adam_config = AdamConfig {
|
||
learning_rate: cfg.lr,
|
||
beta1: 0.9,
|
||
beta2: 0.999,
|
||
epsilon: 1e-8,
|
||
weight_decay: 0.0,
|
||
amsgrad: false,
|
||
};
|
||
let optimizer = AdamOptimizer::new(adam_config)?;
|
||
let scheduler = ReduceLROnPlateauScheduler::builder(cfg.lr)
|
||
.mode(PlateauMode::Min)
|
||
.factor(cfg.scheduler_factor)
|
||
.patience(cfg.scheduler_patience)
|
||
.threshold(1e-4)
|
||
.threshold_mode(ThresholdMode::Rel)
|
||
.min_lr(1e-8)
|
||
.verbose(false)
|
||
.build()?;
|
||
|
||
// Cached PDE tensors - also on capture stream
|
||
let (_, d2udx2) = synthesize_exact_derivatives(&cfg, k_true);
|
||
let cached_u_r_t = Tensor::from_vec_on_stream(
|
||
u_real.iter().map(|&x| x as f32).collect(),
|
||
&[cfg.n_data, 1],
|
||
stream
|
||
)?;
|
||
let cached_u_i_t = Tensor::from_vec_on_stream(
|
||
u_imag.iter().map(|&x| x as f32).collect(),
|
||
&[cfg.n_data, 1],
|
||
stream
|
||
)?;
|
||
let cached_d2u_r_t = Tensor::from_vec_on_stream(
|
||
d2udx2.iter().map(|x| x[0] as f32).collect(),
|
||
&[cfg.n_data, 1],
|
||
stream
|
||
)?;
|
||
let cached_d2u_i_t = Tensor::from_vec_on_stream(
|
||
d2udx2.iter().map(|x| x[1] as f32).collect(),
|
||
&[cfg.n_data, 1],
|
||
stream
|
||
)?;
|
||
let k_sq = k_true * k_true;
|
||
let cached_k_sq_re = k_sq.re as f32;
|
||
let cached_k_sq_im = k_sq.im as f32;
|
||
|
||
let device = Device::cuda(0)?;
|
||
|
||
// === PHASE 8: Initialize zero-sync training infrastructure ===
|
||
let grad_workspace = GradientWorkspace::new(
|
||
cfg.n_data,
|
||
cfg.u_ff_dim,
|
||
cfg.u_hidden,
|
||
cfg.u_layers,
|
||
&device,
|
||
).map_err(|e| anyhow::anyhow!("Failed to create GradientWorkspace: {}", e))?;
|
||
|
||
// Collect parameter shapes for GpuAdam initialization
|
||
let mut param_shapes = Vec::new();
|
||
param_shapes.push(vec![1, cfg.u_ff_dim]); // B matrix
|
||
let mut in_dim = cfg.u_ff_dim * 2;
|
||
for i in 0..=cfg.u_layers {
|
||
let out_dim = if i == cfg.u_layers { 2 } else { cfg.u_hidden };
|
||
param_shapes.push(vec![out_dim, in_dim]); // Weight
|
||
param_shapes.push(vec![out_dim]); // Bias
|
||
in_dim = out_dim;
|
||
}
|
||
|
||
let gpu_adam = GpuAdam::new(
|
||
¶m_shapes,
|
||
&device,
|
||
cfg.lr as f32,
|
||
0.9, 0.999, 1e-8,
|
||
).map_err(|e| anyhow::anyhow!("Failed to create GpuAdam: {}", e))?;
|
||
|
||
// Create a second stream context for the solver's internal use
|
||
// (the passed ctx is returned for external use)
|
||
let solver_stream_ctx = PinnStreamContext::new(0)?;
|
||
|
||
let solver = Self {
|
||
omega,
|
||
k_true,
|
||
u_scale,
|
||
l,
|
||
u_net,
|
||
x_data,
|
||
u_data_target,
|
||
g_prime: cfg.g_prime_true,
|
||
g_double: cfg.g_double_true,
|
||
optimizer,
|
||
scheduler,
|
||
device,
|
||
cfg,
|
||
cached_u_r_t,
|
||
cached_u_i_t,
|
||
cached_d2u_r_t,
|
||
cached_d2u_i_t,
|
||
cached_k_sq_re,
|
||
cached_k_sq_im,
|
||
grad_workspace,
|
||
gpu_adam,
|
||
stream_ctx: solver_stream_ctx,
|
||
};
|
||
|
||
Ok((solver, ws, ctx))
|
||
}
|
||
|
||
/// Training step with CUDA graph acceleration.
|
||
///
|
||
/// This method uses a unified stream context and CUDA graph capture for
|
||
/// near-zero overhead forward passes. The graph is captured on the first
|
||
/// call and replayed on subsequent calls.
|
||
///
|
||
/// ## Performance
|
||
///
|
||
/// - Without graph: ~42µs (11 kernel launches × 4µs each)
|
||
/// - With graph: ~4µs (1 graph replay)
|
||
///
|
||
/// ## Stable Pointer Architecture
|
||
///
|
||
/// The optimizer MUST use in-place updates (`param.sub_()`) to preserve
|
||
/// memory addresses. With stable pointers, the graph NEVER needs recapturing
|
||
/// during training.
|
||
///
|
||
/// # Arguments
|
||
/// * `ws` - Pre-allocated workspace buffers
|
||
/// * `ctx` - Unified stream context (all ops use this stream)
|
||
/// * `graph` - CUDA graph wrapper (captures on first call)
|
||
#[cfg(feature = "cuda")]
|
||
pub fn training_step_with_cuda_graph(
|
||
&mut self,
|
||
ws: &mut ForwardWorkspace,
|
||
ctx: &PinnStreamContext,
|
||
graph: &mut PinnGraph,
|
||
) -> Result<(f64, f64, f64)> {
|
||
clear_tape();
|
||
|
||
// Forward pass: capture on first call, replay on subsequent
|
||
if graph.is_captured() {
|
||
// Fast path: replay captured graph (~4µs)
|
||
graph.launch().map_err(|e| anyhow::anyhow!("{}", e))?;
|
||
} else {
|
||
// First call: capture the forward pass into a graph
|
||
let input_shape: Vec<usize> = self.x_data.shape().dims().iter()
|
||
.map(|&d| d as usize)
|
||
.collect();
|
||
|
||
// Clone references for the closure (avoid borrow issues)
|
||
let x_data = &self.x_data;
|
||
let u_net = &self.u_net;
|
||
|
||
graph.capture(&input_shape, || {
|
||
u_net.forward_on_stream(x_data, ws, ctx)
|
||
.map_err(|e| rtx_tensor::TensorError::runtime(format!("{}", e)))
|
||
}).map_err(|e| anyhow::anyhow!("{}", e))?;
|
||
}
|
||
|
||
// Synchronize before reading output (required for CPU-side loss computation)
|
||
ctx.synchronize().map_err(|e| anyhow::anyhow!("{}", e))?;
|
||
|
||
// Compute losses (still requires GPU→CPU sync for scheduler)
|
||
let loss_data = Self::mse_loss(&ws.output, &self.u_data_target)?;
|
||
let (mse_rr, mse_ri) = self.compute_pde_residual_cached()?;
|
||
let loss_pde = mse_rr + mse_ri;
|
||
let total_loss = self.cfg.data_weight * loss_data + self.cfg.pde_weight * loss_pde;
|
||
|
||
let _new_lr = self.scheduler.step_metric(loss_data);
|
||
|
||
// Note: Optimizer step happens outside this function.
|
||
// With in-place optimizer updates, the graph remains valid.
|
||
|
||
Ok((loss_data, loss_pde, total_loss))
|
||
}
|
||
|
||
/// Run a single training step using pre-allocated workspace (optimized)
|
||
///
|
||
/// This method uses `forward_with_workspace` to minimize allocations during
|
||
/// the forward pass, which is the primary bottleneck in training.
|
||
/// NOTE: Still uses unoptimized compute_pde_residual_tensor()
|
||
pub fn training_step_with_workspace(&mut self, ws: &mut ForwardWorkspace) -> Result<(f64, f64, f64)> {
|
||
clear_tape();
|
||
// Use optimized forward pass that writes to pre-allocated buffers
|
||
self.u_net.forward_with_workspace(&self.x_data, ws)?;
|
||
// Read from workspace output buffer
|
||
let loss_data = Self::mse_loss(&ws.output, &self.u_data_target)?;
|
||
let (mse_rr, mse_ri) = self.compute_pde_residual_tensor()?;
|
||
let loss_pde = mse_rr + mse_ri;
|
||
let total_loss = self.cfg.data_weight * loss_data + self.cfg.pde_weight * loss_pde;
|
||
let _new_lr = self.scheduler.step_metric(loss_data);
|
||
Ok((loss_data, loss_pde, total_loss))
|
||
}
|
||
|
||
/// FULLY OPTIMIZED training step: workspace + cached PDE tensors
|
||
///
|
||
/// Combines both optimizations:
|
||
/// 1. `forward_with_workspace` - zero-alloc forward pass (9x faster)
|
||
/// 2. `compute_pde_residual_cached` - uses pre-computed tensors
|
||
///
|
||
/// This is the fastest possible training step.
|
||
pub fn training_step_fully_optimized(&mut self, ws: &mut ForwardWorkspace) -> Result<(f64, f64, f64)> {
|
||
clear_tape();
|
||
// Use optimized forward pass that writes to pre-allocated buffers
|
||
self.u_net.forward_with_workspace(&self.x_data, ws)?;
|
||
// Read from workspace output buffer
|
||
let loss_data = Self::mse_loss(&ws.output, &self.u_data_target)?;
|
||
let (mse_rr, mse_ri) = self.compute_pde_residual_cached()?; // Use cached PDE!
|
||
let loss_pde = mse_rr + mse_ri;
|
||
let total_loss = self.cfg.data_weight * loss_data + self.cfg.pde_weight * loss_pde;
|
||
let _new_lr = self.scheduler.step_metric(loss_data);
|
||
Ok((loss_data, loss_pde, total_loss))
|
||
}
|
||
|
||
/// ULTIMATE OPTIMIZED training step: CUDA graph + workspace + cached PDE
|
||
///
|
||
/// This is the fastest possible training step on CUDA devices.
|
||
/// Combines all optimizations:
|
||
/// 1. CUDA graph capture - eliminates kernel launch overhead
|
||
/// 2. `forward_with_workspace` - zero-alloc forward pass
|
||
/// 3. `compute_pde_residual_cached` - uses pre-computed tensors
|
||
///
|
||
/// The forward pass is captured into a CUDA graph on the first call,
|
||
/// then replayed with minimal CPU overhead on subsequent calls.
|
||
#[cfg(feature = "cuda")]
|
||
pub fn training_step_with_graph(
|
||
&mut self,
|
||
ws: &mut ForwardWorkspace,
|
||
graph: &mut CapturedForwardGraph,
|
||
) -> Result<(f64, f64, f64)> {
|
||
clear_tape();
|
||
// Use CUDA graph-accelerated forward pass
|
||
self.u_net.forward_with_graph(&self.x_data, ws, graph)?;
|
||
// Read from workspace output buffer
|
||
let loss_data = Self::mse_loss(&ws.output, &self.u_data_target)?;
|
||
let (mse_rr, mse_ri) = self.compute_pde_residual_cached()?;
|
||
let loss_pde = mse_rr + mse_ri;
|
||
let total_loss = self.cfg.data_weight * loss_data + self.cfg.pde_weight * loss_pde;
|
||
let _new_lr = self.scheduler.step_metric(loss_data);
|
||
Ok((loss_data, loss_pde, total_loss))
|
||
}
|
||
|
||
// =========================================================================
|
||
// ZERO-SYNC TRAINING (Eliminates GPU-CPU synchronization overhead)
|
||
// =========================================================================
|
||
|
||
/// DEFERRED-LOSS training step: Only compute loss when needed.
|
||
///
|
||
/// This is a simpler approach than GPU accumulation - just skip loss
|
||
/// computation entirely for most steps, only computing when the scheduler
|
||
/// needs an update.
|
||
///
|
||
/// # Performance Impact
|
||
/// - Eliminates ~35-55µs from loss computation on most steps
|
||
/// - Only computes loss every `sync_interval` steps
|
||
///
|
||
/// # Arguments
|
||
/// * `ws` - Pre-allocated forward pass workspace
|
||
/// * `step` - Current step number (for deciding when to sync)
|
||
/// * `sync_interval` - Compute loss every N steps
|
||
///
|
||
/// # Returns
|
||
/// Some((data_loss, pde_loss, total_loss)) if loss was computed, None otherwise
|
||
pub fn training_step_deferred_loss(
|
||
&mut self,
|
||
ws: &mut ForwardWorkspace,
|
||
step: usize,
|
||
sync_interval: usize,
|
||
) -> Result<Option<(f64, f64, f64)>> {
|
||
clear_tape();
|
||
|
||
// Forward pass (always needed for gradient computation)
|
||
self.u_net.forward_with_workspace(&self.x_data, ws)?;
|
||
|
||
// Only compute loss when needed for scheduler
|
||
if step % sync_interval == 0 {
|
||
// Compute loss with CPU sync (same as fully_optimized)
|
||
let loss_data = Self::mse_loss(&ws.output, &self.u_data_target)?;
|
||
let (mse_rr, mse_ri) = self.compute_pde_residual_cached()?;
|
||
let loss_pde = mse_rr + mse_ri;
|
||
let total_loss = self.cfg.data_weight * loss_data + self.cfg.pde_weight * loss_pde;
|
||
let _new_lr = self.scheduler.step_metric(loss_data);
|
||
Ok(Some((loss_data, loss_pde, total_loss)))
|
||
} else {
|
||
// Skip loss computation entirely - just forward pass
|
||
Ok(None)
|
||
}
|
||
}
|
||
|
||
/// ZERO-SYNC training step (with GPU accumulation).
|
||
///
|
||
/// NOTE: Benchmarking showed this is SLOWER than the original due to
|
||
/// accumulator overhead. Use `training_step_deferred_loss` instead.
|
||
///
|
||
/// This version keeps loss tensors on GPU and accumulates them,
|
||
/// but the extra GPU operations cost more than the to_cpu() savings.
|
||
///
|
||
/// # Arguments
|
||
/// * `ws` - Pre-allocated forward pass workspace
|
||
/// * `accumulator` - GPU-resident loss accumulator
|
||
pub fn training_step_zero_sync(
|
||
&mut self,
|
||
ws: &mut ForwardWorkspace,
|
||
accumulator: &mut GpuLossAccumulator,
|
||
) -> Result<()> {
|
||
clear_tape();
|
||
|
||
// Forward pass (already optimized - no sync)
|
||
self.u_net.forward_with_workspace(&self.x_data, ws)?;
|
||
|
||
// Compute losses on GPU (NO to_cpu() calls!)
|
||
let data_loss_gpu = Self::mse_loss_gpu(&ws.output, &self.u_data_target)?;
|
||
let (pde_re_gpu, pde_im_gpu) = self.compute_pde_residual_gpu()?;
|
||
|
||
// Accumulate on GPU (no sync)
|
||
accumulator.accumulate(
|
||
&data_loss_gpu,
|
||
&pde_re_gpu,
|
||
&pde_im_gpu,
|
||
self.cfg.data_weight as f32,
|
||
self.cfg.pde_weight as f32,
|
||
)?;
|
||
|
||
// Conditionally sync and update scheduler (only every sync_interval steps)
|
||
if let Some((data_loss, _pde_loss, _total_loss)) = accumulator.maybe_sync()? {
|
||
let _new_lr = self.scheduler.step_metric(data_loss);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// DEFERRED-LOSS training loop: maximum performance by skipping loss on most steps.
|
||
///
|
||
/// Only computes loss every `sync_interval` steps, saving ~35-55µs per skipped step.
|
||
/// The scheduler is only updated when loss is computed.
|
||
///
|
||
/// # Arguments
|
||
/// * `sync_interval` - Compute loss every N steps (100 recommended for patience=500)
|
||
///
|
||
/// # Returns
|
||
/// Final total loss value (from last computed loss)
|
||
pub fn train_deferred_loss(&mut self, sync_interval: usize) -> Result<f64> {
|
||
let mut ws = self.create_workspace()?;
|
||
let mut last_loss = 0.0f64;
|
||
|
||
for epoch in 1..=self.cfg.epochs {
|
||
if let Some((_, _, total_loss)) = self.training_step_deferred_loss(&mut ws, epoch, sync_interval)? {
|
||
last_loss = total_loss;
|
||
}
|
||
}
|
||
|
||
Ok(last_loss)
|
||
}
|
||
|
||
/// ZERO-SYNC training loop (with GPU accumulation).
|
||
///
|
||
/// NOTE: Benchmarking showed this is SLOWER than the original due to
|
||
/// accumulator overhead. Use `train_deferred_loss` instead.
|
||
///
|
||
/// # Arguments
|
||
/// * `sync_interval` - Sync to CPU every N steps (100 recommended for patience=500)
|
||
///
|
||
/// # Returns
|
||
/// Final total loss value
|
||
pub fn train_zero_sync(&mut self, sync_interval: usize) -> Result<f64> {
|
||
let mut ws = self.create_workspace()?;
|
||
let mut accumulator = GpuLossAccumulator::new(&self.device, sync_interval)?;
|
||
|
||
for _epoch in 1..=self.cfg.epochs {
|
||
self.training_step_zero_sync(&mut ws, &mut accumulator)?;
|
||
}
|
||
|
||
// Final sync to get loss values
|
||
let (_, _, total_loss) = accumulator.force_sync()?.unwrap_or((0.0, 0.0, 0.0));
|
||
Ok(total_loss)
|
||
}
|
||
|
||
/// ZERO-SYNC training with progress logging.
|
||
///
|
||
/// Same as `train_zero_sync` but logs progress at specified intervals.
|
||
/// Only syncs when logging is needed, minimizing overhead.
|
||
pub fn train_zero_sync_with_logging(&mut self, sync_interval: usize) -> Result<f64> {
|
||
let mut ws = self.create_workspace()?;
|
||
let mut accumulator = GpuLossAccumulator::new(&self.device, sync_interval)?;
|
||
let mut best_loss = f64::MAX;
|
||
|
||
for epoch in 1..=self.cfg.epochs {
|
||
self.training_step_zero_sync(&mut ws, &mut accumulator)?;
|
||
|
||
// Log progress at print_every intervals (requires sync)
|
||
if epoch % self.cfg.print_every == 0 {
|
||
if let Some((data_loss, pde_loss, total_loss)) = accumulator.force_sync()? {
|
||
if total_loss < best_loss {
|
||
best_loss = total_loss;
|
||
}
|
||
// Logging could go here if needed
|
||
let _ = (data_loss, pde_loss); // Suppress unused warnings
|
||
}
|
||
}
|
||
}
|
||
|
||
// Final sync
|
||
if let Some((_, _, total_loss)) = accumulator.force_sync()? {
|
||
if total_loss < best_loss {
|
||
best_loss = total_loss;
|
||
}
|
||
}
|
||
|
||
Ok(best_loss)
|
||
}
|
||
|
||
/// Run training loop (unoptimized - uses compute_pde_residual_tensor)
|
||
pub fn train(&mut self) -> Result<f64> {
|
||
let mut best_loss = f64::MAX;
|
||
for _epoch in 1..=self.cfg.epochs {
|
||
let (_, _, total_loss) = self.training_step()?;
|
||
if total_loss < best_loss {
|
||
best_loss = total_loss;
|
||
}
|
||
}
|
||
Ok(best_loss)
|
||
}
|
||
|
||
/// Run optimized training loop using CACHED PDE tensors
|
||
///
|
||
/// This is the recommended training method. Uses pre-computed tensors
|
||
/// to eliminate ~89% overhead from repeated allocations.
|
||
///
|
||
/// Expected performance: ~35-40ms for 100 epochs (vs ~323ms unoptimized)
|
||
pub fn train_cached(&mut self) -> Result<f64> {
|
||
let mut best_loss = f64::MAX;
|
||
for _epoch in 1..=self.cfg.epochs {
|
||
let (_, _, total_loss) = self.training_step_cached()?;
|
||
if total_loss < best_loss {
|
||
best_loss = total_loss;
|
||
}
|
||
}
|
||
Ok(best_loss)
|
||
}
|
||
|
||
/// Run optimized training loop using workspace buffers
|
||
pub fn train_with_workspace(&mut self) -> Result<f64> {
|
||
let mut ws = self.create_workspace()?;
|
||
let mut best_loss = f64::MAX;
|
||
for _epoch in 1..=self.cfg.epochs {
|
||
let (_, _, total_loss) = self.training_step_with_workspace(&mut ws)?;
|
||
if total_loss < best_loss {
|
||
best_loss = total_loss;
|
||
}
|
||
}
|
||
Ok(best_loss)
|
||
}
|
||
|
||
/// FULLY OPTIMIZED training loop: workspace + cached PDE tensors
|
||
///
|
||
/// This is the recommended training method for maximum performance.
|
||
/// Combines workspace-based forward pass + cached PDE tensors.
|
||
///
|
||
/// Expected performance: ~40ms for 100 epochs (vs ~380ms unoptimized = 9.5x faster)
|
||
pub fn train_fully_optimized(&mut self) -> Result<f64> {
|
||
let mut ws = self.create_workspace()?;
|
||
let mut best_loss = f64::MAX;
|
||
for _epoch in 1..=self.cfg.epochs {
|
||
let (_, _, total_loss) = self.training_step_fully_optimized(&mut ws)?;
|
||
if total_loss < best_loss {
|
||
best_loss = total_loss;
|
||
}
|
||
}
|
||
Ok(best_loss)
|
||
}
|
||
|
||
/// ULTIMATE OPTIMIZED training loop: CUDA graph + workspace + cached PDE
|
||
///
|
||
/// This is the fastest training method on CUDA devices. The forward pass
|
||
/// is captured into a CUDA graph on the first iteration, then replayed
|
||
/// with minimal CPU overhead on subsequent iterations.
|
||
///
|
||
/// Expected performance: Forward pass kernel launch overhead reduced from
|
||
/// ~50-100µs to ~5µs (10-20x improvement for small batches).
|
||
#[cfg(feature = "cuda")]
|
||
pub fn train_with_graph(&mut self) -> Result<f64> {
|
||
let mut ws = self.create_workspace()?;
|
||
let mut graph = CapturedForwardGraph::new();
|
||
let mut best_loss = f64::MAX;
|
||
for _epoch in 1..=self.cfg.epochs {
|
||
let (_, _, total_loss) = self.training_step_with_graph(&mut ws, &mut graph)?;
|
||
if total_loss < best_loss {
|
||
best_loss = total_loss;
|
||
}
|
||
}
|
||
// Log graph statistics
|
||
if graph.is_captured() {
|
||
eprintln!("CUDA graph launched {} times", graph.launch_count());
|
||
}
|
||
Ok(best_loss)
|
||
}
|
||
|
||
// =========================================================================
|
||
// PHASE 8: ANALYTICAL BACKPROP TRAINING (Zero-Sync, Zero-Autograd)
|
||
// =========================================================================
|
||
|
||
/// Single training step with analytical backpropagation.
|
||
///
|
||
/// This method implements a complete training step WITHOUT autograd:
|
||
/// 1. Forward pass - caches intermediate activations in grad_workspace
|
||
/// 2. Loss gradient - computes dL/d_output for MSE loss
|
||
/// 3. Backward pass - propagates gradients layer by layer
|
||
/// 4. Optimizer step - updates all parameters in-place on GPU
|
||
///
|
||
/// ## Performance
|
||
/// - No CPU-GPU sync during the step
|
||
/// - No autograd tape construction
|
||
/// - All buffers pre-allocated (zero heap allocation)
|
||
///
|
||
/// ## Returns
|
||
/// Nothing - all data stays on GPU for zero-sync operation.
|
||
/// Use `compute_data_loss_for_logging()` periodically to check progress.
|
||
#[cfg(feature = "cuda")]
|
||
pub fn train_step_data_only(&mut self) -> Result<()> {
|
||
// --- 1. Forward Pass (Caching Activations) ---
|
||
self.u_net.forward_training(&self.x_data, &mut self.grad_workspace)?;
|
||
|
||
// --- 2. Loss Gradient (dL/d_output) - GPU native, zero-alloc ---
|
||
let num_layers = self.u_net.layers().len();
|
||
let u_pred = self.grad_workspace.h[num_layers - 1].clone();
|
||
mse_backward_inplace(
|
||
&u_pred,
|
||
&self.u_data_target,
|
||
&self.stream_ctx,
|
||
&mut self.grad_workspace.d_loss,
|
||
).map_err(|e| anyhow::anyhow!("mse_backward_inplace failed: {}", e))?;
|
||
|
||
// --- 3. Backward Pass - GPU native, zero-alloc ---
|
||
// Start with d_loss, propagate backward through layers
|
||
//
|
||
// For the last layer (output), dL_dh = d_loss
|
||
// For each layer i: layer_backward_inplace writes:
|
||
// - dW[i], db[i] to workspace
|
||
// - dL_dh_prev[i] as the gradient to pass to layer i-1
|
||
|
||
// Copy d_loss to dL_dh_prev[num_layers-1] as starting point
|
||
// Then we iterate backward, using dL_dh_prev[i+1] as input to layer i
|
||
self.stream_ctx.mul_scalar_out(
|
||
&self.grad_workspace.d_loss,
|
||
1.0,
|
||
&mut self.grad_workspace.dL_dh_prev[num_layers - 1],
|
||
).map_err(|e| anyhow::anyhow!("copy d_loss failed: {}", e))?;
|
||
|
||
for i in (0..num_layers).rev() {
|
||
// Get input gradient (from next layer or d_loss for output layer)
|
||
// After first iteration, use dL_dh_prev[i+1] from prev layer's backward
|
||
// For output layer (i = num_layers-1), we just copied d_loss there
|
||
|
||
// Get input to this layer (h_prev)
|
||
let h_prev = if i == 0 {
|
||
self.grad_workspace.fourier_features.clone()
|
||
} else {
|
||
self.grad_workspace.h[i - 1].clone()
|
||
};
|
||
|
||
// Get cached activation
|
||
let h_i = self.grad_workspace.h[i].clone();
|
||
|
||
// Get dL_dh for this layer
|
||
let dL_dh = if i == num_layers - 1 {
|
||
// Output layer: use d_loss (already in dL_dh_prev[num_layers-1])
|
||
self.grad_workspace.dL_dh_prev[num_layers - 1].clone()
|
||
} else {
|
||
// Hidden layer: use output from next layer's backward
|
||
self.grad_workspace.dL_dh_prev[i + 1].clone()
|
||
};
|
||
|
||
let apply_tanh = i < num_layers - 1;
|
||
let w = self.u_net.layers()[i].weight().clone();
|
||
|
||
layer_backward_inplace(
|
||
&dL_dh,
|
||
&h_i,
|
||
&h_prev,
|
||
&w,
|
||
i,
|
||
apply_tanh,
|
||
&self.stream_ctx,
|
||
&mut self.grad_workspace,
|
||
).map_err(|e| anyhow::anyhow!("layer_backward_inplace failed for layer {}: {}", i, e))?;
|
||
}
|
||
|
||
// --- 4. Fourier Backward: d_h -> dB ---
|
||
// dL_dh_prev[0] now contains gradient w.r.t. fourier features
|
||
// Use fused GPU kernel for zero-allocation backward pass
|
||
fourier_backward_inplace(
|
||
&self.grad_workspace.dL_dh_prev[0],
|
||
&self.x_data,
|
||
self.u_net.b_learnable(),
|
||
&self.stream_ctx,
|
||
&mut self.grad_workspace.dB,
|
||
).map_err(|e| anyhow::anyhow!("fourier_backward_inplace failed: {}", e))?;
|
||
|
||
// --- 5. Optimizer Step - FUSED ---
|
||
// Update parameters in-place using fused kernel
|
||
self.optimizer_step_fused()?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Perform fused optimizer step updating parameters in-place.
|
||
#[cfg(feature = "cuda")]
|
||
fn optimizer_step_fused(&mut self) -> Result<()> {
|
||
self.gpu_adam.t += 1;
|
||
let t = self.gpu_adam.t;
|
||
|
||
// Compute bias corrections
|
||
let bias_correction1: f32 = 1.0 - self.gpu_adam.beta1.powi(t as i32);
|
||
let bias_correction2: f32 = 1.0 - self.gpu_adam.beta2.powi(t as i32);
|
||
let lr_adj = self.gpu_adam.lr * bias_correction2.sqrt() / bias_correction1;
|
||
|
||
// Update B matrix (index 0)
|
||
self.stream_ctx.adam_step(
|
||
self.u_net.b_learnable_mut(),
|
||
&mut self.gpu_adam.m[0],
|
||
&mut self.gpu_adam.v[0],
|
||
&self.grad_workspace.dB,
|
||
lr_adj,
|
||
self.gpu_adam.beta1,
|
||
self.gpu_adam.beta2,
|
||
self.gpu_adam.eps,
|
||
).map_err(|e| anyhow::anyhow!("adam_step for B failed: {}", e))?;
|
||
|
||
// Update layer weights and biases
|
||
let num_layers = self.u_net.layers().len();
|
||
let mut param_idx = 1; // Start after B
|
||
|
||
for i in 0..num_layers {
|
||
// Weight
|
||
let layer = &mut self.u_net.layers_mut()[i];
|
||
self.stream_ctx.adam_step(
|
||
layer.weight_mut(),
|
||
&mut self.gpu_adam.m[param_idx],
|
||
&mut self.gpu_adam.v[param_idx],
|
||
&self.grad_workspace.dW[i],
|
||
lr_adj,
|
||
self.gpu_adam.beta1,
|
||
self.gpu_adam.beta2,
|
||
self.gpu_adam.eps,
|
||
).map_err(|e| anyhow::anyhow!("adam_step for W{} failed: {}", i, e))?;
|
||
param_idx += 1;
|
||
|
||
// Bias
|
||
self.stream_ctx.adam_step(
|
||
layer.bias_mut().expect("Layer should have bias"),
|
||
&mut self.gpu_adam.m[param_idx],
|
||
&mut self.gpu_adam.v[param_idx],
|
||
&self.grad_workspace.db[i],
|
||
lr_adj,
|
||
self.gpu_adam.beta1,
|
||
self.gpu_adam.beta2,
|
||
self.gpu_adam.eps,
|
||
).map_err(|e| anyhow::anyhow!("adam_step for b{} failed: {}", i, e))?;
|
||
param_idx += 1;
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// CPU fallback version (allocating, slower)
|
||
#[cfg(not(feature = "cuda"))]
|
||
pub fn train_step_data_only(&mut self) -> Result<()> {
|
||
// --- 1. Forward Pass (Caching Activations) ---
|
||
self.u_net.forward_training(&self.x_data, &mut self.grad_workspace)?;
|
||
|
||
// --- 2. Loss Gradient (dL/d_output) ---
|
||
let num_layers = self.u_net.layers().len();
|
||
let u_pred = &self.grad_workspace.h[num_layers - 1];
|
||
let d_loss = mse_backward(u_pred, &self.u_data_target)
|
||
.map_err(|e| anyhow::anyhow!("mse_backward failed: {}", e))?;
|
||
|
||
// --- 3. Backward Pass (Propagate Gradients Layer by Layer) ---
|
||
let mut d_h = d_loss;
|
||
|
||
for i in (0..num_layers).rev() {
|
||
let h_prev = if i == 0 {
|
||
&self.grad_workspace.fourier_features
|
||
} else {
|
||
&self.grad_workspace.h[i - 1]
|
||
};
|
||
|
||
let apply_tanh = i < num_layers - 1;
|
||
|
||
let (d_w, d_b, dh_prev) = layer_backward(
|
||
&d_h,
|
||
&self.grad_workspace.h[i],
|
||
h_prev,
|
||
self.u_net.layers()[i].weight(),
|
||
apply_tanh,
|
||
).map_err(|e| anyhow::anyhow!("layer_backward failed: {}", e))?;
|
||
|
||
self.grad_workspace.dW[i] = d_w;
|
||
self.grad_workspace.db[i] = d_b;
|
||
d_h = dh_prev;
|
||
}
|
||
|
||
// --- 4. Fourier Backward: d_h -> dB ---
|
||
self.grad_workspace.dB = fourier_backward(
|
||
&d_h,
|
||
&self.x_data,
|
||
self.u_net.b_learnable(),
|
||
).map_err(|e| anyhow::anyhow!("fourier_backward failed: {}", e))?;
|
||
|
||
// --- 5. Optimizer Step ---
|
||
let (mut params, grads) = self.collect_params_and_grads_cloned();
|
||
self.gpu_adam.step(&mut params, &grads)
|
||
.map_err(|e| anyhow::anyhow!("optimizer step failed: {}", e))?;
|
||
|
||
self.write_params_back(¶ms)?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Collect cloned parameters and gradients for optimizer step.
|
||
///
|
||
/// Returns: (cloned parameters, gradients)
|
||
/// Order: B, then for each layer: (W, b)
|
||
fn collect_params_and_grads_cloned(&self) -> (Vec<Tensor>, Vec<Tensor>) {
|
||
let mut params = Vec::new();
|
||
let mut grads = Vec::new();
|
||
|
||
// B matrix
|
||
params.push(self.u_net.b_learnable().clone());
|
||
grads.push(self.grad_workspace.dB.clone());
|
||
|
||
// Layer weights and biases
|
||
for i in 0..self.u_net.layers().len() {
|
||
params.push(self.u_net.layers()[i].weight().clone());
|
||
grads.push(self.grad_workspace.dW[i].clone());
|
||
|
||
if let Some(bias) = self.u_net.layers()[i].bias() {
|
||
params.push(bias.clone());
|
||
grads.push(self.grad_workspace.db[i].clone());
|
||
}
|
||
}
|
||
|
||
(params, grads)
|
||
}
|
||
|
||
/// Write updated parameters back to the network after optimizer step.
|
||
fn write_params_back(&mut self, params: &[Tensor]) -> Result<()> {
|
||
let mut idx = 0;
|
||
|
||
// B matrix (first parameter)
|
||
*self.u_net.b_learnable_mut() = params[idx].clone();
|
||
idx += 1;
|
||
|
||
// Layer weights and biases
|
||
let num_layers = self.u_net.layers().len();
|
||
for i in 0..num_layers {
|
||
// Weight
|
||
self.u_net.layers_mut()[i].set_weight(params[idx].clone())
|
||
.map_err(|e| anyhow::anyhow!("Failed to set weight for layer {}: {}", i, e))?;
|
||
idx += 1;
|
||
|
||
// Bias (if present)
|
||
if self.u_net.layers()[i].bias().is_some() {
|
||
self.u_net.layers_mut()[i].set_bias(params[idx].clone())
|
||
.map_err(|e| anyhow::anyhow!("Failed to set bias for layer {}: {}", i, e))?;
|
||
idx += 1;
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Compute data loss for logging/debugging (requires GPU-CPU sync).
|
||
///
|
||
/// Only call this periodically (e.g., every 100 steps) to minimize sync overhead.
|
||
pub fn compute_data_loss_for_logging(&self) -> Result<f64> {
|
||
let num_layers = self.u_net.layers().len();
|
||
let u_pred = &self.grad_workspace.h[num_layers - 1];
|
||
Self::mse_loss(u_pred, &self.u_data_target)
|
||
}
|
||
|
||
/// Train for N steps using analytical backprop, returning final loss.
|
||
///
|
||
/// This is the recommended training method for Phase 8.
|
||
/// - No autograd overhead
|
||
/// - No CPU-GPU sync except at the end for loss reporting
|
||
///
|
||
/// # Arguments
|
||
/// * `n_steps` - Number of training steps
|
||
/// * `log_interval` - How often to compute and log loss (0 = never until end)
|
||
pub fn train_analytical(&mut self, n_steps: usize, log_interval: usize) -> Result<f64> {
|
||
let mut best_loss = f64::MAX;
|
||
|
||
for step in 1..=n_steps {
|
||
self.train_step_data_only()?;
|
||
|
||
// Optionally log progress
|
||
if log_interval > 0 && step % log_interval == 0 {
|
||
let loss = self.compute_data_loss_for_logging()?;
|
||
if loss < best_loss {
|
||
best_loss = loss;
|
||
}
|
||
if step % (log_interval * 10) == 0 {
|
||
eprintln!("Step {}: loss = {:.6e}", step, loss);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Always compute final loss
|
||
let final_loss = self.compute_data_loss_for_logging()?;
|
||
if final_loss < best_loss {
|
||
best_loss = final_loss;
|
||
}
|
||
|
||
Ok(best_loss)
|
||
}
|
||
}
|