metrale_storage/
cuda_graph.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: CUDA graph capture of a stream's work, replayed with one
4//! `cuGraphLaunch`.
5//!
6//! Owner: metrale-storage.
7//! Invariants:
8//! - Once `capture` has begun a stream capture, it ends it before returning,
9//!   also when `body` fails.
10
11use anyhow::{Context, Result, bail};
12
13unsafe extern "C" {
14    fn cuStreamBeginCapture_v2(stream: u64, mode: u32) -> i32;
15    fn cuStreamEndCapture(stream: u64, graph_out: *mut u64) -> i32;
16    fn cuGraphInstantiateWithFlags(exec_out: *mut u64, graph: u64, flags: u64) -> i32;
17    fn cuGraphLaunch(exec: u64, stream: u64) -> i32;
18    fn cuGraphExecDestroy(exec: u64) -> i32;
19    fn cuGraphDestroy(graph: u64) -> i32;
20}
21
22const CU_STREAM_CAPTURE_MODE_GLOBAL: u32 = 0;
23
24pub struct CapturedStep {
25    graph: u64,
26    graph_exec: u64,
27}
28
29impl CapturedStep {
30    /// 2026-09-25: Capture the stream work that `body` issues. Pointers passed
31    /// inside `body` are fixed at capture, so each `launch` reruns the work on
32    /// the same memory.
33    pub fn capture<F>(stream: u64, body: F) -> Result<Self>
34    where
35        F: FnOnce() -> Result<()>,
36    {
37        let s = unsafe { cuStreamBeginCapture_v2(stream, CU_STREAM_CAPTURE_MODE_GLOBAL) };
38        if s != 0 {
39            bail!("cuStreamBeginCapture_v2 failed: {s}");
40        }
41        let body_result = body();
42        let mut graph = 0u64;
43        let s2 = unsafe { cuStreamEndCapture(stream, &mut graph) };
44        // 2026-09-25: `body`'s error is returned only after the capture has
45        // ended, so the stream is not left capturing.
46        body_result.context("CUDA graph body")?;
47        if s2 != 0 {
48            bail!("cuStreamEndCapture failed: {s2}");
49        }
50        let mut exec = 0u64;
51        let s3 = unsafe { cuGraphInstantiateWithFlags(&mut exec, graph, 0) };
52        if s3 != 0 {
53            unsafe { cuGraphDestroy(graph) };
54            bail!("cuGraphInstantiateWithFlags failed: {s3}");
55        }
56        Ok(Self {
57            graph,
58            graph_exec: exec,
59        })
60    }
61
62    pub fn launch(&self, stream: u64) -> Result<()> {
63        let s = unsafe { cuGraphLaunch(self.graph_exec, stream) };
64        if s != 0 {
65            bail!("cuGraphLaunch failed: {s}");
66        }
67        Ok(())
68    }
69}
70
71impl Drop for CapturedStep {
72    fn drop(&mut self) {
73        unsafe {
74            let _ = cuGraphExecDestroy(self.graph_exec);
75            let _ = cuGraphDestroy(self.graph);
76        }
77    }
78}
metrale_storage/cuda_graph.rs

metrale_storage/
cuda_graph.rs
