1use std::collections::BTreeSet;
29
30use crate::error::{GraphError, GraphResult};
31use crate::graph::ComputeGraph;
32use crate::node::{KernelConfig, NodeId, NodeKind};
33
34#[derive(Debug, Clone, PartialEq)]
45pub enum NodeParamUpdate {
46 Kernel {
50 function_name: Option<String>,
52 config: KernelConfig,
54 },
55 Memset {
58 value: u8,
60 },
61}
62
63impl NodeParamUpdate {
64 #[must_use]
66 pub fn target_tag(&self) -> &'static str {
67 match self {
68 Self::Kernel { .. } => "kernel",
69 Self::Memset { .. } => "memset",
70 }
71 }
72}
73
74#[derive(Debug, Clone)]
83pub struct ExecGraph {
84 kinds: Vec<NodeKind>,
86 edges: BTreeSet<(u32, u32)>,
88 execution_order: Vec<NodeId>,
90}
91
92impl ExecGraph {
93 pub fn instantiate(graph: &ComputeGraph) -> GraphResult<Self> {
103 let execution_order = graph.topological_order()?;
104 let kinds: Vec<NodeKind> = graph.nodes().iter().map(|n| n.kind.clone()).collect();
105 let edges: BTreeSet<(u32, u32)> =
106 graph.edges().into_iter().map(|(a, b)| (a.0, b.0)).collect();
107 Ok(Self {
108 kinds,
109 edges,
110 execution_order,
111 })
112 }
113
114 #[must_use]
116 pub fn node_count(&self) -> usize {
117 self.kinds.len()
118 }
119
120 #[must_use]
122 pub fn edge_count(&self) -> usize {
123 self.edges.len()
124 }
125
126 #[must_use]
128 pub fn execution_order(&self) -> &[NodeId] {
129 &self.execution_order
130 }
131
132 pub fn node_kind(&self, id: NodeId) -> GraphResult<&NodeKind> {
138 self.kinds
139 .get(id.0 as usize)
140 .ok_or(GraphError::NodeNotFound(id))
141 }
142
143 #[must_use]
145 pub fn has_edge(&self, from: NodeId, to: NodeId) -> bool {
146 self.edges.contains(&(from.0, to.0))
147 }
148
149 pub fn update_node(&mut self, id: NodeId, update: NodeParamUpdate) -> GraphResult<()> {
160 let kind = self
161 .kinds
162 .get_mut(id.0 as usize)
163 .ok_or(GraphError::NodeNotFound(id))?;
164 match (kind, update) {
165 (
166 NodeKind::KernelLaunch {
167 function_name,
168 config,
169 ..
170 },
171 NodeParamUpdate::Kernel {
172 function_name: new_name,
173 config: new_config,
174 },
175 ) => {
176 if let Some(name) = new_name {
177 *function_name = name;
178 }
179 *config = new_config;
180 Ok(())
181 }
182 (NodeKind::Memset { value, .. }, NodeParamUpdate::Memset { value: new_value }) => {
183 *value = new_value;
184 Ok(())
185 }
186 (kind, update) => Err(GraphError::ValidationFailed(format!(
187 "node {id} is a '{}' node but update targets '{}'",
188 kind.tag(),
189 update.target_tag()
190 ))),
191 }
192 }
193
194 pub fn update(&mut self, new_graph: &ComputeGraph) -> GraphResult<()> {
211 let diff = self.diff(new_graph)?;
212 if !diff.is_updatable() {
213 return Err(GraphError::ValidationFailed(diff.reject_reason()));
214 }
215 self.kinds = new_graph.nodes().iter().map(|n| n.kind.clone()).collect();
218 Ok(())
219 }
220
221 pub fn diff(&self, other: &ComputeGraph) -> GraphResult<ExecGraphDiff> {
231 if other.is_empty() {
232 return Err(GraphError::EmptyGraph);
233 }
234 if other.node_count() != self.kinds.len() {
236 return Ok(ExecGraphDiff {
237 node_count_changed: true,
238 topology_changed: true,
239 non_updatable_nodes: Vec::new(),
240 changed_params: Vec::new(),
241 });
242 }
243 let other_edges: BTreeSet<(u32, u32)> =
245 other.edges().into_iter().map(|(a, b)| (a.0, b.0)).collect();
246 let topology_changed = other_edges != self.edges;
247
248 let mut non_updatable_nodes = Vec::new();
249 let mut changed_params = Vec::new();
250 for (i, new_node) in other.nodes().iter().enumerate() {
251 let old = &self.kinds[i];
252 let new = &new_node.kind;
253 match classify_change(old, new) {
254 NodeChange::Same => {}
255 NodeChange::Param => changed_params.push(NodeId(i as u32)),
256 NodeChange::NonUpdatable => non_updatable_nodes.push(NodeId(i as u32)),
257 }
258 }
259
260 Ok(ExecGraphDiff {
261 node_count_changed: false,
262 topology_changed,
263 non_updatable_nodes,
264 changed_params,
265 })
266 }
267}
268
269enum NodeChange {
274 Same,
276 Param,
278 NonUpdatable,
281}
282
283fn classify_change(old: &NodeKind, new: &NodeKind) -> NodeChange {
285 match (old, new) {
286 (
287 NodeKind::KernelLaunch {
288 function_name: of,
289 config: oc,
290 fusible: ob,
291 },
292 NodeKind::KernelLaunch {
293 function_name: nf,
294 config: nc,
295 fusible: nb,
296 },
297 ) => {
298 if ob != nb {
299 NodeChange::NonUpdatable
302 } else if of == nf && oc == nc {
303 NodeChange::Same
304 } else {
305 NodeChange::Param
306 }
307 }
308 (
309 NodeKind::Memset {
310 size_bytes: os,
311 value: ov,
312 },
313 NodeKind::Memset {
314 size_bytes: ns,
315 value: nv,
316 },
317 ) => {
318 if os != ns {
319 NodeChange::NonUpdatable
320 } else if ov == nv {
321 NodeChange::Same
322 } else {
323 NodeChange::Param
324 }
325 }
326 (
329 NodeKind::Memcpy {
330 dir: od,
331 size_bytes: os,
332 },
333 NodeKind::Memcpy {
334 dir: nd,
335 size_bytes: ns,
336 },
337 ) => {
338 if od == nd && os == ns {
339 NodeChange::Same
340 } else {
341 NodeChange::NonUpdatable
342 }
343 }
344 (a, b) if a == b => NodeChange::Same,
346 _ => NodeChange::NonUpdatable,
348 }
349}
350
351#[derive(Debug, Clone, PartialEq, Eq)]
357pub struct ExecGraphDiff {
358 pub node_count_changed: bool,
360 pub topology_changed: bool,
362 pub non_updatable_nodes: Vec<NodeId>,
365 pub changed_params: Vec<NodeId>,
367}
368
369impl ExecGraphDiff {
370 #[must_use]
373 pub fn is_identical(&self) -> bool {
374 !self.node_count_changed
375 && !self.topology_changed
376 && self.non_updatable_nodes.is_empty()
377 && self.changed_params.is_empty()
378 }
379
380 #[must_use]
383 pub fn is_updatable(&self) -> bool {
384 !self.node_count_changed && !self.topology_changed && self.non_updatable_nodes.is_empty()
385 }
386
387 #[must_use]
391 pub fn reject_reason(&self) -> String {
392 if self.node_count_changed {
393 "node count differs — topology must match for in-place update".to_owned()
394 } else if self.topology_changed {
395 "dependency edges differ — topology must match for in-place update".to_owned()
396 } else if !self.non_updatable_nodes.is_empty() {
397 format!(
398 "{} node(s) changed in a non-updatable way (kind / size / fusibility)",
399 self.non_updatable_nodes.len()
400 )
401 } else {
402 "update is possible".to_owned()
403 }
404 }
405}
406
407#[cfg(test)]
412mod tests {
413 use super::*;
414 use crate::builder::GraphBuilder;
415 use crate::node::MemcpyDir;
416
417 fn diamond() -> (ComputeGraph, [NodeId; 4]) {
418 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
419 let a = b.add_kernel("a", 4, 256, 0).finish();
420 let l = b.add_kernel("l", 4, 256, 0).finish();
421 let r = b.add_kernel("r", 4, 256, 0).finish();
422 let d = b.add_kernel("d", 4, 256, 0).finish();
423 b.dep(a, l).dep(a, r).dep(l, d).dep(r, d);
424 (b.build().expect("diamond builds"), [a, l, r, d])
425 }
426
427 #[test]
428 fn instantiate_empty_errors() {
429 let g = ComputeGraph::new();
430 assert!(matches!(
431 ExecGraph::instantiate(&g),
432 Err(GraphError::EmptyGraph)
433 ));
434 }
435
436 #[test]
437 fn instantiate_snapshots_topology() {
438 let (g, _) = diamond();
439 let ex = ExecGraph::instantiate(&g).expect("instantiate");
440 assert_eq!(ex.node_count(), 4);
441 assert_eq!(ex.edge_count(), 4);
442 assert_eq!(ex.execution_order().len(), 4);
443 }
444
445 #[test]
446 fn execution_order_respects_dependencies() {
447 let (g, [a, l, r, d]) = diamond();
448 let ex = ExecGraph::instantiate(&g).expect("instantiate");
449 let order = ex.execution_order();
450 let pos = |n: NodeId| order.iter().position(|&x| x == n).expect("present");
451 assert!(pos(a) < pos(l));
452 assert!(pos(a) < pos(r));
453 assert!(pos(l) < pos(d));
454 assert!(pos(r) < pos(d));
455 }
456
457 #[test]
458 fn clone_preserves_topology() {
459 let (g, [a, _l, _r, d]) = diamond();
460 let ex = ExecGraph::instantiate(&g).expect("instantiate");
461 let cloned = ex.clone();
462 assert_eq!(cloned.node_count(), ex.node_count());
463 assert_eq!(cloned.edge_count(), ex.edge_count());
464 assert!(cloned.has_edge(a, NodeId(1)));
465 assert_eq!(cloned.execution_order(), ex.execution_order());
466 assert!(!cloned.has_edge(d, a));
468 }
469
470 #[test]
471 fn update_node_kernel_config() {
472 let (g, [a, ..]) = diamond();
473 let mut ex = ExecGraph::instantiate(&g).expect("instantiate");
474 ex.update_node(
475 a,
476 NodeParamUpdate::Kernel {
477 function_name: Some("a_v2".into()),
478 config: KernelConfig::linear(8, 128, 512),
479 },
480 )
481 .expect("kernel update");
482 let k = ex.node_kind(a).expect("node a");
483 match k {
484 NodeKind::KernelLaunch {
485 function_name,
486 config,
487 ..
488 } => {
489 assert_eq!(function_name, "a_v2");
490 assert_eq!(config.grid, (8, 1, 1));
491 assert_eq!(config.shared_mem_bytes, 512);
492 }
493 _ => panic!("expected kernel"),
494 }
495 }
496
497 #[test]
498 fn update_node_wrong_kind_rejected() {
499 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
501 let z = b.add_memset("zero", 4096, 0);
502 let g = b.build().expect("builds");
503 let mut ex = ExecGraph::instantiate(&g).expect("instantiate");
504 let res = ex.update_node(
505 z,
506 NodeParamUpdate::Kernel {
507 function_name: None,
508 config: KernelConfig::linear(1, 1, 0),
509 },
510 );
511 assert!(matches!(res, Err(GraphError::ValidationFailed(_))));
512 }
513
514 #[test]
515 fn update_node_out_of_range() {
516 let (g, _) = diamond();
517 let mut ex = ExecGraph::instantiate(&g).expect("instantiate");
518 let res = ex.update_node(NodeId(99), NodeParamUpdate::Memset { value: 1 });
519 assert!(matches!(res, Err(GraphError::NodeNotFound(_))));
520 }
521
522 #[test]
523 fn whole_graph_update_param_only_succeeds() {
524 let (g0, _) = diamond();
525 let mut ex = ExecGraph::instantiate(&g0).expect("instantiate");
526 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
528 let a = b.add_kernel("a", 16, 64, 0).finish(); let l = b.add_kernel("l", 4, 256, 0).finish();
530 let r = b.add_kernel("r", 4, 256, 0).finish();
531 let d = b.add_kernel("d", 4, 256, 0).finish();
532 b.dep(a, l).dep(a, r).dep(l, d).dep(r, d);
533 let g1 = b.build().expect("builds");
534 ex.update(&g1).expect("param-only update should succeed");
535 match ex.node_kind(a).expect("a") {
537 NodeKind::KernelLaunch { config, .. } => assert_eq!(config.grid, (16, 1, 1)),
538 _ => panic!("kernel"),
539 }
540 }
541
542 #[test]
543 fn whole_graph_update_topology_change_rejected() {
544 let (g0, _) = diamond();
545 let mut ex = ExecGraph::instantiate(&g0).expect("instantiate");
546 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
548 let a = b.add_kernel("a", 4, 256, 0).finish();
549 let l = b.add_kernel("l", 4, 256, 0).finish();
550 let r = b.add_kernel("r", 4, 256, 0).finish();
551 let d = b.add_kernel("d", 4, 256, 0).finish();
552 b.dep(a, l).dep(a, r).dep(l, d); let g1 = b.build().expect("builds");
554 let res = ex.update(&g1);
555 assert!(matches!(res, Err(GraphError::ValidationFailed(_))));
556 }
557
558 #[test]
559 fn whole_graph_update_node_count_change_rejected() {
560 let (g0, _) = diamond();
561 let mut ex = ExecGraph::instantiate(&g0).expect("instantiate");
562 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
563 let a = b.add_kernel("a", 4, 256, 0).finish();
564 let l = b.add_kernel("l", 4, 256, 0).finish();
565 b.dep(a, l);
566 let g1 = b.build().expect("builds");
567 let res = ex.update(&g1);
568 assert!(matches!(res, Err(GraphError::ValidationFailed(_))));
569 }
570
571 #[test]
572 fn diff_identical_graph() {
573 let (g, _) = diamond();
574 let ex = ExecGraph::instantiate(&g).expect("instantiate");
575 let d = ex.diff(&g).expect("diff");
576 assert!(d.is_identical());
577 assert!(d.is_updatable());
578 assert!(d.changed_params.is_empty());
579 }
580
581 #[test]
582 fn diff_reports_changed_params() {
583 let (g0, [a, ..]) = diamond();
584 let ex = ExecGraph::instantiate(&g0).expect("instantiate");
585 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
586 let na = b.add_kernel("a", 99, 1, 0).finish(); let l = b.add_kernel("l", 4, 256, 0).finish();
588 let r = b.add_kernel("r", 4, 256, 0).finish();
589 let dd = b.add_kernel("d", 4, 256, 0).finish();
590 b.dep(na, l).dep(na, r).dep(l, dd).dep(r, dd);
591 let g1 = b.build().expect("builds");
592 let diff = ex.diff(&g1).expect("diff");
593 assert!(diff.is_updatable());
594 assert!(!diff.is_identical());
595 assert_eq!(diff.changed_params, vec![a]);
596 }
597
598 #[test]
599 fn diff_memcpy_size_change_non_updatable() {
600 let mut b0 = GraphBuilder::new().with_auto_infer_edges(false);
601 let up0 = b0.add_memcpy("up", MemcpyDir::HostToDevice, 1024);
602 let k0 = b0.add_kernel("k", 1, 32, 0).finish();
603 b0.dep(up0, k0);
604 let g0 = b0.build().expect("builds");
605 let ex = ExecGraph::instantiate(&g0).expect("instantiate");
606
607 let mut b1 = GraphBuilder::new().with_auto_infer_edges(false);
608 let up1 = b1.add_memcpy("up", MemcpyDir::HostToDevice, 2048); let k1 = b1.add_kernel("k", 1, 32, 0).finish();
610 b1.dep(up1, k1);
611 let g1 = b1.build().expect("builds");
612
613 let diff = ex.diff(&g1).expect("diff");
614 assert!(!diff.is_updatable());
615 assert_eq!(diff.non_updatable_nodes, vec![up0]);
616 assert!(matches!(ex.diff(&g1).map(|d| d.is_updatable()), Ok(false)));
617 }
618
619 #[test]
620 fn diff_kind_change_non_updatable() {
621 let mut b0 = GraphBuilder::new().with_auto_infer_edges(false);
622 let n0 = b0.add_kernel("k", 1, 32, 0).finish();
623 let g0 = b0.build().expect("builds");
624 let ex = ExecGraph::instantiate(&g0).expect("instantiate");
625
626 let mut b1 = GraphBuilder::new().with_auto_infer_edges(false);
627 let _n1 = b1.add_memset("z", 4096, 0); let g1 = b1.build().expect("builds");
629
630 let diff = ex.diff(&g1).expect("diff");
631 assert!(!diff.is_updatable());
632 assert_eq!(diff.non_updatable_nodes, vec![n0]);
633 }
634
635 #[test]
636 fn diff_empty_other_errors() {
637 let (g, _) = diamond();
638 let ex = ExecGraph::instantiate(&g).expect("instantiate");
639 let empty = ComputeGraph::new();
640 assert!(matches!(ex.diff(&empty), Err(GraphError::EmptyGraph)));
641 }
642
643 #[test]
644 fn param_update_target_tag() {
645 assert_eq!(NodeParamUpdate::Memset { value: 1 }.target_tag(), "memset");
646 assert_eq!(
647 NodeParamUpdate::Kernel {
648 function_name: None,
649 config: KernelConfig::linear(1, 1, 0)
650 }
651 .target_tag(),
652 "kernel"
653 );
654 }
655
656 #[test]
657 fn stress_large_graph_instantiates_and_diffs() {
658 const N: usize = 10_000;
661 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
662 let ids: Vec<NodeId> = (0..N)
665 .map(|i| b.add_kernel(&format!("k{i}"), 1, 32, 0).finish())
666 .collect();
667 for i in 2..N {
668 b.dep(ids[i - 1], ids[i]);
669 b.dep(ids[i - 2], ids[i]);
670 }
671 let g = b.build().expect("large graph builds");
672 let ex = ExecGraph::instantiate(&g).expect("large graph instantiates");
673 assert_eq!(ex.node_count(), N);
674 assert_eq!(ex.execution_order().len(), N);
675 let diff = ex.diff(&g).expect("self-diff");
677 assert!(diff.is_identical());
678 }
679
680 #[test]
681 fn update_memset_value() {
682 let mut b = GraphBuilder::new().with_auto_infer_edges(false);
683 let z = b.add_memset("z", 4096, 0x00);
684 let g = b.build().expect("builds");
685 let mut ex = ExecGraph::instantiate(&g).expect("instantiate");
686 ex.update_node(z, NodeParamUpdate::Memset { value: 0xff })
687 .expect("memset update");
688 match ex.node_kind(z).expect("z") {
689 NodeKind::Memset { value, .. } => assert_eq!(*value, 0xff),
690 _ => panic!("memset"),
691 }
692 }
693}