1use std::collections::{HashMap, HashSet};
63use std::sync::Arc;
64
65use smallvec::SmallVec;
66use vyre_foundation::ir::{BufferAccess, BufferDecl, Expr, Ident, MemoryKind, Node, Program};
67use vyre_foundation::memory_model::MemoryOrdering;
68
69use crate::backend::{
70 BackendError, DispatchConfig, OutputBuffers, ResidentDispatchStep, ResidentReadRange, Resource,
71 TimedDispatchResult, VyreBackend,
72};
73use crate::binding::{Binding, BindingPlan, BindingRole};
74
75#[derive(Clone, Debug, PartialEq, Eq)]
83enum EntryWrapper {
84 Region { generator: Ident },
85 Block,
86}
87
88struct PlannedGridSyncSegment {
89 program: Program,
90 input_names: Vec<Ident>,
91 output_names: Vec<Ident>,
92}
93
94fn peel_entry_wrappers(program: &Program) -> (Vec<EntryWrapper>, &[Node]) {
95 let mut wrappers = Vec::new();
96 let mut entry = program.entry();
97 loop {
98 if entry.len() == 1 {
99 match &entry[0] {
100 Node::Region {
101 generator, body, ..
102 } => {
103 wrappers.push(EntryWrapper::Region {
104 generator: generator.clone(),
105 });
106 entry = body.as_slice();
107 continue;
108 }
109 Node::Block(body) => {
110 wrappers.push(EntryWrapper::Block);
111 entry = body.as_slice();
112 continue;
113 }
114 _ => {}
115 }
116 }
117 break;
118 }
119 (wrappers, entry)
120}
121
122fn entry_sequence(program: &Program) -> &[Node] {
123 peel_entry_wrappers(program).1
124}
125
126#[must_use]
135pub fn contains_grid_sync(program: &Program) -> bool {
136 if !program.stats().has_node_barrier() {
142 return false;
143 }
144 node_slice_contains_grid_sync(entry_sequence(program))
145}
146
147fn node_slice_contains_grid_sync(nodes: &[Node]) -> bool {
148 nodes.iter().any(node_contains_grid_sync)
149}
150
151fn node_contains_grid_sync(node: &Node) -> bool {
152 match node {
153 Node::Barrier {
154 ordering: MemoryOrdering::GridSync,
155 ..
156 } => true,
157 Node::If {
158 then, otherwise, ..
159 } => node_slice_contains_grid_sync(then) || node_slice_contains_grid_sync(otherwise),
160 Node::Loop { body, .. } | Node::Block(body) => node_slice_contains_grid_sync(body),
161 Node::Region { body, .. } => node_slice_contains_grid_sync(body),
162 _ => false,
163 }
164}
165
166#[must_use]
179pub fn split_on_grid_sync(program: &Program) -> Vec<Program> {
180 match try_split_on_grid_sync(program) {
181 Ok(segments) => segments,
182 Err(_error) => Vec::new(),
183 }
184}
185
186fn hoist_grid_sync_barriers(nodes: &[Node]) -> Vec<Node> {
192 let mut new_nodes = Vec::new();
193 for node in nodes {
194 match node {
195 Node::Block(body) => {
196 let new_body = hoist_grid_sync_barriers(body);
197 let has_barrier = new_body.iter().any(|n| {
198 matches!(
199 n,
200 Node::Barrier {
201 ordering: MemoryOrdering::GridSync,
202 ..
203 }
204 )
205 });
206 if has_barrier {
207 let mut current_segment = Vec::new();
208 for b_node in new_body {
209 if matches!(
210 b_node,
211 Node::Barrier {
212 ordering: MemoryOrdering::GridSync,
213 ..
214 }
215 ) {
216 new_nodes.push(Node::Block(std::mem::take(&mut current_segment)));
217 new_nodes.push(b_node);
218 } else {
219 current_segment.push(b_node);
220 }
221 }
222 new_nodes.push(Node::Block(current_segment));
223 } else {
224 new_nodes.push(Node::Block(new_body));
225 }
226 }
227 Node::Region {
228 generator,
229 source_region,
230 body,
231 } => {
232 let new_body = hoist_grid_sync_barriers(body);
233 let has_barrier = new_body.iter().any(|n| {
234 matches!(
235 n,
236 Node::Barrier {
237 ordering: MemoryOrdering::GridSync,
238 ..
239 }
240 )
241 });
242 if has_barrier {
243 let mut current_segment = Vec::new();
244 for b_node in new_body {
245 if matches!(
246 b_node,
247 Node::Barrier {
248 ordering: MemoryOrdering::GridSync,
249 ..
250 }
251 ) {
252 new_nodes.push(Node::Region {
253 generator: generator.clone(),
254 source_region: source_region.clone(),
255 body: Arc::new(std::mem::take(&mut current_segment)),
256 });
257 new_nodes.push(b_node);
258 } else {
259 current_segment.push(b_node);
260 }
261 }
262 new_nodes.push(Node::Region {
263 generator: generator.clone(),
264 source_region: source_region.clone(),
265 body: Arc::new(current_segment),
266 });
267 } else {
268 new_nodes.push(Node::Region {
269 generator: generator.clone(),
270 source_region: source_region.clone(),
271 body: Arc::new(new_body),
272 });
273 }
274 }
275 other => {
276 new_nodes.push(other.clone());
277 }
278 }
279 }
280 new_nodes
281}
282
283fn collect_global_let_bindings(nodes: &[Node], map: &mut std::collections::HashMap<String, Node>) {
284 for node in nodes {
285 match node {
286 Node::Let { name, .. } => {
287 map.insert(name.as_str().to_string(), node.clone());
288 }
289 Node::If {
290 then, otherwise, ..
291 } => {
292 collect_global_let_bindings(then, map);
293 collect_global_let_bindings(otherwise, map);
294 }
295 Node::Loop { body, .. } | Node::Block(body) => {
296 collect_global_let_bindings(body, map);
297 }
298 Node::Region { body, .. } => {
299 collect_global_let_bindings(&body[..], map);
300 }
301 _ => {}
302 }
303 }
304}
305
306fn collect_locally_defined_vars(nodes: &[Node], vars: &mut std::collections::HashSet<String>) {
307 for node in nodes {
308 match node {
309 Node::Let { name, .. } => {
310 vars.insert(name.as_str().to_string());
311 }
312 Node::Loop { var, body, .. } => {
313 vars.insert(var.as_str().to_string());
314 collect_locally_defined_vars(body, vars);
315 }
316 Node::If {
317 then, otherwise, ..
318 } => {
319 collect_locally_defined_vars(then, vars);
320 collect_locally_defined_vars(otherwise, vars);
321 }
322 Node::Block(body) => {
323 collect_locally_defined_vars(body, vars);
324 }
325 Node::Region { body, .. } => {
326 collect_locally_defined_vars(&body[..], vars);
327 }
328 _ => {}
329 }
330 }
331}
332
333fn collect_referenced_vars(expr: &Expr, vars: &mut std::collections::HashSet<String>) {
334 match expr {
335 Expr::Var(name) => {
336 vars.insert(name.as_str().to_string());
337 }
338 Expr::Load { index, .. } => {
339 collect_referenced_vars(index, vars);
340 }
341 Expr::BinOp { left, right, .. } => {
342 collect_referenced_vars(left, vars);
343 collect_referenced_vars(right, vars);
344 }
345 Expr::UnOp { operand, .. } => {
346 collect_referenced_vars(operand, vars);
347 }
348 Expr::Call { args, .. } => {
349 for arg in args {
350 collect_referenced_vars(arg, vars);
351 }
352 }
353 Expr::Select {
354 cond,
355 true_val,
356 false_val,
357 } => {
358 collect_referenced_vars(cond, vars);
359 collect_referenced_vars(true_val, vars);
360 collect_referenced_vars(false_val, vars);
361 }
362 Expr::Cast { value, .. } => {
363 collect_referenced_vars(value, vars);
364 }
365 Expr::Fma { a, b, c } => {
366 collect_referenced_vars(a, vars);
367 collect_referenced_vars(b, vars);
368 collect_referenced_vars(c, vars);
369 }
370 Expr::Atomic {
371 index,
372 expected,
373 value,
374 ..
375 } => {
376 collect_referenced_vars(index, vars);
377 if let Some(expected) = expected {
378 collect_referenced_vars(expected, vars);
379 }
380 collect_referenced_vars(value, vars);
381 }
382 Expr::SubgroupBallot { cond } => {
383 collect_referenced_vars(cond, vars);
384 }
385 Expr::SubgroupShuffle { value, lane } => {
386 collect_referenced_vars(value, vars);
387 collect_referenced_vars(lane, vars);
388 }
389 Expr::SubgroupReduce { value, .. } => {
390 collect_referenced_vars(value, vars);
391 }
392 _ => {}
393 }
394}
395
396fn collect_node_referenced_vars(node: &Node, vars: &mut std::collections::HashSet<String>) {
397 match node {
398 Node::Let { value, .. } => {
399 collect_referenced_vars(value, vars);
400 }
401 Node::Assign { value, .. } => {
402 collect_referenced_vars(value, vars);
403 }
404 Node::Store { index, value, .. } => {
405 collect_referenced_vars(index, vars);
406 collect_referenced_vars(value, vars);
407 }
408 Node::If {
409 cond,
410 then,
411 otherwise,
412 } => {
413 collect_referenced_vars(cond, vars);
414 for n in then {
415 collect_node_referenced_vars(n, vars);
416 }
417 for n in otherwise {
418 collect_node_referenced_vars(n, vars);
419 }
420 }
421 Node::Loop { from, to, body, .. } => {
422 collect_referenced_vars(from, vars);
423 collect_referenced_vars(to, vars);
424 for n in body {
425 collect_node_referenced_vars(n, vars);
426 }
427 }
428 Node::Block(body) => {
429 for n in body {
430 collect_node_referenced_vars(n, vars);
431 }
432 }
433 Node::Region { body, .. } => {
434 for n in body.as_ref() {
435 collect_node_referenced_vars(n, vars);
436 }
437 }
438 Node::AsyncLoad { offset, size, .. } => {
439 collect_referenced_vars(offset, vars);
440 collect_referenced_vars(size, vars);
441 }
442 Node::AsyncStore { offset, size, .. } => {
443 collect_referenced_vars(offset, vars);
444 collect_referenced_vars(size, vars);
445 }
446 Node::Trap { address, .. } => {
447 collect_referenced_vars(address, vars);
448 }
449 _ => {}
450 }
451}
452
453fn resolve_dependencies(
454 name: &str,
455 global_lets: &std::collections::HashMap<String, Node>,
456 resolved_names: &mut std::collections::HashSet<String>,
457 resolved_lets: &mut Vec<Node>,
458) {
459 if resolved_names.contains(name) {
460 return;
461 }
462 if let Some(let_node) = global_lets.get(name) {
463 resolved_names.insert(name.to_string());
464 let mut deps = std::collections::HashSet::new();
465 collect_node_referenced_vars(let_node, &mut deps);
466 for dep in deps {
467 resolve_dependencies(&dep, global_lets, resolved_names, resolved_lets);
468 }
469 resolved_lets.push(let_node.clone());
470 }
471}
472
473fn propagate_let_bindings(segments: &mut [Vec<Node>], hoisted_inner: &[Node]) {
474 let mut global_lets = std::collections::HashMap::new();
475 collect_global_let_bindings(hoisted_inner, &mut global_lets);
476
477 for segment_nodes in segments {
478 let mut locally_defined = std::collections::HashSet::new();
479 collect_locally_defined_vars(segment_nodes, &mut locally_defined);
480
481 let mut referenced = std::collections::HashSet::new();
482 for node in segment_nodes.iter() {
483 collect_node_referenced_vars(node, &mut referenced);
484 }
485
486 let mut free_vars = Vec::new();
487 for name in referenced {
488 if !locally_defined.contains(&name) {
489 free_vars.push(name);
490 }
491 }
492
493 let mut resolved_lets = Vec::new();
494 let mut resolved_names = std::collections::HashSet::new();
495 for name in free_vars {
496 resolve_dependencies(&name, &global_lets, &mut resolved_names, &mut resolved_lets);
497 }
498
499 if !resolved_lets.is_empty() {
500 resolved_lets.extend(std::mem::take(segment_nodes));
501 *segment_nodes = resolved_lets;
502 }
503 }
504}
505
506pub fn try_split_on_grid_sync(program: &Program) -> Result<Vec<Program>, BackendError> {
513 let (wrappers, inner) = peel_entry_wrappers(program);
514 let hoisted_inner = hoist_grid_sync_barriers(inner);
515 let split_count = hoisted_inner
516 .iter()
517 .filter(|node| {
518 matches!(
519 node,
520 Node::Barrier {
521 ordering: MemoryOrdering::GridSync,
522 ..
523 }
524 )
525 })
526 .count();
527 if split_count == 0 {
528 let mut segments = Vec::new();
529 reserve_grid_sync_vec(&mut segments, 1, "grid-sync no-op segment")?;
530 segments.push(program.clone());
531 return Ok(segments);
532 }
533
534 let segment_count = split_count + 1;
535 let executable_nodes = hoisted_inner.len().checked_sub(split_count).ok_or_else(|| {
536 BackendError::InvalidProgram {
537 fix: format!(
538 "grid-sync split_count {split_count} exceeded entry node count {}. Fix: split_on_grid_sync must count barriers from the same entry sequence it segments.",
539 hoisted_inner.len()
540 ),
541 }
542 })?;
543 let segment_capacity = executable_nodes.div_ceil(segment_count);
544
545 let mut raw_segments = Vec::new();
546 let mut current = Vec::new();
547 reserve_grid_sync_vec(&mut current, segment_capacity, "grid-sync current segment")?;
548 for node in &hoisted_inner {
549 match node {
550 Node::Barrier {
551 ordering: MemoryOrdering::GridSync,
552 ..
553 } => {
554 let mut next = Vec::new();
555 reserve_grid_sync_vec(&mut next, segment_capacity, "grid-sync next segment")?;
556 let entry = std::mem::replace(&mut current, next);
557 raw_segments.push(entry);
558 }
559 other => {
560 current.push(other.clone());
561 }
562 }
563 }
564 raw_segments.push(current);
565
566 propagate_let_bindings(&mut raw_segments, &hoisted_inner);
567
568 let mut segments = Vec::new();
569 reserve_grid_sync_vec(
570 &mut segments,
571 raw_segments.len(),
572 "grid-sync split segments",
573 )?;
574 for entry in raw_segments {
575 segments.push(wrap_split_segment(program, &wrappers, entry));
576 }
577 Ok(segments)
578}
579
580fn wrap_split_segment(program: &Program, wrappers: &[EntryWrapper], entry: Vec<Node>) -> Program {
581 let mut wrapped_entry = entry;
585 for wrapper in wrappers.iter().rev() {
586 match wrapper {
587 EntryWrapper::Region { generator } => {
588 wrapped_entry = vec![Node::Region {
589 generator: generator.clone(),
590 source_region: None,
591 body: Arc::new(wrapped_entry),
592 }];
593 }
594 EntryWrapper::Block => {
595 wrapped_entry = vec![Node::Block(wrapped_entry)];
596 }
597 }
598 }
599 program.with_rewritten_entry(wrapped_entry)
600}
601
602pub fn plan_host_grid_sync_segment_programs(
612 program: &Program,
613) -> Result<Vec<Program>, BackendError> {
614 Ok(plan_host_grid_sync_segments(program)?
615 .into_iter()
616 .map(|segment| segment.program)
617 .collect())
618}
619
620fn plan_host_grid_sync_segments(
621 program: &Program,
622) -> Result<Vec<PlannedGridSyncSegment>, BackendError> {
623 let split = try_split_on_grid_sync(program)?;
624 let first_writer = first_writer_segment_per_buffer(&split, program)?;
625 let mut planned = Vec::new();
626 reserve_grid_sync_vec(&mut planned, split.len(), "grid-sync planned host segments")?;
627 for (segment_idx, segment) in split.into_iter().enumerate() {
628 let rewritten =
629 rewrite_segment_buffers_for_host_split(program, &segment, segment_idx, &first_writer)?;
630 let input_names = segment_input_names(&rewritten)?;
631 let output_names = segment_output_names(&rewritten)?;
632 planned.push(PlannedGridSyncSegment {
633 program: rewritten,
634 input_names,
635 output_names,
636 });
637 }
638 Ok(planned)
639}
640
641fn first_writer_segment_per_buffer(
654 split: &[Program],
655 program: &Program,
656) -> Result<HashMap<Ident, usize>, BackendError> {
657 let mut first_writer: HashMap<Ident, usize> = HashMap::new();
658 reserve_grid_sync_hash_map(
659 &mut first_writer,
660 program.buffers().len(),
661 "grid-sync first-writer map",
662 )?;
663 for (segment_idx, segment) in split.iter().enumerate() {
664 let mut reads = HashSet::new();
665 let mut writes = HashSet::new();
666 reserve_grid_sync_hash_set(
667 &mut reads,
668 program.buffers().len(),
669 "grid-sync first-writer read scan",
670 )?;
671 reserve_grid_sync_hash_set(
672 &mut writes,
673 program.buffers().len(),
674 "grid-sync first-writer write scan",
675 )?;
676 for node in entry_sequence(segment) {
677 collect_segment_buffer_targets(node, &mut reads, &mut writes);
678 }
679 for name in writes {
680 first_writer.entry(name).or_insert(segment_idx);
681 }
682 }
683 Ok(first_writer)
684}
685
686fn rewrite_segment_buffers_for_host_split(
687 source: &Program,
688 segment: &Program,
689 segment_idx: usize,
690 first_writer: &HashMap<Ident, usize>,
691) -> Result<Program, BackendError> {
692 let mut reads = HashSet::new();
693 let mut writes = HashSet::new();
694 reserve_grid_sync_hash_set(
695 &mut reads,
696 source.buffers().len(),
697 "grid-sync segment read set",
698 )?;
699 reserve_grid_sync_hash_set(
700 &mut writes,
701 source.buffers().len(),
702 "grid-sync segment write set",
703 )?;
704 for node in entry_sequence(segment) {
705 collect_segment_buffer_targets(node, &mut reads, &mut writes);
706 }
707
708 let mut buffers = Vec::new();
709 reserve_grid_sync_vec(
710 &mut buffers,
711 source.buffers().len(),
712 "grid-sync segment buffers",
713 )?;
714 for buffer in source.buffers() {
715 let name = Ident::from(buffer.name());
716 let reads_this = reads.contains(&name);
717 let writes_this = writes.contains(&name);
718 let readwrite_passthrough = matches!(buffer.access(), BufferAccess::ReadWrite)
719 && !buffer.is_output()
720 && !buffer.is_pipeline_live_out()
721 && !reads_this
722 && !writes_this;
723
724 if !reads_this && !writes_this && !readwrite_passthrough {
725 continue;
726 }
727
728 let mut rewritten = buffer.clone();
729 if matches!(rewritten.access(), BufferAccess::Workgroup) {
730 buffers.push(rewritten);
731 continue;
732 }
733
734 let is_source_output = buffer.is_output() || buffer.is_pipeline_live_out();
741 let earlier_segment_wrote_output = is_source_output
742 && first_writer
743 .get(&name)
744 .is_some_and(|&first| first < segment_idx);
745
746 let access = if readwrite_passthrough {
747 BufferAccess::ReadWrite
748 } else if earlier_segment_wrote_output && writes_this {
749 BufferAccess::ReadWrite
753 } else {
754 match (reads_this, writes_this) {
755 (true, true) => BufferAccess::ReadWrite,
756 (true, false) => BufferAccess::ReadOnly,
757 (false, true) => BufferAccess::WriteOnly,
758 (false, false) => BufferAccess::ReadWrite,
759 }
760 };
761 rewrite_segment_buffer_access(&mut rewritten, access);
762 rewritten.is_output = false;
771 rewritten.pipeline_live_out = false;
772 buffers.push(rewritten);
773 }
774
775 Ok(segment.with_rewritten_buffers(buffers))
776}
777
778fn rewrite_segment_buffer_access(buffer: &mut BufferDecl, access: BufferAccess) {
779 buffer.kind = match &access {
780 BufferAccess::ReadOnly => MemoryKind::Readonly,
781 BufferAccess::Uniform => MemoryKind::Uniform,
782 BufferAccess::Workgroup => MemoryKind::Shared,
783 _ => MemoryKind::Global,
784 };
785 buffer.access = access;
786}
787
788fn segment_input_names(segment: &Program) -> Result<Vec<Ident>, BackendError> {
789 let mut names = Vec::new();
790 reserve_grid_sync_vec(
791 &mut names,
792 segment.buffers().len(),
793 "grid-sync segment input names",
794 )?;
795 for buffer in segment.buffers() {
796 if matches!(buffer.access(), BufferAccess::Workgroup) {
797 continue;
798 }
799 if segment_buffer_consumes_input(buffer) {
800 names.push(Ident::from(buffer.name()));
801 }
802 }
803 Ok(names)
804}
805
806fn segment_output_names(segment: &Program) -> Result<Vec<Ident>, BackendError> {
807 let mut names = Vec::new();
808 reserve_grid_sync_vec(
809 &mut names,
810 segment.buffers().len(),
811 "grid-sync segment output names",
812 )?;
813 for buffer in segment.buffers() {
814 if matches!(buffer.access(), BufferAccess::Workgroup) {
815 continue;
816 }
817 if segment_buffer_produces_output(buffer) {
818 names.push(Ident::from(buffer.name()));
819 }
820 }
821 Ok(names)
822}
823
824fn original_input_names(program: &Program) -> Result<Vec<Ident>, BackendError> {
825 segment_input_names(program)
826}
827
828fn original_output_names(program: &Program) -> Result<Vec<Ident>, BackendError> {
829 segment_output_names(program)
830}
831
832fn segment_buffer_consumes_input(buffer: &BufferDecl) -> bool {
833 if buffer.is_output() || buffer.is_pipeline_live_out() {
834 return false;
835 }
836 matches!(
837 buffer.access(),
838 BufferAccess::ReadOnly | BufferAccess::ReadWrite | BufferAccess::Uniform
839 )
840}
841
842fn segment_buffer_produces_output(buffer: &BufferDecl) -> bool {
843 buffer.is_output()
844 || buffer.is_pipeline_live_out()
845 || matches!(
846 buffer.access(),
847 BufferAccess::ReadWrite | BufferAccess::WriteOnly
848 )
849}
850
851fn collect_segment_buffer_targets(
852 node: &Node,
853 reads: &mut HashSet<Ident>,
854 writes: &mut HashSet<Ident>,
855) {
856 match node {
857 Node::Let { value, .. } | Node::Assign { value, .. } => {
858 collect_segment_expr_targets(value, reads, writes);
859 }
860 Node::Store {
861 buffer,
862 index,
863 value,
864 } => {
865 writes.insert(Ident::from(buffer));
866 collect_segment_expr_targets(index, reads, writes);
867 collect_segment_expr_targets(value, reads, writes);
868 }
869 Node::If {
870 cond,
871 then,
872 otherwise,
873 } => {
874 collect_segment_expr_targets(cond, reads, writes);
875 for child in then.iter().chain(otherwise.iter()) {
876 collect_segment_buffer_targets(child, reads, writes);
877 }
878 }
879 Node::Loop { from, to, body, .. } => {
880 collect_segment_expr_targets(from, reads, writes);
881 collect_segment_expr_targets(to, reads, writes);
882 for child in body {
883 collect_segment_buffer_targets(child, reads, writes);
884 }
885 }
886 Node::Block(body) => {
887 for child in body {
888 collect_segment_buffer_targets(child, reads, writes);
889 }
890 }
891 Node::Region { body, .. } => {
892 for child in body.iter() {
893 collect_segment_buffer_targets(child, reads, writes);
894 }
895 }
896 Node::AllReduce { buffer, .. } | Node::Broadcast { buffer, .. } => {
897 reads.insert(buffer.clone());
898 writes.insert(buffer.clone());
899 }
900 Node::AllGather { input, output, .. } | Node::ReduceScatter { input, output, .. } => {
901 reads.insert(input.clone());
902 writes.insert(output.clone());
903 }
904 Node::IndirectDispatch { .. }
905 | Node::Return
906 | Node::Barrier { .. }
907 | Node::AsyncLoad { .. }
908 | Node::AsyncStore { .. }
909 | Node::AsyncWait { .. }
910 | Node::Trap { .. }
911 | Node::Resume { .. }
912 | Node::Opaque(_) => {}
913 _ => {}
914 }
915}
916
917fn collect_segment_expr_targets(
918 expr: &Expr,
919 reads: &mut HashSet<Ident>,
920 writes: &mut HashSet<Ident>,
921) {
922 vyre_foundation::visit::visit_expr_buffer_accesses(expr, |access, buffer| {
923 reads.insert(buffer.clone());
924 if access == vyre_foundation::visit::ExprBufferAccess::Atomic {
925 writes.insert(buffer.clone());
926 }
927 });
928}
929
930pub fn dispatch_with_grid_sync_split(
950 backend: &dyn VyreBackend,
951 program: &Program,
952 inputs: &[&[u8]],
953 config: &DispatchConfig,
954) -> Result<Vec<Vec<u8>>, BackendError> {
955 let mut outputs = Vec::new();
956 reserve_grid_sync_vec(
957 &mut outputs,
958 program.output_buffer_indices().len().max(1),
959 "grid-sync final outputs",
960 )?;
961 dispatch_with_grid_sync_split_into(backend, program, inputs, config, &mut outputs)?;
962 Ok(outputs)
963}
964
965pub fn dispatch_with_grid_sync_split_timed(
970 backend: &dyn VyreBackend,
971 program: &Program,
972 inputs: &[&[u8]],
973 config: &DispatchConfig,
974) -> Result<TimedDispatchResult, BackendError> {
975 let started = std::time::Instant::now();
976 let outputs = dispatch_with_grid_sync_split(backend, program, inputs, config)?;
977 Ok(TimedDispatchResult {
978 outputs,
979 wall_ns: elapsed_wall_ns(started)?,
980 device_ns: None,
981 enqueue_ns: None,
982 wait_ns: None,
983 })
984}
985
986pub fn dispatch_resident_with_grid_sync_split_timed(
996 backend: &dyn VyreBackend,
997 program: &Program,
998 resources: &[Resource],
999 config: &DispatchConfig,
1000) -> Result<TimedDispatchResult, BackendError> {
1001 if !contains_grid_sync(program) {
1008 return backend.dispatch_resident_timed(program, resources, config);
1009 }
1010 let segments = try_split_on_grid_sync(program)?;
1011 if segments.is_empty() {
1012 return Err(BackendError::InvalidProgram {
1013 fix: "Fix: program contains GridSync barrier but split_on_grid_sync produced 0 \
1014 segments. This is a grid_sync invariant bug - split_on_grid_sync must \
1015 always return at least one segment."
1016 .to_string(),
1017 });
1018 }
1019 let started = std::time::Instant::now();
1020 let mut final_outputs = Vec::new();
1021 let mut device_ns = Some(0_u64);
1022 let mut enqueue_ns = Some(0_u64);
1023 let mut wait_ns = Some(0_u64);
1024 for (segment_idx, segment) in segments.iter().enumerate() {
1025 let timed = backend
1026 .dispatch_resident_timed(segment, resources, config)
1027 .map_err(|error| grid_sync_segment_error(error, segment_idx, segments.len()))?;
1028 if segment_idx + 1 == segments.len() {
1029 final_outputs = timed.outputs;
1030 }
1031 device_ns = crate::accounting::sum_optional_timing(
1032 device_ns,
1033 timed.device_ns,
1034 "device timing",
1035 "grid-sync segmented",
1036 "per-segment",
1037 )?;
1038 enqueue_ns = crate::accounting::sum_optional_timing(
1039 enqueue_ns,
1040 timed.enqueue_ns,
1041 "enqueue timing",
1042 "grid-sync segmented",
1043 "per-segment",
1044 )?;
1045 wait_ns = crate::accounting::sum_optional_timing(
1046 wait_ns,
1047 timed.wait_ns,
1048 "wait timing",
1049 "grid-sync segmented",
1050 "per-segment",
1051 )?;
1052 }
1053 Ok(TimedDispatchResult {
1054 outputs: final_outputs,
1055 wall_ns: elapsed_wall_ns(started)?,
1056 device_ns,
1057 enqueue_ns,
1058 wait_ns,
1059 })
1060}
1061
1062fn elapsed_wall_ns(started: std::time::Instant) -> Result<u64, BackendError> {
1063 u64::try_from(started.elapsed().as_nanos()).map_err(|error| BackendError::InvalidProgram {
1064 fix: format!(
1065 "Fix: grid-sync segmented wall timing cannot fit u64 nanoseconds: {error}. Split telemetry windows or report per-segment timing."
1066 ),
1067 })
1068}
1069
1070fn seed_backend_allocated_segment_inputs<'a>(
1071 program: &Program,
1072 segments: &[PlannedGridSyncSegment],
1073 current_inputs: &mut HashMap<Ident, GridSyncInput<'a>>,
1074) -> Result<(), BackendError> {
1075 for name in segments
1076 .iter()
1077 .flat_map(|segment| segment.input_names.iter())
1078 {
1079 if current_inputs.contains_key(name) {
1080 continue;
1081 }
1082 let Some(buffer) = program.buffer(name.as_str()) else {
1083 return Err(BackendError::InvalidProgram {
1084 fix: format!(
1085 "Fix: grid-sync segment references undeclared input `{name}`. Rebuild the split from a Program whose buffer table covers every expression dependency."
1086 ),
1087 });
1088 };
1089 if !buffer.is_backend_allocated_output() {
1090 continue;
1091 }
1092 let static_len = buffer
1093 .static_byte_len()
1094 .map_err(|error| BackendError::InvalidProgram {
1095 fix: format!("Fix: cannot seed grid-sync output `{name}`: {error}"),
1096 })?;
1097 let byte_len = static_len
1098 .or_else(|| buffer.output_byte_range().map(|range| range.end))
1099 .ok_or_else(|| BackendError::InvalidProgram {
1100 fix: format!(
1101 "Fix: grid-sync output `{name}` is read-modify-written before its first split output but has no static byte size. Declare a count or output byte range."
1102 ),
1103 })?;
1104 let mut zeroed = Vec::new();
1105 reserve_grid_sync_vec(
1106 &mut zeroed,
1107 byte_len,
1108 "grid-sync backend-allocated output seed",
1109 )?;
1110 zeroed.resize(byte_len, 0);
1111 current_inputs.insert(name.clone(), GridSyncInput::Owned(zeroed));
1112 }
1113 Ok(())
1114}
1115
1116fn dispatch_grid_sync_split_generic<D>(
1122 program: &Program,
1123 inputs: &[&[u8]],
1124 config: &DispatchConfig,
1125 outputs: &mut OutputBuffers,
1126 mut dispatch_segment: D,
1127) -> Result<(), BackendError>
1128where
1129 D: FnMut(&Program, &[&[u8]], &DispatchConfig, &mut OutputBuffers) -> Result<(), BackendError>,
1130{
1131 if !contains_grid_sync(program) {
1138 return dispatch_segment(program, inputs, config, outputs);
1139 }
1140 let segments = plan_host_grid_sync_segments(program)?;
1141 if segments.is_empty() {
1142 return Err(BackendError::InvalidProgram {
1143 fix: "Fix: program contains GridSync barrier but split_on_grid_sync produced 0 \
1144 segments. This is a grid_sync invariant bug - split_on_grid_sync must \
1145 always return at least one segment."
1146 .to_string(),
1147 });
1148 }
1149 crate::observability::record_grid_sync_split(segments.len());
1150 let initial_input_names = original_input_names(program)?;
1157 if inputs.len() != initial_input_names.len() {
1158 return Err(BackendError::InvalidProgram {
1159 fix: format!(
1160 "Fix: grid-sync split expected {} initial input buffer(s) but received {}. Rebuild the dispatch inputs from the Program buffer declarations before splitting.",
1161 initial_input_names.len(),
1162 inputs.len()
1163 ),
1164 });
1165 }
1166 let mut current_inputs: HashMap<Ident, GridSyncInput<'_>> = HashMap::new();
1167 reserve_grid_sync_hash_map(
1168 &mut current_inputs,
1169 program.buffers().len(),
1170 "grid-sync rotating input map",
1171 )?;
1172 for (name, bytes) in initial_input_names.into_iter().zip(inputs.iter().copied()) {
1173 current_inputs.insert(name, GridSyncInput::Borrowed(bytes));
1174 }
1175 seed_backend_allocated_segment_inputs(program, &segments, &mut current_inputs)?;
1176 let mut segment_outputs = Vec::new();
1177 reserve_grid_sync_vec(
1178 &mut segment_outputs,
1179 outputs.capacity().max(1),
1180 "grid-sync intermediate outputs",
1181 )?;
1182 let final_output_names = original_output_names(program)?;
1183
1184 let iterations =
1206 crate::fixpoint_iterations::resolve_fixpoint_iterations(config, "grid-sync split")?;
1207 let mut segment_config = config.clone();
1208 segment_config.fixpoint_iterations = Some(1);
1209
1210 let mut prev_fingerprint: Option<u64> = None;
1218 for _ in 0..iterations {
1219 for (segment_idx, segment) in segments.iter().enumerate() {
1220 let borrowed = borrowed_grid_sync_inputs_by_name(segment, ¤t_inputs)?;
1221 dispatch_segment(
1222 &segment.program,
1223 borrowed.as_slice(),
1224 &segment_config,
1225 &mut segment_outputs,
1226 )
1227 .map_err(|error| grid_sync_segment_error(error, segment_idx, segments.len()))?;
1228 drop(borrowed);
1229 refresh_named_outputs(segment, &mut segment_outputs, &mut current_inputs)?;
1230 }
1231 let fingerprint = owned_accumulator_fingerprint(¤t_inputs);
1232 if prev_fingerprint == Some(fingerprint) {
1233 break;
1234 }
1235 prev_fingerprint = Some(fingerprint);
1236 }
1237 collect_final_named_outputs(&final_output_names, &mut current_inputs, outputs)?;
1238 Ok(())
1239}
1240
1241pub fn dispatch_with_grid_sync_split_into(
1252 backend: &dyn VyreBackend,
1253 program: &Program,
1254 inputs: &[&[u8]],
1255 config: &DispatchConfig,
1256 outputs: &mut OutputBuffers,
1257) -> Result<(), BackendError> {
1258 dispatch_grid_sync_split_generic(program, inputs, config, outputs, |p, i, c, o| {
1259 backend.dispatch_borrowed_into(p, i, c, o)
1260 })
1261}
1262
1263pub fn dispatch_with_grid_sync_split_via_into<F>(
1283 program: &Program,
1284 inputs: &[&[u8]],
1285 config: &DispatchConfig,
1286 dispatch: &F,
1287 outputs: &mut OutputBuffers,
1288) -> Result<(), BackendError>
1289where
1290 F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
1291{
1292 dispatch_grid_sync_split_generic(program, inputs, config, outputs, |p, i, c, o| {
1293 dispatch(p, i, c.grid_override, o).map_err(BackendError::new)
1294 })
1295}
1296
1297pub fn dispatch_with_grid_sync_split_via<F>(
1302 program: &Program,
1303 inputs: &[&[u8]],
1304 config: &DispatchConfig,
1305 dispatch: &F,
1306) -> Result<Vec<Vec<u8>>, BackendError>
1307where
1308 F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
1309{
1310 let mut outputs = Vec::new();
1311 reserve_grid_sync_vec(
1312 &mut outputs,
1313 program.output_buffer_indices().len().max(1),
1314 "grid-sync via final outputs",
1315 )?;
1316 dispatch_with_grid_sync_split_via_into(program, inputs, config, dispatch, &mut outputs)?;
1317 Ok(outputs)
1318}
1319
1320pub fn dispatch_resident_grid_sync_fixpoint_into(
1358 backend: &dyn VyreBackend,
1359 program: &Program,
1360 inputs: &[&[u8]],
1361 config: &DispatchConfig,
1362 outputs: &mut OutputBuffers,
1363) -> Result<(), BackendError> {
1364 if !contains_grid_sync(program) {
1371 return backend.dispatch_borrowed_into(program, inputs, config, outputs);
1372 }
1373 let segments = try_split_on_grid_sync(program)?;
1374 if segments.is_empty() {
1375 return Err(BackendError::InvalidProgram {
1376 fix: "Fix: program contains GridSync barrier but split_on_grid_sync produced 0 \
1377 segments. This is a grid_sync invariant bug - split_on_grid_sync must \
1378 always return at least one segment."
1379 .to_string(),
1380 });
1381 }
1382 crate::observability::record_grid_sync_split(segments.len());
1383
1384 let resident = allocate_resident_program_resources(backend, program, inputs)?;
1388 let result =
1389 run_resident_grid_sync_fixpoint(backend, program, &segments, &resident, config, outputs);
1390 let free_result = free_resident_program_resources(backend, resident);
1392 result.and(free_result)
1393}
1394
1395struct ResidentProgramResources {
1399 ordered: Vec<Resource>,
1402 by_name: HashMap<Ident, (Resource, usize)>,
1405}
1406
1407fn allocate_resident_program_resources(
1413 backend: &dyn VyreBackend,
1414 program: &Program,
1415 inputs: &[&[u8]],
1416) -> Result<ResidentProgramResources, BackendError> {
1417 let plan = BindingPlan::from_borrowed_inputs(program, inputs)?;
1418 let mut ordered = Vec::new();
1419 reserve_grid_sync_vec(
1420 &mut ordered,
1421 plan.bindings.len(),
1422 "resident grid-sync resources",
1423 )?;
1424 let mut by_name = HashMap::new();
1425 reserve_grid_sync_hash_map(
1426 &mut by_name,
1427 plan.bindings.len(),
1428 "resident grid-sync resource name map",
1429 )?;
1430 for binding in &plan.bindings {
1431 if binding.role == BindingRole::Shared {
1432 continue;
1433 }
1434 let byte_len = resident_binding_byte_len(binding, inputs)?;
1443 let alloc_len = byte_len.max(binding.element_size.max(1));
1444 let resource = backend.allocate_resident(alloc_len)?;
1445 match binding.input_index {
1449 Some(index) if !inputs.get(index).copied().unwrap_or(&[]).is_empty() => {
1450 let bytes = inputs[index];
1451 backend.upload_resident(&resource, bytes)?;
1452 }
1453 _ => {
1454 let zeros = zeroed_upload_buffer(alloc_len)?;
1455 backend.upload_resident(&resource, &zeros)?;
1456 }
1457 }
1458 by_name.insert(
1459 Ident::from(binding.name.as_ref()),
1460 (resource.clone(), byte_len),
1461 );
1462 ordered.push(resource);
1463 }
1464 Ok(ResidentProgramResources { ordered, by_name })
1465}
1466
1467fn resident_binding_byte_len(binding: &Binding, inputs: &[&[u8]]) -> Result<usize, BackendError> {
1470 if let Some(index) = binding.input_index {
1471 if let Some(bytes) = inputs.get(index) {
1472 return Ok(bytes.len());
1473 }
1474 }
1475 binding.static_byte_len.ok_or_else(|| BackendError::InvalidProgram {
1476 fix: format!(
1477 "Fix: resident grid-sync output buffer `{}` has no static byte length; dynamic-sized outputs are not supported on the resident grid-sync path. Declare a fixed `count` on the buffer or route this program through dispatch_with_grid_sync_split_into.",
1478 binding.name
1479 ),
1480 })
1481}
1482
1483fn zeroed_upload_buffer(byte_len: usize) -> Result<Vec<u8>, BackendError> {
1486 let mut zeros = Vec::new();
1487 crate::allocation::try_reserve_vec_to_capacity(&mut zeros, byte_len).map_err(|error| {
1488 BackendError::InvalidProgram {
1489 fix: format!(
1490 "Fix: failed to reserve a {byte_len}-byte zero-init staging buffer for a resident grid-sync output: {error}. Shard the program into smaller buffers."
1491 ),
1492 }
1493 })?;
1494 zeros.resize(byte_len, 0);
1495 Ok(zeros)
1496}
1497
1498fn run_resident_grid_sync_fixpoint(
1502 backend: &dyn VyreBackend,
1503 program: &Program,
1504 segments: &[Program],
1505 resident: &ResidentProgramResources,
1506 config: &DispatchConfig,
1507 outputs: &mut OutputBuffers,
1508) -> Result<(), BackendError> {
1509 let iterations = crate::fixpoint_iterations::resolve_fixpoint_iterations(
1510 config,
1511 "resident grid-sync split",
1512 )?;
1513 let repeat_count = u32::try_from(iterations).map_err(|error| BackendError::InvalidProgram {
1514 fix: format!(
1515 "Fix: resident grid-sync fixpoint iteration count {iterations} does not fit u32: {error}."
1516 ),
1517 })?;
1518
1519 let mut steps = Vec::new();
1522 reserve_grid_sync_vec(&mut steps, segments.len(), "resident grid-sync steps")?;
1523 for segment in segments {
1524 steps.push(ResidentDispatchStep {
1525 program: segment,
1526 resources: resident.ordered.as_slice(),
1527 grid_override: config.grid_override,
1528 workgroup_override: config.workgroup_override,
1532 });
1533 }
1534
1535 let output_names = original_output_names(program)?;
1538 let mut read_ranges = Vec::new();
1539 reserve_grid_sync_vec(
1540 &mut read_ranges,
1541 output_names.len(),
1542 "resident grid-sync read ranges",
1543 )?;
1544 for name in &output_names {
1545 let (resource, byte_len) =
1546 resident.by_name.get(name).ok_or_else(|| BackendError::InvalidProgram {
1547 fix: format!(
1548 "Fix: resident grid-sync final output `{name}` has no resident resource; it was not declared as a non-shared program buffer."
1549 ),
1550 })?;
1551 read_ranges.push(ResidentReadRange {
1552 resource,
1553 byte_offset: 0,
1554 byte_len: *byte_len,
1555 });
1556 }
1557
1558 while outputs.len() < output_names.len() {
1561 outputs.push(Vec::new());
1562 }
1563 outputs.truncate(output_names.len());
1564 for slot in outputs.iter_mut() {
1565 slot.clear();
1566 }
1567 let mut output_refs: Vec<&mut Vec<u8>> = outputs.iter_mut().collect();
1568
1569 backend.dispatch_resident_repeated_sequence_read_ranges_into(
1570 &[],
1571 &steps,
1572 repeat_count,
1573 &read_ranges,
1574 output_refs.as_mut_slice(),
1575 )
1576}
1577
1578fn free_resident_program_resources(
1582 backend: &dyn VyreBackend,
1583 resident: ResidentProgramResources,
1584) -> Result<(), BackendError> {
1585 let ResidentProgramResources { ordered, by_name } = resident;
1586 drop(by_name);
1589 let mut first_error: Option<BackendError> = None;
1590 for resource in ordered {
1591 if let Err(error) = backend.free_resident(resource) {
1592 if first_error.is_none() {
1593 first_error = Some(error);
1594 }
1595 }
1596 }
1597 match first_error {
1598 Some(error) => Err(error),
1599 None => Ok(()),
1600 }
1601}
1602
1603fn reserve_grid_sync_vec<T>(
1604 vec: &mut Vec<T>,
1605 capacity: usize,
1606 field: &'static str,
1607) -> Result<(), BackendError> {
1608 crate::allocation::try_reserve_vec_to_capacity(vec, capacity).map_err(|error| {
1609 BackendError::InvalidProgram {
1610 fix: format!(
1611 "Fix: failed to reserve {field} for {capacity} entries during grid-sync dispatch splitting: {error}. Split the program into fewer grid-sync segments or run on a backend with native grid sync."
1612 ),
1613 }
1614 })
1615}
1616
1617fn reserve_grid_sync_hash_map<K, V>(
1618 map: &mut HashMap<K, V>,
1619 capacity: usize,
1620 field: &'static str,
1621) -> Result<(), BackendError>
1622where
1623 K: Eq + std::hash::Hash,
1624{
1625 map.try_reserve(capacity)
1626 .map_err(|error| BackendError::InvalidProgram {
1627 fix: format!(
1628 "Fix: failed to reserve {field} for {capacity} entries during grid-sync dispatch splitting: {error}. Split the program into fewer grid-sync segments or run on a backend with native grid sync."
1629 ),
1630 })
1631}
1632
1633fn reserve_grid_sync_hash_set<T>(
1634 set: &mut HashSet<T>,
1635 capacity: usize,
1636 field: &'static str,
1637) -> Result<(), BackendError>
1638where
1639 T: Eq + std::hash::Hash,
1640{
1641 set.try_reserve(capacity)
1642 .map_err(|error| BackendError::InvalidProgram {
1643 fix: format!(
1644 "Fix: failed to reserve {field} for {capacity} entries during grid-sync dispatch splitting: {error}. Split the program into fewer grid-sync segments or run on a backend with native grid sync."
1645 ),
1646 })
1647}
1648
1649fn borrowed_grid_sync_inputs<'a>(
1650 inputs: &'a [GridSyncInput<'a>],
1651) -> Result<SmallVec<[&'a [u8]; 8]>, BackendError> {
1652 let mut borrowed = SmallVec::<[&[u8]; 8]>::new();
1653 borrowed.try_reserve(inputs.len()).map_err(|error| {
1654 BackendError::InvalidProgram {
1655 fix: format!(
1656 "Fix: failed to reserve grid-sync borrowed input slices for {} input(s): {error}. Split the program into fewer grid-sync live buffers or run on a backend with native grid sync.",
1657 inputs.len()
1658 ),
1659 }
1660 })?;
1661 borrowed.extend(inputs.iter().map(GridSyncInput::as_slice));
1662 Ok(borrowed)
1663}
1664
1665fn borrowed_grid_sync_inputs_by_name<'a>(
1666 segment: &PlannedGridSyncSegment,
1667 inputs: &'a HashMap<Ident, GridSyncInput<'a>>,
1668) -> Result<SmallVec<[&'a [u8]; 8]>, BackendError> {
1669 let mut borrowed = SmallVec::<[&[u8]; 8]>::new();
1670 borrowed
1671 .try_reserve(segment.input_names.len())
1672 .map_err(|error| BackendError::InvalidProgram {
1673 fix: format!(
1674 "Fix: failed to reserve grid-sync borrowed input slices for {} segment input(s): {error}. Split the program into fewer grid-sync live buffers or run on a backend with native grid sync.",
1675 segment.input_names.len()
1676 ),
1677 })?;
1678 for name in &segment.input_names {
1679 let input = inputs.get(name).ok_or_else(|| BackendError::InvalidProgram {
1680 fix: format!(
1681 "Fix: grid-sync segment input `{name}` has no bytes from caller input or a prior segment output. Ensure every cross-segment read is written before the GridSync barrier."
1682 ),
1683 })?;
1684 borrowed.push(input.as_slice());
1685 }
1686 Ok(borrowed)
1687}
1688
1689fn owned_accumulator_fingerprint(inputs: &HashMap<Ident, GridSyncInput<'_>>) -> u64 {
1701 const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
1702 const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
1703 let mut combined: u64 = 0;
1704 for (name, input) in inputs {
1705 let GridSyncInput::Owned(bytes) = input else {
1706 continue;
1707 };
1708 let mut hash = FNV_OFFSET;
1709 for byte in name.as_str().as_bytes() {
1710 hash ^= u64::from(*byte);
1711 hash = hash.wrapping_mul(FNV_PRIME);
1712 }
1713 hash ^= 0xff;
1715 hash = hash.wrapping_mul(FNV_PRIME);
1716 for byte in bytes.iter() {
1717 hash ^= u64::from(*byte);
1718 hash = hash.wrapping_mul(FNV_PRIME);
1719 }
1720 combined ^= hash;
1721 }
1722 combined
1723}
1724
1725fn grid_sync_segment_error(
1726 error: BackendError,
1727 segment_idx: usize,
1728 segment_count: usize,
1729) -> BackendError {
1730 match error {
1731 BackendError::InvalidProgram { fix } => BackendError::InvalidProgram {
1732 fix: format!(
1733 "Fix: grid-sync split segment {segment_idx} of {segment_count} dispatch failed: {fix}"
1734 ),
1735 },
1736 other => other,
1737 }
1738}
1739
1740enum GridSyncInput<'a> {
1741 Borrowed(&'a [u8]),
1742 Owned(Vec<u8>),
1743}
1744
1745impl GridSyncInput<'_> {
1746 fn as_slice(&self) -> &[u8] {
1747 match self {
1748 Self::Borrowed(bytes) => bytes,
1749 Self::Owned(bytes) => bytes.as_slice(),
1750 }
1751 }
1752
1753 fn refresh_from_output(&mut self, bytes: &mut Vec<u8>) -> Result<(), BackendError> {
1754 match self {
1755 Self::Borrowed(_) => {
1756 let mut owned = Vec::new();
1757 reserve_grid_sync_vec(&mut owned, bytes.len(), "grid-sync readwrite input")?;
1758 owned.extend_from_slice(bytes);
1759 *self = Self::Owned(owned);
1760 }
1761 Self::Owned(owned) => {
1762 std::mem::swap(owned, bytes);
1763 }
1764 }
1765 Ok(())
1766 }
1767}
1768
1769fn refresh_named_outputs<'a>(
1770 segment: &PlannedGridSyncSegment,
1771 outputs: &mut Vec<Vec<u8>>,
1772 inputs: &mut HashMap<Ident, GridSyncInput<'a>>,
1773) -> Result<(), BackendError> {
1774 if outputs.len() != segment.output_names.len() {
1775 return Err(BackendError::InvalidProgram {
1776 fix: format!(
1777 "Fix: grid-sync split segment produced {} output slot(s) but the planned buffer map expected {}. Preserve segment output declaration order when dispatching split kernels.",
1778 outputs.len(),
1779 segment.output_names.len()
1780 ),
1781 });
1782 }
1783 for (name, bytes) in segment.output_names.iter().cloned().zip(outputs.iter_mut()) {
1784 match inputs.get_mut(&name) {
1785 Some(slot) => slot.refresh_from_output(bytes)?,
1786 None => {
1787 let mut owned = GridSyncInput::Owned(Vec::new());
1788 owned.refresh_from_output(bytes)?;
1789 inputs.insert(name, owned);
1790 }
1791 }
1792 }
1793 for output in outputs {
1794 output.clear();
1795 }
1796 Ok(())
1797}
1798
1799fn collect_final_named_outputs<'a>(
1800 final_output_names: &[Ident],
1801 inputs: &mut HashMap<Ident, GridSyncInput<'a>>,
1802 outputs: &mut OutputBuffers,
1803) -> Result<(), BackendError> {
1804 let mut final_outputs = Vec::new();
1805 reserve_grid_sync_vec(
1806 &mut final_outputs,
1807 final_output_names.len(),
1808 "grid-sync final named outputs",
1809 )?;
1810 for name in final_output_names {
1811 let output = inputs
1812 .remove(name)
1813 .ok_or_else(|| BackendError::InvalidProgram {
1814 fix: format!(
1815 "Fix: grid-sync final output `{name}` was not produced by any split segment."
1816 ),
1817 })?;
1818 match output {
1819 GridSyncInput::Owned(bytes) => final_outputs.push(bytes),
1820 GridSyncInput::Borrowed(bytes) => {
1821 let mut owned = Vec::new();
1822 reserve_grid_sync_vec(&mut owned, bytes.len(), "grid-sync borrowed final output")?;
1823 owned.extend_from_slice(bytes);
1824 final_outputs.push(owned);
1825 }
1826 }
1827 }
1828 crate::replace_output_buffers_preserving_slots(final_outputs, outputs);
1829 Ok(())
1830}
1831
1832fn refresh_readwrite_inputs(
1839 segment: &Program,
1840 outputs: &mut Vec<Vec<u8>>,
1841 inputs: &mut [GridSyncInput<'_>],
1842) -> Result<(), BackendError> {
1843 use vyre_foundation::ir::BufferAccess;
1844 let mut input_idx = 0usize;
1850 let mut output_idx = 0usize;
1851 for buffer in segment.buffers() {
1852 if matches!(buffer.access(), BufferAccess::Workgroup) {
1853 continue;
1854 }
1855 let is_output_buffer = buffer.is_output();
1856 let is_readwrite = matches!(buffer.access(), BufferAccess::ReadWrite);
1857
1858 if is_readwrite && !is_output_buffer {
1862 if let (Some(slot), Some(bytes)) =
1863 (inputs.get_mut(input_idx), outputs.get_mut(output_idx))
1864 {
1865 slot.refresh_from_output(bytes)?;
1866 }
1867 }
1868
1869 if !is_output_buffer {
1871 input_idx += 1;
1872 }
1873 if is_readwrite {
1876 output_idx += 1;
1877 }
1878 }
1879 for output in outputs {
1880 output.clear();
1881 }
1882 Ok(())
1883}
1884
1885#[cfg(test)]
1886mod tests {
1887 use super::*;
1888 use std::sync::atomic::{AtomicUsize, Ordering};
1889 use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr};
1890
1891 fn buffer() -> BufferDecl {
1892 BufferDecl::storage("buf", 0, BufferAccess::ReadWrite, DataType::U32).with_count(4)
1893 }
1894
1895 fn region(generator: &str, body: Vec<Node>) -> Node {
1896 Node::Region {
1897 generator: Ident::from(generator),
1898 source_region: None,
1899 body: Arc::new(body),
1900 }
1901 }
1902
1903 #[test]
1904 fn grid_sync_release_paths_use_fallible_split_storage() {
1905 let source = include_str!("grid_sync.rs");
1906 let production = source
1907 .split("#[cfg(test)]")
1908 .next()
1909 .expect("Fix: grid-sync production source must precede tests");
1910
1911 assert!(
1912 production.contains("pub fn try_split_on_grid_sync")
1913 && production.contains("fn reserve_grid_sync_vec")
1914 && production.contains("try_reserve_vec_to_capacity"),
1915 "Fix: grid-sync splitting must expose fallible segment/input/output scratch reservation."
1916 );
1917 assert!(
1918 production.contains("let segments = try_split_on_grid_sync(program)?")
1919 && !production.contains("let segments = split_on_grid_sync(program);"),
1920 "Fix: production grid-sync dispatch paths must use fallible splitting, not the legacy infallible helper."
1921 );
1922 assert!(
1923 !production.contains("Vec::with_capacity"),
1924 "Fix: production grid-sync splitting must not allocate dispatch scratch infallibly."
1925 );
1926 assert!(
1927 !production.contains(".as_nanos() as u64")
1928 && !production.contains("segmented timing overflowed u64"),
1929 "Fix: production grid-sync timing telemetry must return typed errors instead of truncating or panicking."
1930 );
1931 }
1932
1933 fn inner_len(program: &Program) -> usize {
1935 entry_sequence(program).len()
1936 }
1937
1938 #[test]
1939 fn no_grid_sync_returns_single_segment() {
1940 let program = Program::wrapped(
1941 vec![buffer()],
1942 [1, 1, 1],
1943 vec![region(
1944 "a",
1945 vec![Node::store("buf", Expr::u32(0), Expr::u32(1))],
1946 )],
1947 );
1948 assert!(!contains_grid_sync(&program));
1949 let segments = split_on_grid_sync(&program);
1950 assert_eq!(segments.len(), 1);
1951 assert_eq!(inner_len(&segments[0]), 1);
1953 }
1954
1955 #[test]
1956 fn one_grid_sync_splits_into_two() {
1957 let program = Program::wrapped(
1958 vec![buffer()],
1959 [1, 1, 1],
1960 vec![
1961 region("a", vec![Node::store("buf", Expr::u32(0), Expr::u32(1))]),
1962 Node::barrier_with_ordering(MemoryOrdering::GridSync),
1963 region("b", vec![Node::store("buf", Expr::u32(1), Expr::u32(2))]),
1964 ],
1965 );
1966 assert!(contains_grid_sync(&program));
1967 let segments = split_on_grid_sync(&program);
1968 assert_eq!(segments.len(), 2);
1969 assert_eq!(inner_len(&segments[0]), 1);
1970 assert_eq!(inner_len(&segments[1]), 1);
1971 }
1972
1973 #[test]
1974 fn block_nested_grid_sync_splits_into_two() {
1975 let program = Program::wrapped(
1976 vec![buffer()],
1977 [1, 1, 1],
1978 vec![Node::Block(vec![
1979 region("a", vec![Node::store("buf", Expr::u32(0), Expr::u32(1))]),
1980 Node::barrier_with_ordering(MemoryOrdering::GridSync),
1981 region("b", vec![Node::store("buf", Expr::u32(1), Expr::u32(2))]),
1982 ])],
1983 );
1984 assert!(contains_grid_sync(&program));
1985 let segments = split_on_grid_sync(&program);
1986 assert_eq!(segments.len(), 2);
1987 assert_eq!(inner_len(&segments[0]), 1);
1988 assert_eq!(inner_len(&segments[1]), 1);
1989 }
1990
1991 #[test]
1992 fn three_grid_syncs_split_into_four() {
1993 let program = Program::wrapped(
1994 vec![buffer()],
1995 [1, 1, 1],
1996 vec![
1997 region("a", vec![Node::Return]),
1998 Node::barrier_with_ordering(MemoryOrdering::GridSync),
1999 region("b", vec![Node::Return]),
2000 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2001 region("c", vec![Node::Return]),
2002 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2003 region("d", vec![Node::Return]),
2004 ],
2005 );
2006 let segments = split_on_grid_sync(&program);
2007 assert_eq!(segments.len(), 4);
2008 }
2009
2010 #[test]
2011 fn workgroup_barrier_does_not_split() {
2012 let program = Program::wrapped(
2013 vec![buffer()],
2014 [1, 1, 1],
2015 vec![
2016 region("a", vec![Node::Return]),
2017 Node::barrier_with_ordering(MemoryOrdering::SeqCst),
2018 region("b", vec![Node::Return]),
2019 ],
2020 );
2021 assert!(!contains_grid_sync(&program));
2022 let segments = split_on_grid_sync(&program);
2023 assert_eq!(segments.len(), 1);
2024 assert_eq!(inner_len(&segments[0]), 3);
2026 }
2027
2028 #[test]
2029 fn buffers_and_workgroup_size_propagate_to_each_segment() {
2030 let program = Program::wrapped(
2031 vec![buffer()],
2032 [256, 1, 1],
2033 vec![
2034 region("a", vec![Node::Return]),
2035 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2036 region("b", vec![Node::Return]),
2037 ],
2038 );
2039 let segments = split_on_grid_sync(&program);
2040 for seg in &segments {
2041 assert_eq!(seg.workgroup_size(), [256, 1, 1]);
2042 assert_eq!(seg.buffers().len(), 1);
2043 assert_eq!(seg.buffers()[0].name(), "buf");
2044 }
2045 }
2046
2047 #[test]
2048 fn refresh_readwrite_inputs_swaps_owned_buffers_after_first_segment() {
2049 let segment = Program::wrapped(vec![buffer()], [1, 1, 1], vec![Node::Return]);
2050 let initial = [1u8, 0, 0, 0];
2051 let mut inputs = [GridSyncInput::Borrowed(initial.as_slice())];
2052 let mut outputs = vec![Vec::with_capacity(8)];
2053 let output_ptr = outputs[0].as_ptr() as usize;
2054 outputs[0].extend_from_slice(&[2, 0, 0, 0]);
2055
2056 refresh_readwrite_inputs(&segment, &mut outputs, &mut inputs)
2057 .expect("Fix: test readwrite refresh should fit borrowed promotion storage");
2058
2059 let first_owned_ptr = match &inputs[0] {
2060 GridSyncInput::Owned(bytes) => {
2061 assert_eq!(bytes, &[2, 0, 0, 0]);
2062 bytes.as_ptr() as usize
2063 }
2064 GridSyncInput::Borrowed(_) => panic!("ReadWrite input must become owned after refresh"),
2065 };
2066 assert_eq!(outputs[0].as_ptr() as usize, output_ptr);
2067 assert!(outputs[0].is_empty());
2068
2069 outputs[0].extend_from_slice(&[3, 0, 0, 0]);
2070 let second_output_ptr = outputs[0].as_ptr() as usize;
2071 refresh_readwrite_inputs(&segment, &mut outputs, &mut inputs)
2072 .expect("Fix: test readwrite refresh should reuse owned storage");
2073
2074 match &inputs[0] {
2075 GridSyncInput::Owned(bytes) => {
2076 assert_eq!(bytes, &[3, 0, 0, 0]);
2077 assert_eq!(
2078 bytes.as_ptr() as usize,
2079 second_output_ptr,
2080 "owned ReadWrite input should take the backend output allocation instead of copying"
2081 );
2082 }
2083 GridSyncInput::Borrowed(_) => panic!("ReadWrite input must remain owned"),
2084 }
2085 assert_eq!(
2086 outputs[0].as_ptr() as usize,
2087 first_owned_ptr,
2088 "backend output slot should receive the previous owned input allocation for reuse"
2089 );
2090 }
2091
2092 struct ReuseCheckingBackend {
2093 calls: AtomicUsize,
2094 final_outputs_addr: usize,
2095 final_slot_addr: usize,
2096 }
2097
2098 impl crate::backend::private::Sealed for ReuseCheckingBackend {}
2099
2100 impl VyreBackend for ReuseCheckingBackend {
2101 fn id(&self) -> &'static str {
2102 "grid-sync-reuse-checking"
2103 }
2104
2105 fn dispatch(
2106 &self,
2107 _program: &Program,
2108 _inputs: &[Vec<u8>],
2109 _config: &DispatchConfig,
2110 ) -> Result<Vec<Vec<u8>>, BackendError> {
2111 unreachable!("test uses dispatch_borrowed_into")
2112 }
2113
2114 fn dispatch_borrowed_into(
2115 &self,
2116 _program: &Program,
2117 inputs: &[&[u8]],
2118 _config: &DispatchConfig,
2119 outputs: &mut OutputBuffers,
2120 ) -> Result<(), BackendError> {
2121 let call = self.calls.fetch_add(1, Ordering::SeqCst);
2122 if call == 1 && self.final_outputs_addr != 0 {
2123 assert_eq!(outputs.as_ptr() as usize, self.final_outputs_addr);
2124 assert_eq!(outputs[0].as_ptr() as usize, self.final_slot_addr);
2125 }
2126 if outputs.is_empty() {
2127 outputs.push(Vec::new());
2128 }
2129 outputs[0].clear();
2130 outputs[0].extend_from_slice(inputs[0]);
2131 if call == 0 {
2132 outputs[0][0] = 7;
2133 } else {
2134 outputs[0][0] = outputs[0][0].saturating_add(1);
2135 }
2136 Ok(())
2137 }
2138 }
2139
2140 #[test]
2141 fn split_into_preserves_caller_output_slot_after_named_output_collection() {
2142 let program = Program::wrapped(
2143 vec![buffer()],
2144 [1, 1, 1],
2145 vec![
2146 region("a", vec![Node::Return]),
2147 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2148 region("b", vec![Node::Return]),
2149 ],
2150 );
2151 let mut outputs = vec![Vec::with_capacity(8)];
2152 let outputs_addr = outputs.as_ptr() as usize;
2153 let slot_addr = outputs[0].as_ptr() as usize;
2154 let backend = ReuseCheckingBackend {
2155 calls: AtomicUsize::new(0),
2156 final_outputs_addr: 0,
2157 final_slot_addr: 0,
2158 };
2159 let input = [0u8, 0, 0, 0];
2160 dispatch_with_grid_sync_split_into(
2161 &backend,
2162 &program,
2163 &[input.as_slice()],
2164 &DispatchConfig::default(),
2165 &mut outputs,
2166 )
2167 .expect("Fix: grid-sync split should write into caller-owned output storage");
2168
2169 assert_eq!(backend.calls.load(Ordering::SeqCst), 2);
2170 assert_eq!(outputs, vec![vec![8, 0, 0, 0]]);
2171 assert_eq!(outputs.as_ptr() as usize, outputs_addr);
2172 assert_eq!(outputs[0].as_ptr() as usize, slot_addr);
2173 }
2174
2175 struct IncrementingBackend {
2183 calls: AtomicUsize,
2184 }
2185
2186 impl crate::backend::private::Sealed for IncrementingBackend {}
2187
2188 impl VyreBackend for IncrementingBackend {
2189 fn id(&self) -> &'static str {
2190 "grid-sync-incrementing"
2191 }
2192
2193 fn dispatch(
2194 &self,
2195 _program: &Program,
2196 _inputs: &[Vec<u8>],
2197 _config: &DispatchConfig,
2198 ) -> Result<Vec<Vec<u8>>, BackendError> {
2199 unreachable!("test uses dispatch_borrowed_into")
2200 }
2201
2202 fn dispatch_borrowed_into(
2203 &self,
2204 _program: &Program,
2205 inputs: &[&[u8]],
2206 config: &DispatchConfig,
2207 outputs: &mut OutputBuffers,
2208 ) -> Result<(), BackendError> {
2209 self.calls.fetch_add(1, Ordering::SeqCst);
2210 assert_eq!(
2213 config.fixpoint_iterations,
2214 Some(1),
2215 "segment dispatch must receive fixpoint_iterations=1; the outer split loop owns the iteration count"
2216 );
2217 if outputs.is_empty() {
2218 outputs.push(Vec::new());
2219 }
2220 outputs[0].clear();
2221 outputs[0].extend_from_slice(inputs[0]);
2222 outputs[0][0] = outputs[0][0].saturating_add(1);
2223 Ok(())
2224 }
2225 }
2226
2227 #[test]
2228 fn split_into_loops_whole_sequence_fixpoint_iterations_times() {
2229 let program = Program::wrapped(
2231 vec![buffer()],
2232 [1, 1, 1],
2233 vec![
2234 region("a", vec![Node::Return]),
2235 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2236 region("b", vec![Node::Return]),
2237 ],
2238 );
2239
2240 let backend = IncrementingBackend {
2242 calls: AtomicUsize::new(0),
2243 };
2244 let mut outputs = vec![Vec::new()];
2245 dispatch_with_grid_sync_split_into(
2246 &backend,
2247 &program,
2248 &[[0u8, 0, 0, 0].as_slice()],
2249 &DispatchConfig::default(),
2250 &mut outputs,
2251 )
2252 .expect("single-pass split dispatch");
2253 assert_eq!(backend.calls.load(Ordering::SeqCst), 2);
2254 assert_eq!(outputs, vec![vec![2, 0, 0, 0]]);
2255
2256 let backend = IncrementingBackend {
2261 calls: AtomicUsize::new(0),
2262 };
2263 let config = DispatchConfig {
2264 fixpoint_iterations: Some(3),
2265 ..DispatchConfig::default()
2266 };
2267 let mut outputs = vec![Vec::new()];
2268 dispatch_with_grid_sync_split_into(
2269 &backend,
2270 &program,
2271 &[[0u8, 0, 0, 0].as_slice()],
2272 &config,
2273 &mut outputs,
2274 )
2275 .expect("multi-pass split dispatch");
2276 assert_eq!(
2277 backend.calls.load(Ordering::SeqCst),
2278 6,
2279 "split must re-run the whole 2-segment sequence 3 times"
2280 );
2281 assert_eq!(
2282 outputs,
2283 vec![vec![6, 0, 0, 0]],
2284 "accumulator must advance one hop per fixpoint pass (2 segments × 3 passes)"
2285 );
2286 }
2287
2288 #[test]
2291 fn split_seeds_first_atomic_output_without_caller_bytes() {
2292 let program = Program::wrapped(
2293 vec![BufferDecl::output("out", 0, DataType::U32).with_count(1)],
2294 [1, 1, 1],
2295 vec![
2296 Node::let_bind(
2297 "prior",
2298 Expr::atomic_add("out", Expr::u32(0), Expr::u32(1)),
2299 ),
2300 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2301 Node::Return,
2302 ],
2303 );
2304 let dispatch = |segment: &Program,
2305 inputs: &[&[u8]],
2306 _grid: Option<[u32; 3]>,
2307 outputs: &mut Vec<Vec<u8>>|
2308 -> Result<(), String> {
2309 outputs.clear();
2310 if !segment_output_names(segment)
2311 .map_err(|error| error.to_string())?
2312 .is_empty()
2313 {
2314 assert_eq!(inputs, &[&[0, 0, 0, 0][..]]);
2315 outputs.push(1_u32.to_le_bytes().to_vec());
2316 }
2317 Ok(())
2318 };
2319
2320 let outputs = dispatch_with_grid_sync_split_via(
2321 &program,
2322 &[],
2323 &DispatchConfig::default(),
2324 &dispatch,
2325 )
2326 .expect("backend-allocated atomic output must receive its zero seed");
2327
2328 assert_eq!(outputs, vec![1_u32.to_le_bytes().to_vec()]);
2329 }
2330
2331 #[test]
2332 fn split_via_closure_entry_matches_backend_entry_on_the_same_grid_sync_program() {
2333 let program = Program::wrapped(
2340 vec![buffer()],
2341 [1, 1, 1],
2342 vec![
2343 region("a", vec![Node::Return]),
2344 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2345 region("b", vec![Node::Return]),
2346 ],
2347 );
2348 let config = DispatchConfig {
2349 fixpoint_iterations: Some(3),
2350 ..DispatchConfig::default()
2351 };
2352 let inputs: [&[u8]; 1] = [[0u8, 0, 0, 0].as_slice()];
2353
2354 let backend = IncrementingBackend {
2355 calls: AtomicUsize::new(0),
2356 };
2357 let mut backend_outputs = vec![Vec::new()];
2358 dispatch_with_grid_sync_split_into(
2359 &backend,
2360 &program,
2361 &inputs,
2362 &config,
2363 &mut backend_outputs,
2364 )
2365 .expect("backend split dispatch");
2366
2367 let closure_backend = IncrementingBackend {
2371 calls: AtomicUsize::new(0),
2372 };
2373 let dispatch = |program: &Program,
2374 inputs: &[&[u8]],
2375 grid: Option<[u32; 3]>,
2376 outputs: &mut Vec<Vec<u8>>|
2377 -> Result<(), String> {
2378 let segment_config = DispatchConfig {
2379 grid_override: grid,
2380 fixpoint_iterations: Some(1),
2381 ..DispatchConfig::default()
2382 };
2383 closure_backend
2384 .dispatch_borrowed_into(program, inputs, &segment_config, outputs)
2385 .map_err(|error| error.to_string())
2386 };
2387 let via_outputs = dispatch_with_grid_sync_split_via(&program, &inputs, &config, &dispatch)
2388 .expect("closure split dispatch");
2389
2390 assert_eq!(
2391 via_outputs, backend_outputs,
2392 "closure and backend split entries must produce identical output"
2393 );
2394 assert_eq!(
2395 closure_backend.calls.load(Ordering::SeqCst),
2396 backend.calls.load(Ordering::SeqCst),
2397 "both entries must drive the same number of segment dispatches (3 passes x 2 segments)"
2398 );
2399 assert_eq!(via_outputs, vec![vec![6u8, 0, 0, 0]]);
2400 }
2401
2402 struct OwnedFinalReserveBackend {
2403 calls: AtomicUsize,
2404 }
2405
2406 impl crate::backend::private::Sealed for OwnedFinalReserveBackend {}
2407
2408 impl VyreBackend for OwnedFinalReserveBackend {
2409 fn id(&self) -> &'static str {
2410 "grid-sync-owned-final-reserve"
2411 }
2412
2413 fn dispatch(
2414 &self,
2415 _program: &Program,
2416 _inputs: &[Vec<u8>],
2417 _config: &DispatchConfig,
2418 ) -> Result<Vec<Vec<u8>>, BackendError> {
2419 unreachable!("test uses dispatch_borrowed_into")
2420 }
2421
2422 fn dispatch_borrowed_into(
2423 &self,
2424 _program: &Program,
2425 inputs: &[&[u8]],
2426 _config: &DispatchConfig,
2427 outputs: &mut OutputBuffers,
2428 ) -> Result<(), BackendError> {
2429 let call = self.calls.fetch_add(1, Ordering::SeqCst);
2430 if call == 1 {
2431 assert!(
2432 outputs.capacity() >= 1,
2433 "owned grid-sync split wrapper must pre-reserve final output slots before the final segment dispatch"
2434 );
2435 }
2436 if outputs.is_empty() {
2437 outputs.push(Vec::new());
2438 }
2439 outputs[0].clear();
2440 outputs[0].extend_from_slice(inputs[0]);
2441 outputs[0][0] = outputs[0][0].saturating_add(1);
2442 Ok(())
2443 }
2444 }
2445
2446 #[test]
2447 fn split_owned_wrapper_reserves_final_output_vector_before_final_segment() {
2448 let program = Program::wrapped(
2449 vec![buffer()],
2450 [1, 1, 1],
2451 vec![
2452 region("a", vec![Node::Return]),
2453 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2454 region("b", vec![Node::Return]),
2455 ],
2456 );
2457 let backend = OwnedFinalReserveBackend {
2458 calls: AtomicUsize::new(0),
2459 };
2460 let input = [4u8, 0, 0, 0];
2461
2462 let outputs = dispatch_with_grid_sync_split(
2463 &backend,
2464 &program,
2465 &[input.as_slice()],
2466 &DispatchConfig::default(),
2467 )
2468 .expect("Fix: owned grid-sync split should reserve and return final outputs");
2469
2470 assert_eq!(backend.calls.load(Ordering::SeqCst), 2);
2471 assert_eq!(outputs, vec![vec![6, 0, 0, 0]]);
2472 }
2473
2474 #[test]
2475 fn grid_sync_split_records_segment_telemetry() {
2476 let program = Program::wrapped(
2477 vec![buffer()],
2478 [1, 1, 1],
2479 vec![
2480 region("a", vec![Node::Return]),
2481 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2482 region("b", vec![Node::Return]),
2483 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2484 region("c", vec![Node::Return]),
2485 ],
2486 );
2487 let backend = ReuseCheckingBackend {
2488 calls: AtomicUsize::new(0),
2489 final_outputs_addr: 0,
2490 final_slot_addr: 0,
2491 };
2492 let before = crate::observability::snapshot_dispatch_telemetry();
2493 let input = [0u8, 0, 0, 0];
2494 let mut outputs = Vec::new();
2495
2496 dispatch_with_grid_sync_split_into(
2497 &backend,
2498 &program,
2499 &[input.as_slice()],
2500 &DispatchConfig::default(),
2501 &mut outputs,
2502 )
2503 .expect("Fix: grid-sync split should dispatch every segment");
2504
2505 let after = crate::observability::snapshot_dispatch_telemetry();
2506 assert_eq!(backend.calls.load(Ordering::SeqCst), 3);
2507 assert!(after.grid_sync_splits >= before.grid_sync_splits + 1);
2508 assert!(after.grid_sync_segments >= before.grid_sync_segments + 3);
2509 assert!(after.grid_sync_points >= before.grid_sync_points + 2);
2510 }
2511
2512 struct IntermediateReuseBackend {
2513 calls: AtomicUsize,
2514 first_outputs_addr: AtomicUsize,
2515 first_slot_addr: AtomicUsize,
2516 }
2517
2518 impl crate::backend::private::Sealed for IntermediateReuseBackend {}
2519
2520 impl VyreBackend for IntermediateReuseBackend {
2521 fn id(&self) -> &'static str {
2522 "grid-sync-intermediate-reuse"
2523 }
2524
2525 fn dispatch(
2526 &self,
2527 _program: &Program,
2528 _inputs: &[Vec<u8>],
2529 _config: &DispatchConfig,
2530 ) -> Result<Vec<Vec<u8>>, BackendError> {
2531 unreachable!("test uses dispatch_borrowed_into")
2532 }
2533
2534 fn dispatch_borrowed_into(
2535 &self,
2536 _program: &Program,
2537 inputs: &[&[u8]],
2538 _config: &DispatchConfig,
2539 outputs: &mut OutputBuffers,
2540 ) -> Result<(), BackendError> {
2541 let call = self.calls.fetch_add(1, Ordering::SeqCst);
2542 if outputs.is_empty() {
2543 outputs.push(Vec::with_capacity(8));
2544 }
2545 if call == 0 {
2546 self.first_outputs_addr
2547 .store(outputs.as_ptr() as usize, Ordering::SeqCst);
2548 self.first_slot_addr
2549 .store(outputs[0].as_ptr() as usize, Ordering::SeqCst);
2550 } else if call == 1 {
2551 assert_eq!(
2552 outputs.as_ptr() as usize,
2553 self.first_outputs_addr.load(Ordering::SeqCst)
2554 );
2555 assert_eq!(
2556 outputs[0].as_ptr() as usize,
2557 self.first_slot_addr.load(Ordering::SeqCst)
2558 );
2559 }
2560 outputs[0].clear();
2561 outputs[0].extend_from_slice(inputs[0]);
2562 outputs[0][0] = outputs[0][0].saturating_add(1);
2563 Ok(())
2564 }
2565 }
2566
2567 #[test]
2568 fn split_reuses_intermediate_output_slot_between_segments() {
2569 let program = Program::wrapped(
2570 vec![buffer()],
2571 [1, 1, 1],
2572 vec![
2573 region("a", vec![Node::Return]),
2574 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2575 region("b", vec![Node::Return]),
2576 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2577 region("c", vec![Node::Return]),
2578 ],
2579 );
2580 let backend = IntermediateReuseBackend {
2581 calls: AtomicUsize::new(0),
2582 first_outputs_addr: AtomicUsize::new(0),
2583 first_slot_addr: AtomicUsize::new(0),
2584 };
2585 let input = [1u8, 0, 0, 0];
2586 let mut outputs = vec![Vec::with_capacity(8)];
2587
2588 dispatch_with_grid_sync_split_into(
2589 &backend,
2590 &program,
2591 &[input.as_slice()],
2592 &DispatchConfig::default(),
2593 &mut outputs,
2594 )
2595 .expect("Fix: grid-sync split should reuse intermediate output scratch");
2596
2597 assert_eq!(backend.calls.load(Ordering::SeqCst), 3);
2598 assert_eq!(outputs, vec![vec![4, 0, 0, 0]]);
2599 }
2600
2601 #[test]
2602 fn split_keeps_multi_segment_output_as_readwrite_accumulator() {
2603 let out = BufferDecl::output("out", 0, DataType::U32).with_count(4);
2612 let program = Program::wrapped(
2613 vec![out],
2614 [1, 1, 1],
2615 vec![
2616 region("a", vec![Node::store("out", Expr::u32(0), Expr::u32(0xAA))]),
2617 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2618 region("b", vec![Node::store("out", Expr::u32(2), Expr::u32(0xBB))]),
2619 ],
2620 );
2621 let segments =
2622 plan_host_grid_sync_segment_programs(&program).expect("plan host grid-sync segments");
2623 assert_eq!(segments.len(), 2, "one GridSync barrier -> two segments");
2624
2625 let seg0_out = segments[0]
2626 .buffers()
2627 .iter()
2628 .find(|b| b.name() == "out")
2629 .expect("segment 0 must declare the output it writes");
2630 assert_eq!(
2631 seg0_out.access(),
2632 BufferAccess::WriteOnly,
2633 "the first writer establishes the accumulator as write-only"
2634 );
2635 assert!(
2636 !seg0_out.is_output() && !seg0_out.is_pipeline_live_out(),
2637 "split segment buffers must never be marked program-output; final values are reassembled by name"
2638 );
2639
2640 let seg1_out = segments[1]
2641 .buffers()
2642 .iter()
2643 .find(|b| b.name() == "out")
2644 .expect("segment 1 must declare the output it writes");
2645 assert_eq!(
2646 seg1_out.access(),
2647 BufferAccess::ReadWrite,
2648 "a later writer of a multi-segment output must read+merge the accumulated value, not overwrite it"
2649 );
2650 assert!(
2651 !seg1_out.is_output() && !seg1_out.is_pipeline_live_out(),
2652 "the later writer must consume its forwarded prior value, which `segment_buffer_consumes_input` refuses for is_output buffers"
2653 );
2654 assert!(
2655 segment_input_names(&segments[1])
2656 .expect("segment 1 input names")
2657 .iter()
2658 .any(|n| n.as_str() == "out"),
2659 "the accumulated output must be forwarded as an input to the later writing segment"
2660 );
2661 }
2662
2663 struct SlotStoringBackend {
2669 calls: AtomicUsize,
2670 }
2671
2672 impl crate::backend::private::Sealed for SlotStoringBackend {}
2673
2674 impl VyreBackend for SlotStoringBackend {
2675 fn id(&self) -> &'static str {
2676 "grid-sync-slot-storing"
2677 }
2678
2679 fn dispatch(
2680 &self,
2681 _program: &Program,
2682 _inputs: &[Vec<u8>],
2683 _config: &DispatchConfig,
2684 ) -> Result<Vec<Vec<u8>>, BackendError> {
2685 unreachable!("test uses dispatch_borrowed_into")
2686 }
2687
2688 fn dispatch_borrowed_into(
2689 &self,
2690 program: &Program,
2691 inputs: &[&[u8]],
2692 _config: &DispatchConfig,
2693 outputs: &mut OutputBuffers,
2694 ) -> Result<(), BackendError> {
2695 let mut in_pos = None;
2698 let mut cur_in = 0usize;
2699 let mut out_pos = None;
2700 let mut cur_out = 0usize;
2701 for buffer in program.buffers() {
2702 if matches!(buffer.access(), BufferAccess::Workgroup) {
2703 continue;
2704 }
2705 let consumes = segment_buffer_consumes_input(buffer);
2706 let produces = segment_buffer_produces_output(buffer);
2707 if buffer.name() == "out" {
2708 if consumes {
2709 in_pos = Some(cur_in);
2710 }
2711 if produces {
2712 out_pos = Some(cur_out);
2713 }
2714 }
2715 if consumes {
2716 cur_in += 1;
2717 }
2718 if produces {
2719 cur_out += 1;
2720 }
2721 }
2722 let out_pos = out_pos.expect("every writing segment must produce `out`");
2723 let mut state = match in_pos {
2724 Some(i) => inputs[i].to_vec(),
2725 None => vec![0u8; 16],
2726 };
2727
2728 fn apply(nodes: &[Node], state: &mut [u8]) {
2729 for node in nodes {
2730 match node {
2731 Node::Store {
2732 buffer,
2733 index: Expr::LitU32(i),
2734 value: Expr::LitU32(v),
2735 } if buffer.as_str() == "out" => {
2736 let off = (*i as usize) * 4;
2737 state[off] = (*v & 0xff) as u8;
2738 }
2739 Node::Region { body, .. } => apply(body, state),
2740 Node::Block(body) => apply(body, state),
2741 Node::If {
2742 then, otherwise, ..
2743 } => {
2744 apply(then, state);
2745 apply(otherwise, state);
2746 }
2747 Node::Loop { body, .. } => apply(body, state),
2748 _ => {}
2749 }
2750 }
2751 }
2752 apply(entry_sequence(program), &mut state);
2753
2754 self.calls.fetch_add(1, Ordering::SeqCst);
2755 while outputs.len() <= out_pos {
2756 outputs.push(Vec::new());
2757 }
2758 outputs[out_pos].clear();
2759 outputs[out_pos].extend_from_slice(&state);
2760 Ok(())
2761 }
2762 }
2763
2764 #[test]
2765 fn split_preserves_earlier_segment_output_slots_end_to_end() {
2766 let out = BufferDecl::output("out", 0, DataType::U32).with_count(4);
2773 let program = Program::wrapped(
2774 vec![out],
2775 [1, 1, 1],
2776 vec![
2777 region("a", vec![Node::store("out", Expr::u32(0), Expr::u32(0xAA))]),
2778 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2779 region("b", vec![Node::store("out", Expr::u32(2), Expr::u32(0xBB))]),
2780 ],
2781 );
2782 let backend = SlotStoringBackend {
2783 calls: AtomicUsize::new(0),
2784 };
2785 let mut outputs = vec![Vec::new()];
2786 dispatch_with_grid_sync_split_into(
2787 &backend,
2788 &program,
2789 &[],
2790 &DispatchConfig::default(),
2791 &mut outputs,
2792 )
2793 .expect("split dispatch");
2794 assert_eq!(
2795 backend.calls.load(Ordering::SeqCst),
2796 2,
2797 "two segments, single fixpoint pass"
2798 );
2799 assert_eq!(outputs.len(), 1);
2800 assert_eq!(outputs[0].len(), 16, "output buffer is 4 × u32 = 16 bytes");
2801 assert_eq!(
2802 outputs[0][0], 0xAA,
2803 "segment 0's slot (element 0) must survive the final segment's write"
2804 );
2805 assert_eq!(
2806 outputs[0][8], 0xBB,
2807 "the final segment's slot (element 2) is also present"
2808 );
2809 }
2810
2811 struct SaturatingBackend {
2815 calls: AtomicUsize,
2816 cap: u8,
2817 }
2818
2819 impl crate::backend::private::Sealed for SaturatingBackend {}
2820
2821 impl VyreBackend for SaturatingBackend {
2822 fn id(&self) -> &'static str {
2823 "grid-sync-saturating"
2824 }
2825
2826 fn dispatch(
2827 &self,
2828 _program: &Program,
2829 _inputs: &[Vec<u8>],
2830 _config: &DispatchConfig,
2831 ) -> Result<Vec<Vec<u8>>, BackendError> {
2832 unreachable!("test uses dispatch_borrowed_into")
2833 }
2834
2835 fn dispatch_borrowed_into(
2836 &self,
2837 _program: &Program,
2838 inputs: &[&[u8]],
2839 _config: &DispatchConfig,
2840 outputs: &mut OutputBuffers,
2841 ) -> Result<(), BackendError> {
2842 self.calls.fetch_add(1, Ordering::SeqCst);
2843 if outputs.is_empty() {
2844 outputs.push(Vec::new());
2845 }
2846 outputs[0].clear();
2847 outputs[0].extend_from_slice(inputs[0]);
2848 if outputs[0][0] < self.cap {
2849 outputs[0][0] += 1;
2850 }
2851 Ok(())
2852 }
2853 }
2854
2855 #[test]
2856 fn split_outer_loop_early_exits_when_accumulator_reaches_fixpoint() {
2857 let program = Program::wrapped(
2862 vec![buffer()],
2863 [1, 1, 1],
2864 vec![
2865 region("a", vec![Node::Return]),
2866 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2867 region("b", vec![Node::Return]),
2868 ],
2869 );
2870 let backend = SaturatingBackend {
2871 calls: AtomicUsize::new(0),
2872 cap: 3,
2873 };
2874 let config = DispatchConfig {
2875 fixpoint_iterations: Some(10),
2876 ..DispatchConfig::default()
2877 };
2878 let mut outputs = vec![Vec::new()];
2879 dispatch_with_grid_sync_split_into(
2880 &backend,
2881 &program,
2882 &[[0u8, 0, 0, 0].as_slice()],
2883 &config,
2884 &mut outputs,
2885 )
2886 .expect("converging split dispatch");
2887 assert_eq!(
2890 backend.calls.load(Ordering::SeqCst),
2891 6,
2892 "outer loop must early-exit one pass after the accumulator stops changing, not run all 10 iterations"
2893 );
2894 assert_eq!(
2895 outputs,
2896 vec![vec![3, 0, 0, 0]],
2897 "early-exit must return the converged fixpoint value, identical to running every iteration"
2898 );
2899 }
2900
2901 #[test]
2902 fn split_non_converging_accumulator_runs_full_iteration_budget() {
2903 let program = Program::wrapped(
2907 vec![buffer()],
2908 [1, 1, 1],
2909 vec![
2910 region("a", vec![Node::Return]),
2911 Node::barrier_with_ordering(MemoryOrdering::GridSync),
2912 region("b", vec![Node::Return]),
2913 ],
2914 );
2915 let backend = SaturatingBackend {
2917 calls: AtomicUsize::new(0),
2918 cap: 255,
2919 };
2920 let config = DispatchConfig {
2921 fixpoint_iterations: Some(4),
2922 ..DispatchConfig::default()
2923 };
2924 let mut outputs = vec![Vec::new()];
2925 dispatch_with_grid_sync_split_into(
2926 &backend,
2927 &program,
2928 &[[0u8, 0, 0, 0].as_slice()],
2929 &config,
2930 &mut outputs,
2931 )
2932 .expect("non-converging split dispatch");
2933 assert_eq!(
2934 backend.calls.load(Ordering::SeqCst),
2935 8,
2936 "a still-advancing accumulator must run the full 4 iterations x 2 segments"
2937 );
2938 assert_eq!(outputs, vec![vec![8, 0, 0, 0]]);
2939 }
2940
2941 struct ResidentReuseBackend {
2942 calls: AtomicUsize,
2943 owner: crate::ResidentOwner,
2944 }
2945
2946 impl crate::backend::private::Sealed for ResidentReuseBackend {}
2947
2948 impl VyreBackend for ResidentReuseBackend {
2949 fn id(&self) -> &'static str {
2950 "grid-sync-resident-reuse"
2951 }
2952
2953 fn dispatch(
2954 &self,
2955 _program: &Program,
2956 _inputs: &[Vec<u8>],
2957 _config: &DispatchConfig,
2958 ) -> Result<Vec<Vec<u8>>, BackendError> {
2959 unreachable!("test uses dispatch_resident_timed")
2960 }
2961
2962 fn dispatch_borrowed_into(
2963 &self,
2964 _program: &Program,
2965 _inputs: &[&[u8]],
2966 _config: &DispatchConfig,
2967 _outputs: &mut OutputBuffers,
2968 ) -> Result<(), BackendError> {
2969 unreachable!("resident grid-sync split must not refresh through host borrowed inputs")
2970 }
2971
2972 fn dispatch_resident_timed(
2973 &self,
2974 _program: &Program,
2975 resources: &[Resource],
2976 _config: &DispatchConfig,
2977 ) -> Result<TimedDispatchResult, BackendError> {
2978 let bound: Vec<u64> = resources
2979 .iter()
2980 .map(|resource| match resource {
2981 Resource::Resident(handle) => self
2982 .owner
2983 .resolve(*handle, "grid-sync resident reuse test backend")
2984 .expect("Fix: test backend must only receive its own resident handles"),
2985 Resource::Borrowed(_) => panic!(
2986 "Fix: resident grid-sync split must bind device handles, not host bytes."
2987 ),
2988 })
2989 .collect();
2990 assert_eq!(
2991 bound,
2992 vec![11, 22],
2993 "Fix: resident grid-sync split must keep the original device handles bound across every segment."
2994 );
2995 let call = self.calls.fetch_add(1, Ordering::SeqCst);
2996 Ok(TimedDispatchResult {
2997 outputs: vec![vec![call as u8]],
2998 wall_ns: 10,
2999 device_ns: Some(2),
3000 enqueue_ns: Some(3),
3001 wait_ns: Some(4),
3002 })
3003 }
3004 }
3005
3006 #[test]
3007 fn resident_split_reuses_same_device_resources_across_segments() {
3008 let program = Program::wrapped(
3009 vec![buffer()],
3010 [1, 1, 1],
3011 vec![
3012 region("a", vec![Node::Return]),
3013 Node::barrier_with_ordering(MemoryOrdering::GridSync),
3014 region("b", vec![Node::Return]),
3015 Node::barrier_with_ordering(MemoryOrdering::GridSync),
3016 region("c", vec![Node::Return]),
3017 ],
3018 );
3019 let owner = crate::ResidentOwner::new().expect("Fix: owner ids must be available");
3020 let backend = ResidentReuseBackend {
3021 calls: AtomicUsize::new(0),
3022 owner,
3023 };
3024
3025 let timed = dispatch_resident_with_grid_sync_split_timed(
3026 &backend,
3027 &program,
3028 &[
3029 Resource::Resident(owner.handle(11)),
3030 Resource::Resident(owner.handle(22)),
3031 ],
3032 &DispatchConfig::default(),
3033 )
3034 .expect("Fix: resident grid-sync split should run each segment on the same device handles");
3035
3036 assert_eq!(backend.calls.load(Ordering::SeqCst), 3);
3037 assert_eq!(timed.outputs, vec![vec![2]]);
3038 assert_eq!(timed.device_ns, Some(6));
3039 assert_eq!(timed.enqueue_ns, Some(9));
3040 assert_eq!(timed.wait_ns, Some(12));
3041 }
3042
3043 struct ResidentDeviceBackend {
3049 owner: crate::ResidentOwner,
3050 next_id: std::sync::atomic::AtomicU64,
3051 buffers: std::sync::Mutex<HashMap<u64, Vec<u8>>>,
3052 freed: std::sync::Mutex<Vec<u64>>,
3053 dispatches: AtomicUsize,
3054 }
3055
3056 impl ResidentDeviceBackend {
3057 fn new() -> Self {
3058 Self {
3059 owner: crate::ResidentOwner::new().expect("Fix: owner ids must be available"),
3060 next_id: std::sync::atomic::AtomicU64::new(1),
3061 buffers: std::sync::Mutex::new(HashMap::new()),
3062 freed: std::sync::Mutex::new(Vec::new()),
3063 dispatches: AtomicUsize::new(0),
3064 }
3065 }
3066
3067 fn resident_id(&self, resource: &Resource) -> u64 {
3068 match resource {
3069 Resource::Resident(handle) => self
3070 .owner
3071 .resolve(*handle, "grid-sync resident fixpoint test backend")
3072 .expect("Fix: test backend must only receive its own resident handles"),
3073 Resource::Borrowed(_) => {
3074 panic!(
3075 "Fix: resident grid-sync fixpoint must bind Resident handles, not Borrowed"
3076 )
3077 }
3078 }
3079 }
3080 }
3081
3082 impl crate::backend::private::Sealed for ResidentDeviceBackend {}
3083
3084 impl VyreBackend for ResidentDeviceBackend {
3085 fn id(&self) -> &'static str {
3086 "grid-sync-resident-device"
3087 }
3088
3089 fn dispatch(
3090 &self,
3091 _program: &Program,
3092 _inputs: &[Vec<u8>],
3093 _config: &DispatchConfig,
3094 ) -> Result<Vec<Vec<u8>>, BackendError> {
3095 unreachable!("resident fixpoint test uses resident dispatch")
3096 }
3097
3098 fn dispatch_borrowed_into(
3099 &self,
3100 _program: &Program,
3101 _inputs: &[&[u8]],
3102 _config: &DispatchConfig,
3103 _outputs: &mut OutputBuffers,
3104 ) -> Result<(), BackendError> {
3105 unreachable!("resident fixpoint must thread device handles, never host borrowed inputs")
3106 }
3107
3108 fn allocate_resident(&self, byte_len: usize) -> Result<Resource, BackendError> {
3109 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
3110 self.buffers
3113 .lock()
3114 .unwrap()
3115 .insert(id, vec![0xFFu8; byte_len]);
3116 Ok(Resource::Resident(self.owner.handle(id)))
3117 }
3118
3119 fn upload_resident(&self, resource: &Resource, bytes: &[u8]) -> Result<(), BackendError> {
3120 let id = self.resident_id(resource);
3121 let mut buffers = self.buffers.lock().unwrap();
3122 let buf = buffers.get_mut(&id).expect("resident handle exists");
3123 assert!(
3124 bytes.len() <= buf.len(),
3125 "upload {} bytes into a {}-byte resident buffer",
3126 bytes.len(),
3127 buf.len()
3128 );
3129 buf[..bytes.len()].copy_from_slice(bytes);
3130 Ok(())
3131 }
3132
3133 fn download_resident_range_into(
3134 &self,
3135 resource: &Resource,
3136 byte_offset: usize,
3137 byte_len: usize,
3138 output: &mut Vec<u8>,
3139 ) -> Result<(), BackendError> {
3140 let id = self.resident_id(resource);
3141 let buffers = self.buffers.lock().unwrap();
3142 let buf = buffers.get(&id).expect("resident handle exists");
3143 output.clear();
3144 output.extend_from_slice(&buf[byte_offset..byte_offset + byte_len]);
3145 Ok(())
3146 }
3147
3148 fn free_resident(&self, resource: Resource) -> Result<(), BackendError> {
3149 let id = self.resident_id(&resource);
3150 self.buffers.lock().unwrap().remove(&id);
3151 self.freed.lock().unwrap().push(id);
3152 Ok(())
3153 }
3154
3155 fn dispatch_resident_timed(
3156 &self,
3157 program: &Program,
3158 resources: &[Resource],
3159 _config: &DispatchConfig,
3160 ) -> Result<TimedDispatchResult, BackendError> {
3161 self.dispatches.fetch_add(1, Ordering::SeqCst);
3162 let plan = BindingPlan::build(program)?;
3165 let mut out_slot = None;
3166 let mut pos = 0usize;
3167 for binding in &plan.bindings {
3168 if binding.role == BindingRole::Shared {
3169 continue;
3170 }
3171 if binding.name.as_ref() == "out" {
3172 out_slot = Some(pos);
3173 }
3174 pos += 1;
3175 }
3176 let out_slot = out_slot.expect("program declares `out`");
3177 let id = self.resident_id(&resources[out_slot]);
3178 let mut buffers = self.buffers.lock().unwrap();
3179 let buf = buffers.get_mut(&id).expect("resident `out` handle exists");
3180
3181 fn apply(nodes: &[Node], state: &mut [u8]) {
3184 for node in nodes {
3185 match node {
3186 Node::Store {
3187 buffer,
3188 index: Expr::LitU32(i),
3189 value: Expr::LitU32(v),
3190 } if buffer.as_str() == "out" => {
3191 state[(*i as usize) * 4] = (*v & 0xff) as u8;
3192 }
3193 Node::Region { body, .. } => apply(body, state),
3194 Node::Block(body) => apply(body, state),
3195 Node::If {
3196 then, otherwise, ..
3197 } => {
3198 apply(then, state);
3199 apply(otherwise, state);
3200 }
3201 Node::Loop { body, .. } => apply(body, state),
3202 _ => {}
3203 }
3204 }
3205 }
3206 apply(entry_sequence(program), buf.as_mut_slice());
3207
3208 Ok(TimedDispatchResult {
3209 outputs: Vec::new(),
3210 wall_ns: 1,
3211 device_ns: Some(1),
3212 enqueue_ns: Some(1),
3213 wait_ns: Some(1),
3214 })
3215 }
3216 }
3217
3218 #[test]
3219 fn resident_fixpoint_accumulates_across_segments_zero_inits_and_frees() {
3220 let out = BufferDecl::output("out", 0, DataType::U32).with_count(4);
3226 let program = Program::wrapped(
3227 vec![out],
3228 [1, 1, 1],
3229 vec![
3230 region("a", vec![Node::store("out", Expr::u32(0), Expr::u32(0xAA))]),
3231 Node::barrier_with_ordering(MemoryOrdering::GridSync),
3232 region("b", vec![Node::store("out", Expr::u32(2), Expr::u32(0xBB))]),
3233 ],
3234 );
3235 let backend = ResidentDeviceBackend::new();
3236 let mut outputs = vec![Vec::new()];
3237 dispatch_resident_grid_sync_fixpoint_into(
3238 &backend,
3239 &program,
3240 &[],
3241 &DispatchConfig::default(),
3242 &mut outputs,
3243 )
3244 .expect("resident grid-sync fixpoint dispatch");
3245
3246 assert_eq!(
3247 backend.dispatches.load(Ordering::SeqCst),
3248 2,
3249 "two segments, single fixpoint pass under the default config"
3250 );
3251 assert_eq!(outputs.len(), 1, "one output buffer (`out`)");
3252 assert_eq!(outputs[0].len(), 16, "4 × u32 = 16 bytes");
3253 assert_eq!(
3254 outputs[0][0], 0xAA,
3255 "segment 0's slot survives - resident accumulation, no clobber"
3256 );
3257 assert_eq!(outputs[0][8], 0xBB, "the final segment's slot is present");
3258 assert_eq!(outputs[0][4], 0x00, "untouched slot 1 was zero-initialized");
3262 assert_eq!(
3263 outputs[0][12], 0x00,
3264 "untouched slot 3 was zero-initialized"
3265 );
3266 assert_eq!(
3268 backend.freed.lock().unwrap().len(),
3269 1,
3270 "the single `out` resident buffer is freed"
3271 );
3272 assert!(
3273 backend.buffers.lock().unwrap().is_empty(),
3274 "no resident buffer leaks after dispatch"
3275 );
3276 }
3277
3278 #[test]
3279 fn resident_fixpoint_repeats_to_fixpoint_bound() {
3280 let out = BufferDecl::output("out", 0, DataType::U32).with_count(4);
3284 let program = Program::wrapped(
3285 vec![out],
3286 [1, 1, 1],
3287 vec![
3288 region("a", vec![Node::store("out", Expr::u32(0), Expr::u32(0xAA))]),
3289 Node::barrier_with_ordering(MemoryOrdering::GridSync),
3290 region("b", vec![Node::store("out", Expr::u32(2), Expr::u32(0xBB))]),
3291 ],
3292 );
3293 let backend = ResidentDeviceBackend::new();
3294 let mut config = DispatchConfig::default();
3295 config.fixpoint_iterations = Some(3);
3296 let mut outputs = vec![Vec::new()];
3297 dispatch_resident_grid_sync_fixpoint_into(&backend, &program, &[], &config, &mut outputs)
3298 .expect("resident grid-sync fixpoint dispatch");
3299 assert_eq!(
3300 backend.dispatches.load(Ordering::SeqCst),
3301 6,
3302 "2 segments × 3 fixpoint passes"
3303 );
3304 assert_eq!(outputs[0][0], 0xAA);
3305 assert_eq!(outputs[0][8], 0xBB);
3306 }
3307}