93 lines
3.1 KiB
Rust
93 lines
3.1 KiB
Rust
//! CUDA Graphs: capture a sequence of kernel launches and replay them with
|
|
//! near-zero host-side overhead (~3-5 us per launch eliminated).
|
|
//!
|
|
//! Usage:
|
|
//! ```ignore
|
|
//! let stream = CudaStream::new()?;
|
|
//! let mut graph = CudaGraph::new();
|
|
//!
|
|
//! // First call: capture
|
|
//! graph.begin_capture(&stream)?;
|
|
//! // ... launch kernels on `stream` ...
|
|
//! graph.end_capture(&stream)?;
|
|
//!
|
|
//! // Subsequent calls: replay
|
|
//! graph.launch(&stream)?;
|
|
//! ```
|
|
//!
|
|
//! Requirements for captured kernels:
|
|
//! - All tensor shapes must be identical between capture and replay.
|
|
//! - No host-side branching during the captured section.
|
|
//! - Memory addresses used during capture must remain valid during replay.
|
|
|
|
use crate::error::{self, Result};
|
|
use crate::ffi;
|
|
use crate::stream::CudaStream;
|
|
|
|
/// RAII wrapper around a captured CUDA graph and its executable instance.
|
|
pub struct CudaGraph {
|
|
graph: ffi::CudaGraph,
|
|
exec: ffi::CudaGraphExec,
|
|
}
|
|
|
|
impl CudaGraph {
|
|
/// Create an empty graph handle (not yet captured).
|
|
pub fn new() -> Self {
|
|
Self {
|
|
graph: std::ptr::null_mut(),
|
|
exec: std::ptr::null_mut(),
|
|
}
|
|
}
|
|
|
|
/// Returns true if a graph has been captured and instantiated.
|
|
pub fn is_ready(&self) -> bool {
|
|
!self.exec.is_null()
|
|
}
|
|
|
|
/// Begin capturing kernel launches on `stream`.
|
|
/// All subsequent kernel launches on this stream are recorded into the
|
|
/// graph instead of being executed.
|
|
pub fn begin_capture(&mut self, stream: &CudaStream) -> Result<()> {
|
|
// If we have an old graph, destroy it first
|
|
self.destroy_inner();
|
|
// THREAD_LOCAL: only "potentially unsafe" CUDA calls (cudaMalloc etc.)
|
|
// made by THIS thread invalidate the capture. With GLOBAL mode, TP rank
|
|
// threads capturing concurrently would poison each other's captures.
|
|
error::check(unsafe {
|
|
ffi::cudaStreamBeginCapture(stream.as_raw(), ffi::CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL)
|
|
})
|
|
}
|
|
|
|
/// End capture and instantiate the executable graph.
|
|
pub fn end_capture(&mut self, stream: &CudaStream) -> Result<()> {
|
|
error::check(unsafe { ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph) })?;
|
|
error::check(unsafe { ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0) })
|
|
}
|
|
|
|
/// Replay the captured graph on `stream`.
|
|
/// Panics if no graph has been captured yet.
|
|
pub fn launch(&self, stream: &CudaStream) -> Result<()> {
|
|
assert!(self.is_ready(), "CudaGraph::launch called before capture");
|
|
error::check(unsafe { ffi::cudaGraphLaunch(self.exec, stream.as_raw()) })
|
|
}
|
|
|
|
fn destroy_inner(&mut self) {
|
|
if !self.exec.is_null() {
|
|
unsafe { ffi::cudaGraphExecDestroy(self.exec) };
|
|
self.exec = std::ptr::null_mut();
|
|
}
|
|
if !self.graph.is_null() {
|
|
unsafe { ffi::cudaGraphDestroy(self.graph) };
|
|
self.graph = std::ptr::null_mut();
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for CudaGraph {
|
|
fn drop(&mut self) {
|
|
self.destroy_inner();
|
|
}
|
|
}
|
|
|
|
unsafe impl Send for CudaGraph {}
|