oxicuda_graph/executor/
cuda_graph.rs

1//! CUDA Graph capture backend.
2//!
3//! Converts an [`ExecutionPlan`] into an `oxicuda_driver::graph::Graph`,
4//! which can be instantiated and launched with minimal CPU overhead via the
5//! CUDA driver's graph API.
6//!
7//! # Mapping
8//!
9//! | `PlanStep`          | `driver::graph::GraphNode`         |
10//! |---------------------|------------------------------------|
11//! | `KernelLaunch`      | `GraphNode::KernelLaunch`          |
12//! | `Memcpy`            | `GraphNode::Memcpy`                |
13//! | `Memset`            | `GraphNode::Memset`                |
14//! | `Barrier`           | `GraphNode::Empty`                 |
15//! | `EventRecord/Wait`  | `GraphNode::Empty` (sync barrier)  |
16//! | `HostCallback`      | `GraphNode::Empty` (placeholder)   |
17//!
18//! The dependency edges from the original plan are preserved: each step
19//! depends on all steps that produced an event it is waiting for.
20//!
21//! # Limitations
22//!
23//! * Stream-parallel execution is not yet mapped to separate CUDA stream
24//!   arguments — all nodes are enqueued on the same graph-internal stream.
25//!   Proper multi-stream CUDA graph support requires the `cuGraphAddNode`
26//!   API with `cuGraphNodeParams`, which is future work.
27//! * `EventRecord`/`EventWait` steps collapse into `Empty` nodes with a
28//!   dependency edge, which enforces ordering without the overhead of true
29//!   event signalling within a CUDA graph.
30
31use oxicuda_driver::graph::{Graph, MemcpyDirection};
32
33use crate::error::GraphResult;
34use crate::executor::plan::{ExecutionPlan, PlanStep};
35use crate::node::MemcpyDir;
36
37/// Converts a `MemcpyDir` to the driver's `MemcpyDirection`.
38fn to_driver_dir(dir: &MemcpyDir) -> MemcpyDirection {
39    match dir {
40        MemcpyDir::HostToDevice => MemcpyDirection::HostToDevice,
41        MemcpyDir::DeviceToHost => MemcpyDirection::DeviceToHost,
42        MemcpyDir::DeviceToDevice | MemcpyDir::PeerToPeer { .. } => MemcpyDirection::DeviceToDevice,
43    }
44}
45
46/// Captures an [`ExecutionPlan`] into a `driver::graph::Graph`.
47///
48/// Returns the resulting graph, ready to be instantiated via
49/// [`Graph::instantiate`].
50///
51/// # Errors
52///
53/// Returns a [`GraphError`] if any step cannot be added to the graph.
54///
55/// [`GraphError`]: crate::error::GraphError
56pub fn capture(plan: &ExecutionPlan) -> GraphResult<Graph> {
57    let mut graph = Graph::new();
58
59    // Map from plan step index → driver graph node index.
60    let mut step_to_node: Vec<Option<usize>> = vec![None; plan.steps.len()];
61
62    // Map from event_id → the step index that recorded it.
63    let mut event_record_step: std::collections::HashMap<usize, usize> =
64        std::collections::HashMap::new();
65
66    // First pass: add all nodes to the driver graph.
67    for (i, step) in plan.steps.iter().enumerate() {
68        let node_idx = match step {
69            PlanStep::KernelLaunch {
70                function_name,
71                config,
72                ..
73            } => graph.add_kernel_node(
74                function_name,
75                config.grid,
76                config.block,
77                config.shared_mem_bytes,
78            ),
79            PlanStep::Memcpy {
80                dir, size_bytes, ..
81            } => graph.add_memcpy_node(to_driver_dir(dir), *size_bytes),
82            PlanStep::Memset {
83                size_bytes, value, ..
84            } => graph.add_memset_node(*size_bytes, *value),
85            PlanStep::EventRecord { event_id, .. } => {
86                // Record which step emitted this event.
87                event_record_step.insert(*event_id, i);
88                graph.add_empty_node()
89            }
90            PlanStep::EventWait { .. } => graph.add_empty_node(),
91            PlanStep::HostCallback { .. } | PlanStep::Barrier { .. } => graph.add_empty_node(),
92        };
93        step_to_node[i] = Some(node_idx);
94    }
95
96    // Second pass: add dependency edges.
97    // Strategy: the plan steps are in topological submission order,
98    // so for each step we add an edge from the immediately preceding
99    // step on the same stream (sequential ordering within a stream)
100    // and from the EventRecord step for each EventWait.
101
102    // Track the most recent step on each stream.
103    let mut last_on_stream: std::collections::HashMap<u32, usize> =
104        std::collections::HashMap::new();
105
106    for (i, step) in plan.steps.iter().enumerate() {
107        let sid = step.stream().0;
108        let my_node = step_to_node[i].ok_or_else(|| {
109            crate::error::GraphError::Internal(format!("step {i} node not populated"))
110        })?;
111
112        // Within-stream ordering: depend on the previous step on this stream.
113        if let Some(&prev_step) = last_on_stream.get(&sid) {
114            let prev_node = step_to_node[prev_step].ok_or_else(|| {
115                crate::error::GraphError::Internal(format!("step {prev_step} node not populated"))
116            })?;
117            graph.add_dependency(prev_node, my_node).ok();
118        }
119        last_on_stream.insert(sid, i);
120
121        // Cross-stream ordering: EventWait depends on the EventRecord.
122        if let PlanStep::EventWait { event_id, .. } = step {
123            if let Some(&record_step) = event_record_step.get(event_id) {
124                let record_node = step_to_node[record_step].ok_or_else(|| {
125                    crate::error::GraphError::Internal(format!(
126                        "event record step {record_step} node not populated"
127                    ))
128                })?;
129                graph.add_dependency(record_node, my_node).ok();
130            }
131        }
132    }
133
134    Ok(graph)
135}
136
137// ---------------------------------------------------------------------------
138// Tests
139// ---------------------------------------------------------------------------
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::builder::GraphBuilder;
145    use crate::executor::plan::ExecutionPlan;
146    use crate::graph::ComputeGraph;
147    use crate::node::MemcpyDir;
148
149    fn simple_graph() -> ComputeGraph {
150        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
151        let up = b.add_memcpy("up", MemcpyDir::HostToDevice, 1024);
152        let k = b.add_kernel("gemm", 4, 256, 0).fusible(false).finish();
153        let dn = b.add_memcpy("dn", MemcpyDir::DeviceToHost, 1024);
154        b.chain(&[up, k, dn]);
155        b.build().expect("test graph builds successfully")
156    }
157
158    #[test]
159    fn cuda_graph_capture_no_panic() {
160        let g = simple_graph();
161        let plan = ExecutionPlan::build(&g, 1).expect("execution plan builds from valid graph");
162        let dg = capture(&plan).expect("CUDA graph capture succeeds on valid execution plan");
163        assert!(dg.node_count() > 0);
164    }
165
166    #[test]
167    fn cuda_graph_node_count_matches_steps() {
168        let g = simple_graph();
169        let plan = ExecutionPlan::build(&g, 1).expect("execution plan builds from valid graph");
170        let step_count = plan.steps.len();
171        let dg = capture(&plan).expect("CUDA graph capture succeeds on valid execution plan");
172        assert_eq!(dg.node_count(), step_count);
173    }
174
175    #[test]
176    fn cuda_graph_has_dependencies() {
177        // Linear chain → must have edges.
178        let g = simple_graph();
179        let plan = ExecutionPlan::build(&g, 1).expect("execution plan builds from valid graph");
180        let dg = capture(&plan).expect("CUDA graph capture succeeds on valid execution plan");
181        assert!(dg.dependency_count() > 0);
182    }
183
184    #[test]
185    fn cuda_graph_is_dag_instantiatable() {
186        let g = simple_graph();
187        let plan = ExecutionPlan::build(&g, 1).expect("execution plan builds from valid graph");
188        let dg = capture(&plan).expect("CUDA graph capture succeeds on valid execution plan");
189        // The graph should be a valid DAG (instantiate runs topological sort).
190        let exec = dg.instantiate();
191        assert!(exec.is_ok(), "captured graph must be a valid DAG");
192    }
193
194    #[test]
195    fn cuda_graph_event_sync_adds_edge() {
196        // Manually build a plan with EventRecord → EventWait across streams.
197        let steps = vec![
198            PlanStep::KernelLaunch {
199                nodes: vec![crate::node::NodeId(0)],
200                function_name: "k0".into(),
201                config: crate::node::KernelConfig::linear(1, 32, 0),
202                stream: crate::node::StreamId(0),
203            },
204            PlanStep::EventRecord {
205                event_id: 0,
206                stream: crate::node::StreamId(0),
207            },
208            PlanStep::EventWait {
209                event_id: 0,
210                stream: crate::node::StreamId(1),
211            },
212            PlanStep::KernelLaunch {
213                nodes: vec![crate::node::NodeId(1)],
214                function_name: "k1".into(),
215                config: crate::node::KernelConfig::linear(1, 32, 0),
216                stream: crate::node::StreamId(1),
217            },
218        ];
219        let plan = ExecutionPlan {
220            steps,
221            num_streams: 2,
222            pool_bytes: 0,
223            kernel_count_original: 2,
224            kernel_count_fused: 2,
225            event_count: 1,
226        };
227        let dg = capture(&plan).expect("CUDA graph capture succeeds on valid execution plan");
228        // The EventWait node must depend on the EventRecord node.
229        assert!(dg.dependency_count() >= 1);
230        let exec = dg.instantiate();
231        assert!(exec.is_ok());
232    }
233
234    #[test]
235    fn cuda_graph_memset_node_recorded() {
236        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
237        b.add_memset("zero", 4096, 0x00);
238        let g = b.build().expect("test graph builds successfully");
239        let plan = ExecutionPlan::build(&g, 1).expect("execution plan builds from valid graph");
240        let dg = capture(&plan).expect("CUDA graph capture succeeds on valid execution plan");
241        assert_eq!(dg.node_count(), 1);
242        // Should be a Memset node in the driver graph.
243        let nodes = dg.nodes();
244        assert!(matches!(
245            nodes[0],
246            oxicuda_driver::graph::GraphNode::Memset { .. }
247        ));
248    }
249
250    #[test]
251    fn cuda_graph_barrier_becomes_empty_node() {
252        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
253        b.add_barrier("sync");
254        let g = b.build().expect("test graph builds successfully");
255        let plan = ExecutionPlan::build(&g, 1).expect("execution plan builds from valid graph");
256        let dg = capture(&plan).expect("CUDA graph capture succeeds on valid execution plan");
257        let nodes = dg.nodes();
258        // Barrier → Empty node.
259        assert!(matches!(nodes[0], oxicuda_driver::graph::GraphNode::Empty));
260    }
261}
oxicuda_graph/executor/cuda_graph.rs

oxicuda_graph/executor/
cuda_graph.rs
