419 lines
14 KiB
Rust
419 lines
14 KiB
Rust
//! CUDA Graph Capture and Replay for PINN Forward Pass
|
|
//!
|
|
//! This module provides graph capture functionality that records the forward pass
|
|
//! into a CUDA graph for near-zero-overhead replay during training.
|
|
//!
|
|
//! ## Key Insight: Stable Pointers
|
|
//!
|
|
//! CUDA Graphs capture POINTERS, not VALUES. The graph records:
|
|
//! "Launch Kernel A, reading from Address X, writing to Address Y."
|
|
//!
|
|
//! This means:
|
|
//! - **Wrong**: `weights = weights - lr * grad` (allocates new memory → graph breaks)
|
|
//! - **Right**: `weights.sub_assign_(lr * grad)` (in-place update → graph valid)
|
|
//!
|
|
//! With proper in-place optimizer updates, the graph NEVER needs recapturing
|
|
//! during training.
|
|
//!
|
|
//! ## Usage
|
|
//!
|
|
//! ```ignore
|
|
//! let ctx = Arc::new(PinnStreamContext::new(0)?);
|
|
//! let mut graph = PinnGraph::new(ctx.clone());
|
|
//!
|
|
//! // Training loop - graph captured once, replayed 49,999 times
|
|
//! for epoch in 1..=50_000 {
|
|
//! if graph.is_captured() {
|
|
//! graph.launch()?; // Fast: ~4µs
|
|
//! } else {
|
|
//! graph.capture(|| solver.forward_on_stream(&x, &mut ws, &ctx))?;
|
|
//! }
|
|
//!
|
|
//! // Optimizer uses in-place updates - graph remains valid!
|
|
//! optimizer.step_inplace(&mut params)?;
|
|
//! }
|
|
//! ```
|
|
//!
|
|
//! ## When to Invalidate
|
|
//!
|
|
//! Only invalidate the graph for rare structural changes:
|
|
//! - Batch size change (workspace buffer addresses change)
|
|
//! - Model topology change (add/remove layers)
|
|
//! - Manual buffer reallocation
|
|
|
|
use rtx_tensor::{TensorError, Result};
|
|
use std::sync::Arc;
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use super::cuda_stream_context::PinnStreamContext;
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use cudarc::driver::safe::CudaGraph;
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use cudarc::driver::sys::{CUstreamCaptureMode, CUgraphInstantiate_flags_enum};
|
|
|
|
// =============================================================================
|
|
// PINN GRAPH
|
|
// =============================================================================
|
|
|
|
/// CUDA graph for accelerated PINN forward pass.
|
|
///
|
|
/// Captures the forward pass operations into a graph that can be replayed
|
|
/// with minimal CPU overhead (~4µs vs ~44µs for 11 kernel launches).
|
|
#[cfg(feature = "cuda")]
|
|
pub struct PinnGraph {
|
|
/// The instantiated graph (ready for replay)
|
|
graph: Option<CudaGraph>,
|
|
/// Stream context used for capture and launch
|
|
ctx: Arc<PinnStreamContext>,
|
|
/// Input shape (graphs are shape-specific)
|
|
input_shape: Vec<usize>,
|
|
/// Launch count for statistics
|
|
launch_count: u64,
|
|
/// Whether capture is in progress
|
|
capturing: bool,
|
|
}
|
|
|
|
#[cfg(feature = "cuda")]
|
|
impl PinnGraph {
|
|
/// Create a new (uncaptured) graph wrapper.
|
|
///
|
|
/// # Arguments
|
|
/// * `ctx` - The unified stream context for all operations
|
|
pub fn new(ctx: Arc<PinnStreamContext>) -> Self {
|
|
Self {
|
|
graph: None,
|
|
ctx,
|
|
input_shape: Vec::new(),
|
|
launch_count: 0,
|
|
capturing: false,
|
|
}
|
|
}
|
|
|
|
/// Check if the graph has been captured.
|
|
pub fn is_captured(&self) -> bool {
|
|
self.graph.is_some()
|
|
}
|
|
|
|
/// Get the number of times this graph has been launched.
|
|
pub fn launch_count(&self) -> u64 {
|
|
self.launch_count
|
|
}
|
|
|
|
/// Get the input shape the graph was captured with.
|
|
pub fn input_shape(&self) -> &[usize] {
|
|
&self.input_shape
|
|
}
|
|
|
|
/// Capture operations into a CUDA graph.
|
|
///
|
|
/// The closure `capture_fn` should execute the forward pass operations.
|
|
/// All operations must use the same stream (via PinnStreamContext).
|
|
///
|
|
/// # Arguments
|
|
/// * `input_shape` - Shape of the input tensor (for validation)
|
|
/// * `capture_fn` - Closure that executes the operations to capture
|
|
///
|
|
/// # Example
|
|
/// ```ignore
|
|
/// graph.capture(&x.shape().dims(), || {
|
|
/// solver.u_net.forward_on_stream(&x, &mut ws, &ctx)
|
|
/// })?;
|
|
/// ```
|
|
pub fn capture<F>(&mut self, input_shape: &[usize], capture_fn: F) -> Result<()>
|
|
where
|
|
F: FnOnce() -> Result<()>,
|
|
{
|
|
// Check if shape changed
|
|
if self.is_captured() && self.input_shape != input_shape {
|
|
self.invalidate();
|
|
}
|
|
|
|
// Already captured with same shape
|
|
if self.is_captured() {
|
|
return Ok(());
|
|
}
|
|
|
|
// CRITICAL: Full device synchronization BEFORE capture.
|
|
// This ensures all tensor allocations and operations on ALL streams are complete.
|
|
// Without this, CUDA graph capture fails with STREAM_CAPTURE_ISOLATION because
|
|
// the capture stream doesn't "own" the tensor data that was written by other streams.
|
|
//
|
|
// Note: This is a CPU-blocking operation, but it only happens once during the
|
|
// first capture. Subsequent replays are ultra-fast (~4µs).
|
|
self.ctx.device_synchronize()?;
|
|
|
|
// CRITICAL: Re-establish stream dependency edge AFTER all tensor allocations.
|
|
// The join() in PinnStreamContext::new() may have happened before the solver/model
|
|
// was created. By joining again here, we ensure the capture stream has dependency
|
|
// edges to ALL tensor writes, including model weights created after context init.
|
|
//
|
|
// Without this, even with device_synchronize(), CUDA graph capture fails because
|
|
// the capture stream's dependency graph doesn't include the model weight allocations.
|
|
self.ctx.join_with_default_stream()?;
|
|
|
|
// Begin stream capture in GLOBAL mode
|
|
// GLOBAL mode allows all operations from any stream to be captured into the graph.
|
|
// This is necessary when input tensors were allocated on a different stream.
|
|
// With our singleton context, all memory is in the same virtual address space.
|
|
self.capturing = true;
|
|
self.ctx.stream()
|
|
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_GLOBAL)
|
|
.map_err(|e| TensorError::device(format!("Failed to begin graph capture: {:?}", e)))?;
|
|
|
|
// Execute the forward pass (operations are recorded, not executed)
|
|
let result = capture_fn();
|
|
|
|
// End capture regardless of result
|
|
let graph_result = self.ctx.stream()
|
|
.end_capture(CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
|
|
|
|
self.capturing = false;
|
|
|
|
// Check if capture_fn succeeded
|
|
result?;
|
|
|
|
// Check if end_capture succeeded
|
|
let graph = graph_result
|
|
.map_err(|e| TensorError::device(format!("Failed to end graph capture: {:?}", e)))?
|
|
.ok_or_else(|| TensorError::device("Graph capture returned empty (no operations captured)"))?;
|
|
|
|
// Store the captured graph
|
|
self.graph = Some(graph);
|
|
self.input_shape = input_shape.to_vec();
|
|
|
|
#[cfg(debug_assertions)]
|
|
eprintln!(
|
|
"[PinnGraph] CUDA graph captured for input shape {:?}",
|
|
self.input_shape
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Capture and immediately execute (for first iteration).
|
|
///
|
|
/// This is a convenience method that captures the graph and then
|
|
/// executes it, ensuring the first forward pass produces correct output.
|
|
///
|
|
/// # Arguments
|
|
/// * `input_shape` - Shape of the input tensor
|
|
/// * `capture_fn` - Closure that executes the operations to capture
|
|
pub fn capture_and_launch<F>(&mut self, input_shape: &[usize], capture_fn: F) -> Result<()>
|
|
where
|
|
F: FnOnce() -> Result<()>,
|
|
{
|
|
self.capture(input_shape, capture_fn)?;
|
|
// Note: capture executes the operations during recording,
|
|
// so we don't need an additional launch here.
|
|
// The graph will be used for subsequent calls.
|
|
Ok(())
|
|
}
|
|
|
|
/// Replay the captured graph.
|
|
///
|
|
/// This launches all captured operations with a single driver call,
|
|
/// eliminating kernel launch overhead.
|
|
///
|
|
/// # Returns
|
|
/// Error if the graph hasn't been captured yet.
|
|
pub fn launch(&mut self) -> Result<()> {
|
|
let graph = self.graph.as_ref()
|
|
.ok_or_else(|| TensorError::runtime("Graph not captured - call capture() first"))?;
|
|
|
|
graph.launch()
|
|
.map_err(|e| TensorError::device(format!("Graph launch failed: {:?}", e)))?;
|
|
|
|
self.launch_count += 1;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Invalidate the captured graph.
|
|
///
|
|
/// Call this when:
|
|
/// - Batch size changes
|
|
/// - Model topology changes
|
|
/// - Workspace buffers are reallocated
|
|
///
|
|
/// Note: You do NOT need to invalidate after optimizer steps if the
|
|
/// optimizer uses in-place updates (which preserve memory addresses).
|
|
pub fn invalidate(&mut self) {
|
|
#[cfg(debug_assertions)]
|
|
if self.graph.is_some() {
|
|
eprintln!("[PinnGraph] Invalidating CUDA graph (was captured for {:?})", self.input_shape);
|
|
}
|
|
self.graph = None;
|
|
self.input_shape.clear();
|
|
// Note: launch_count is preserved for statistics
|
|
}
|
|
|
|
/// Check if capture is currently in progress.
|
|
///
|
|
/// This is mainly for debugging and error messages.
|
|
pub fn is_capturing(&self) -> bool {
|
|
self.capturing
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "cuda")]
|
|
impl std::fmt::Debug for PinnGraph {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("PinnGraph")
|
|
.field("captured", &self.is_captured())
|
|
.field("input_shape", &self.input_shape)
|
|
.field("launch_count", &self.launch_count)
|
|
.field("capturing", &self.capturing)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "cuda")]
|
|
impl Drop for PinnGraph {
|
|
fn drop(&mut self) {
|
|
self.invalidate();
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// NON-CUDA STUB
|
|
// =============================================================================
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
pub struct PinnGraph {
|
|
_private: (),
|
|
}
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
impl PinnGraph {
|
|
pub fn new(_ctx: std::sync::Arc<super::cuda_stream_context::PinnStreamContext>) -> Self {
|
|
Self { _private: () }
|
|
}
|
|
|
|
pub fn is_captured(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
pub fn launch_count(&self) -> u64 {
|
|
0
|
|
}
|
|
|
|
pub fn capture<F>(&mut self, _input_shape: &[usize], _capture_fn: F) -> Result<()>
|
|
where
|
|
F: FnOnce() -> Result<()>,
|
|
{
|
|
Err(TensorError::device("CUDA feature not enabled"))
|
|
}
|
|
|
|
pub fn launch(&mut self) -> Result<()> {
|
|
Err(TensorError::device("CUDA feature not enabled"))
|
|
}
|
|
|
|
pub fn invalidate(&mut self) {}
|
|
}
|
|
|
|
// =============================================================================
|
|
// TESTS
|
|
// =============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_graph_creation() {
|
|
use super::super::cuda_stream_context::PinnStreamContext;
|
|
|
|
let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
|
|
let graph = PinnGraph::new(ctx);
|
|
|
|
assert!(!graph.is_captured());
|
|
assert_eq!(graph.launch_count(), 0);
|
|
assert!(graph.input_shape().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_graph_capture_simple() {
|
|
use super::super::cuda_stream_context::PinnStreamContext;
|
|
use rtx_tensor::{Tensor, Device};
|
|
|
|
let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
|
|
let mut graph = PinnGraph::new(ctx.clone());
|
|
let device = Device::cuda(0).unwrap();
|
|
|
|
// Create test tensors
|
|
let x = Tensor::randn(&[100, 1], &device).unwrap();
|
|
let b = Tensor::randn(&[1, 64], &device).unwrap();
|
|
let mut output = Tensor::zeros(&[100, 128], &device).unwrap();
|
|
|
|
// CRITICAL: Device-wide synchronize before capture to ensure all tensor
|
|
// allocations (which may use other streams) are complete. CUDA graph
|
|
// capture cannot have cross-stream dependencies.
|
|
ctx.device_synchronize().unwrap();
|
|
|
|
// Capture a simple operation
|
|
let result = graph.capture(&[100, 1], || {
|
|
ctx.fourier_features_out(&x, &b, 6.28, &mut output)
|
|
});
|
|
|
|
assert!(result.is_ok(), "Capture failed: {:?}", result.err());
|
|
assert!(graph.is_captured());
|
|
assert_eq!(graph.input_shape(), &[100, 1]);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_graph_replay() {
|
|
use super::super::cuda_stream_context::PinnStreamContext;
|
|
use rtx_tensor::{Tensor, Device};
|
|
|
|
let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
|
|
let mut graph = PinnGraph::new(ctx.clone());
|
|
let device = Device::cuda(0).unwrap();
|
|
|
|
let x = Tensor::randn(&[100, 1], &device).unwrap();
|
|
let b = Tensor::randn(&[1, 64], &device).unwrap();
|
|
let mut output = Tensor::zeros(&[100, 128], &device).unwrap();
|
|
|
|
// CRITICAL: Device-wide synchronize before capture to ensure all tensor
|
|
// allocations (which may use other streams) are complete. CUDA graph
|
|
// capture cannot have cross-stream dependencies.
|
|
ctx.device_synchronize().unwrap();
|
|
|
|
// Capture
|
|
graph.capture(&[100, 1], || {
|
|
ctx.fourier_features_out(&x, &b, 6.28, &mut output)
|
|
}).unwrap();
|
|
|
|
// Replay multiple times
|
|
for _ in 0..10 {
|
|
graph.launch().unwrap();
|
|
}
|
|
|
|
ctx.synchronize().unwrap();
|
|
assert_eq!(graph.launch_count(), 10);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_graph_invalidation() {
|
|
use super::super::cuda_stream_context::PinnStreamContext;
|
|
|
|
let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
|
|
let mut graph = PinnGraph::new(ctx.clone());
|
|
|
|
// Capture with a no-op (empty graph will fail, but that's OK for this test)
|
|
// We just want to test the invalidation logic
|
|
let _ = graph.capture(&[100, 1], || Ok(()));
|
|
|
|
// Invalidate
|
|
graph.invalidate();
|
|
|
|
assert!(!graph.is_captured());
|
|
assert!(graph.input_shape().is_empty());
|
|
}
|
|
}
|