spark_storage/
cuda_graph.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// CUDA graph capture + replay primitive (Phase 4). Captures a sequence of
4// stream operations into a single graph and replays them via one
5// `cuGraphLaunch` call (~2 µs vs ~5–10 µs per kernel in eager mode).
6
7use anyhow::{Context, Result, bail};
8
9unsafe extern "C" {
10    fn cuStreamBeginCapture_v2(stream: u64, mode: u32) -> i32;
11    fn cuStreamEndCapture(stream: u64, graph_out: *mut u64) -> i32;
12    fn cuGraphInstantiateWithFlags(exec_out: *mut u64, graph: u64, flags: u64) -> i32;
13    fn cuGraphLaunch(exec: u64, stream: u64) -> i32;
14    fn cuGraphExecDestroy(exec: u64) -> i32;
15    fn cuGraphDestroy(graph: u64) -> i32;
16}
17
18const CU_STREAM_CAPTURE_MODE_GLOBAL: u32 = 0;
19
20pub struct CapturedStep {
21    graph: u64,
22    graph_exec: u64,
23}
24
25impl CapturedStep {
26    /// Capture a sequence of stream operations issued by `body`. Device
27    /// pointers passed to launches inside `body` are baked in at capture
28    /// time; subsequent `launch` calls reread the same memory locations.
29    pub fn capture<F>(stream: u64, body: F) -> Result<Self>
30    where
31        F: FnOnce() -> Result<()>,
32    {
33        let s = unsafe { cuStreamBeginCapture_v2(stream, CU_STREAM_CAPTURE_MODE_GLOBAL) };
34        if s != 0 {
35            bail!("cuStreamBeginCapture_v2 failed: {s}");
36        }
37        let body_result = body();
38        let mut graph = 0u64;
39        let s2 = unsafe { cuStreamEndCapture(stream, &mut graph) };
40        // Surface body errors *after* end_capture so the stream isn't left
41        // mid-capture (which makes any further use of it fail).
42        body_result.context("CUDA graph body")?;
43        if s2 != 0 {
44            bail!("cuStreamEndCapture failed: {s2}");
45        }
46        let mut exec = 0u64;
47        let s3 = unsafe { cuGraphInstantiateWithFlags(&mut exec, graph, 0) };
48        if s3 != 0 {
49            unsafe { cuGraphDestroy(graph) };
50            bail!("cuGraphInstantiateWithFlags failed: {s3}");
51        }
52        Ok(Self {
53            graph,
54            graph_exec: exec,
55        })
56    }
57
58    pub fn launch(&self, stream: u64) -> Result<()> {
59        let s = unsafe { cuGraphLaunch(self.graph_exec, stream) };
60        if s != 0 {
61            bail!("cuGraphLaunch failed: {s}");
62        }
63        Ok(())
64    }
65}
66
67impl Drop for CapturedStep {
68    fn drop(&mut self) {
69        unsafe {
70            let _ = cuGraphExecDestroy(self.graph_exec);
71            let _ = cuGraphDestroy(self.graph);
72        }
73    }
74}
spark_storage/cuda_graph.rs

spark_storage/
cuda_graph.rs
