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].unwrap();
109
110        // Within-stream ordering: depend on the previous step on this stream.
111        if let Some(&prev_step) = last_on_stream.get(&sid) {
112            let prev_node = step_to_node[prev_step].unwrap();
113            graph.add_dependency(prev_node, my_node).ok();
114        }
115        last_on_stream.insert(sid, i);
116
117        // Cross-stream ordering: EventWait depends on the EventRecord.
118        if let PlanStep::EventWait { event_id, .. } = step {
119            if let Some(&record_step) = event_record_step.get(event_id) {
120                let record_node = step_to_node[record_step].unwrap();
121                graph.add_dependency(record_node, my_node).ok();
122            }
123        }
124    }
125
126    Ok(graph)
127}
128
129// ---------------------------------------------------------------------------
130// Tests
131// ---------------------------------------------------------------------------
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use crate::builder::GraphBuilder;
137    use crate::executor::plan::ExecutionPlan;
138    use crate::graph::ComputeGraph;
139    use crate::node::MemcpyDir;
140
141    fn simple_graph() -> ComputeGraph {
142        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
143        let up = b.add_memcpy("up", MemcpyDir::HostToDevice, 1024);
144        let k = b.add_kernel("gemm", 4, 256, 0).fusible(false).finish();
145        let dn = b.add_memcpy("dn", MemcpyDir::DeviceToHost, 1024);
146        b.chain(&[up, k, dn]);
147        b.build().unwrap()
148    }
149
150    #[test]
151    fn cuda_graph_capture_no_panic() {
152        let g = simple_graph();
153        let plan = ExecutionPlan::build(&g, 1).unwrap();
154        let dg = capture(&plan).unwrap();
155        assert!(dg.node_count() > 0);
156    }
157
158    #[test]
159    fn cuda_graph_node_count_matches_steps() {
160        let g = simple_graph();
161        let plan = ExecutionPlan::build(&g, 1).unwrap();
162        let step_count = plan.steps.len();
163        let dg = capture(&plan).unwrap();
164        assert_eq!(dg.node_count(), step_count);
165    }
166
167    #[test]
168    fn cuda_graph_has_dependencies() {
169        // Linear chain → must have edges.
170        let g = simple_graph();
171        let plan = ExecutionPlan::build(&g, 1).unwrap();
172        let dg = capture(&plan).unwrap();
173        assert!(dg.dependency_count() > 0);
174    }
175
176    #[test]
177    fn cuda_graph_is_dag_instantiatable() {
178        let g = simple_graph();
179        let plan = ExecutionPlan::build(&g, 1).unwrap();
180        let dg = capture(&plan).unwrap();
181        // The graph should be a valid DAG (instantiate runs topological sort).
182        let exec = dg.instantiate();
183        assert!(exec.is_ok(), "captured graph must be a valid DAG");
184    }
185
186    #[test]
187    fn cuda_graph_event_sync_adds_edge() {
188        // Manually build a plan with EventRecord → EventWait across streams.
189        let steps = vec![
190            PlanStep::KernelLaunch {
191                nodes: vec![crate::node::NodeId(0)],
192                function_name: "k0".into(),
193                config: crate::node::KernelConfig::linear(1, 32, 0),
194                stream: crate::node::StreamId(0),
195            },
196            PlanStep::EventRecord {
197                event_id: 0,
198                stream: crate::node::StreamId(0),
199            },
200            PlanStep::EventWait {
201                event_id: 0,
202                stream: crate::node::StreamId(1),
203            },
204            PlanStep::KernelLaunch {
205                nodes: vec![crate::node::NodeId(1)],
206                function_name: "k1".into(),
207                config: crate::node::KernelConfig::linear(1, 32, 0),
208                stream: crate::node::StreamId(1),
209            },
210        ];
211        let plan = ExecutionPlan {
212            steps,
213            num_streams: 2,
214            pool_bytes: 0,
215            kernel_count_original: 2,
216            kernel_count_fused: 2,
217            event_count: 1,
218        };
219        let dg = capture(&plan).unwrap();
220        // The EventWait node must depend on the EventRecord node.
221        assert!(dg.dependency_count() >= 1);
222        let exec = dg.instantiate();
223        assert!(exec.is_ok());
224    }
225
226    #[test]
227    fn cuda_graph_memset_node_recorded() {
228        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
229        b.add_memset("zero", 4096, 0x00);
230        let g = b.build().unwrap();
231        let plan = ExecutionPlan::build(&g, 1).unwrap();
232        let dg = capture(&plan).unwrap();
233        assert_eq!(dg.node_count(), 1);
234        // Should be a Memset node in the driver graph.
235        let nodes = dg.nodes();
236        assert!(matches!(
237            nodes[0],
238            oxicuda_driver::graph::GraphNode::Memset { .. }
239        ));
240    }
241
242    #[test]
243    fn cuda_graph_barrier_becomes_empty_node() {
244        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
245        b.add_barrier("sync");
246        let g = b.build().unwrap();
247        let plan = ExecutionPlan::build(&g, 1).unwrap();
248        let dg = capture(&plan).unwrap();
249        let nodes = dg.nodes();
250        // Barrier → Empty node.
251        assert!(matches!(nodes[0], oxicuda_driver::graph::GraphNode::Empty));
252    }
253}
oxicuda_graph/executor/cuda_graph.rs

oxicuda_graph/executor/
cuda_graph.rs
