Skip to main content

vyre_driver/
grid_sync.rs

1//! Grid-sync kernel splitting.
2//!
3//! Op id: `vyre-driver::grid_sync`. Soundness: `Exact` over the
4//! cross-grid barrier contract.
5//!
6//! ## Why this lives in vyre-driver, not the backend
7//!
8//! Every backend that lacks a native cooperative whole-grid launch
9//! needs the same kernel-split semantics for
10//! `Node::Barrier { ordering: GridSync }`: split the program at the
11//! barrier, dispatch each segment as its own kernel launch, and
12//! re-feed the prior segment's outputs as inputs to the next. The
13//! kernel-launch boundary itself is the grid-level fence  -  every
14//! prior write becomes globally visible before the next launch reads.
15//!
16//! Backends route through [`crate::grid_sync::dispatch_with_grid_sync_split`] when
17//! [`VyreBackend::supports_grid_sync`] is `false` and the program
18//! contains any `Node::Barrier { ordering: GridSync }`. Backends that
19//! return `true` emit one kernel and satisfy the barrier device-side.
20//!
21//! ## Algorithm
22//!
23//! 1. Walk the program's top-level entry sequence.
24//! 2. Each prefix-suffix split at a `Node::Barrier { GridSync }`
25//!    becomes one segment.
26//! 3. For each segment, build a `Program` with a segment-local buffer
27//!    table: buffers read or written by that segment plus passthrough
28//!    read-write buffers that must preserve caller-visible storage.
29//! 4. Dispatch segments in order, threading live buffers by buffer name
30//!    rather than positional output slot. Segment read-only inputs are
31//!    assembled from the caller's original bytes or prior segment
32//!    outputs; final host-visible output slots are reassembled in the
33//!    original program's output declaration order.
34//!
35//! ## Device-resident variant
36//!
37//! [`crate::grid_sync::dispatch_with_grid_sync_split_into`] round-trips every live buffer
38//! host↔device between each segment and on every fixpoint pass. For a fused
39//! multi-rule program whose shared output accumulator is hundreds of MiB and
40//! which splits into hundreds of segments, that transfer, not launch
41//! latency, dominates wall time. [`crate::grid_sync::dispatch_resident_grid_sync_fixpoint_into`]
42//! is the device-resident counterpart: it uploads inputs into backend-resident
43//! resources once, keeps them bound across every segment and fixpoint pass (so
44//! the accumulator threads in place on-device, since resident dispatch never
45//! clears a bound buffer between launches), and reads back only the final
46//! outputs. It requires [`VyreBackend::supports_resident_dispatch`]; callers
47//! route to it on resident-capable backends and to the host split otherwise.
48//! Both paths are recall- and proof-identical (proven by a host/resident
49//! differential gate); the choice is purely a host↔device-traffic optimization.
50//!
51//! ## Soundness
52//!
53//! - Atomicity preserved: every `atomic_or` that fired in segment N
54//!   has flushed to global memory by the time segment N+1 launches  -
55//!   backend launch APIs issue an implicit grid-level fence at
56//!   submission boundaries.
57//! - Ordering preserved: the original program's host-visible output
58//!   is byte-identical to the un-split version, modulo timing.
59//! - No re-validation surprise: each split segment validates against
60//!   the same backend supported-ops set as the original.
61
62use 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/// Walk past `Program::wrapped`'s synthetic outer Region. Real
76/// programs are constructed via `wrapped`, which inserts a single
77/// outer Region around the user's entry sequence; the split logic
78/// must operate on the inner sequence so a `GridSync` barrier inside
79/// the wrapper actually splits the program. Programs constructed
80/// via `Program::new` use the entry sequence directly  -  in that
81/// case we just return it unchanged.
82#[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/// Whether `program` contains any `Node::Barrier { ordering: GridSync }`
127/// in its dispatch-level entry sequence (peeled past any synthetic
128/// outer Region).
129///
130/// The check is intentionally shallow: nested grid-sync barriers
131/// inside `Node::Loop` or inner `Node::Region` bodies are a contract
132/// violation (`validate::barrier` rejects them) and never reach this
133/// path. The split operates at the dispatch-level granularity.
134#[must_use]
135pub fn contains_grid_sync(program: &Program) -> bool {
136    // O(1) negative gate: if the cached ProgramStats bitset records no
137    // Barrier of any kind in the entire tree, there is definitely no
138    // top-level GridSync barrier either. Skip the entry-sequence walk
139    // (which itself is shallow but still pays a buffers/buffer_index
140    // dispatch on every backend dispatch path).
141    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/// Split `program` at every top-level `Node::Barrier { GridSync }`.
167///
168/// Returns a vector of segments in execution order. The barrier nodes
169/// themselves are dropped from the segments  -  the kernel-launch
170/// boundary between segments takes their place.
171///
172/// Each returned segment is a complete `Program` that shares the
173/// original's buffer table, workgroup size, and metadata; only the
174/// entry sequence changes. Segments without any executable nodes are
175/// preserved (an empty segment between two adjacent barriers becomes
176/// a no-op kernel that completes with byte-identical inputs and
177/// outputs).
178#[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
186/// Fallible variant of [`split_on_grid_sync`] for production dispatch paths.
187///
188/// # Errors
189/// Returns an actionable [`BackendError`] if segment storage cannot be
190/// reserved or if split accounting overflows.
191fn 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
506/// Fallible variant of [`split_on_grid_sync`] for production dispatch paths.
507///
508/// # Errors
509/// Returns an actionable [`BackendError`] if segment storage cannot be
510/// reserved or if split accounting overflows.
511
512pub 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    // Re-wrap each segment in the same wrapper stack the source had,
582    // so tagged/fused programs keep provenance and structure while the
583    // executable body is split at launch boundaries.
584    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
602/// Diagnostics: the host-split segment **programs** (post buffer-rewrite) that
603/// the host-split dispatch path (`dispatch_with_grid_sync_split*`) validates and
604/// launches when the backend lacks native grid-sync. Exposed so tooling and
605/// tests can inspect or validate each segment without a live backend, the
606/// raw [`try_split_on_grid_sync`] output omits the per-segment buffer
607/// access/role rewrite, so it is not what the backend actually sees.
608///
609/// # Errors
610/// Propagates any [`BackendError`] from splitting or buffer rewriting.
611pub 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
641/// For each buffer name, the index of the FIRST split segment that writes it.
642///
643/// A source-output buffer written by more than one segment is an
644/// **accumulator**: each segment writes only its own slots (e.g. a fused
645/// multi-rule `results_packed`, where every rule's result-store lands in a
646/// different grid-sync segment). A LATER writer must therefore read+merge the
647/// value forwarded from earlier segments via `current_inputs`, never overwrite
648/// it with a fresh WriteOnly buffer, which would silently zero every earlier
649/// segment's slots (recall=0 for every rule whose store is not in the final
650/// segment). `rewrite_segment_buffers_for_host_split` uses this map to keep an
651/// already-produced output buffer as a `ReadWrite` accumulator in later
652/// segments instead of a write-only output.
653fn 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        // A source-output buffer that an EARLIER segment already wrote is an
735        // accumulator across the split: this segment must read the value
736        // forwarded via `current_inputs` and merge its own slots, never
737        // overwrite it with a fresh WriteOnly buffer (which zeroes the earlier
738        // segments' slots, the silent recall=0 mode for any fused rule whose
739        // result-store does not land in the final segment).
740        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            // Later writer of a multi-segment output accumulator: read the
750            // accumulated prior value (uploaded as input) and merge this
751            // segment's slots in place.
752            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        // Never mark a split segment's buffer as the program output: a
763        // multi-segment output accumulator must CONSUME its forwarded prior
764        // value as input in later segments, and `segment_buffer_consumes_input`
765        // refuses any `is_output` buffer. Each writing segment still produces
766        // the buffer (WriteOnly/ReadWrite both produce output), so its bytes
767        // are captured into `current_inputs`; the final host-visible values are
768        // reassembled by name from the SOURCE program's output set in
769        // `collect_final_named_outputs`, independent of any per-segment flag.
770        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
930/// Universal dispatch helper that satisfies `Node::Barrier { ordering:
931/// GridSync }` on any backend by splitting at the barrier and running
932/// each segment as its own kernel launch.
933///
934/// Backends with native cooperative-launch grid sync (advertised via
935/// [`VyreBackend::supports_grid_sync`]) bypass the split  -  the
936/// program is dispatched once. Backends without it route here so the
937/// kernel-launch boundary becomes the grid-level fence: every prior
938/// write is globally visible to subsequent launches.
939///
940/// # Inputs
941/// `inputs` matches the input slice the caller would have passed to
942/// `dispatch_borrowed`. After each segment, the helper refreshes
943/// every ReadWrite buffer's slot from the segment's readback so the
944/// next segment sees the prior writes.
945///
946/// # Errors
947/// Propagates any `BackendError` raised by `dispatch_borrowed` on a
948/// segment, prefixed with the segment index for diagnosability.
949pub 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
965/// Timed variant of [`dispatch_with_grid_sync_split`].
966///
967/// # Errors
968/// Propagates any [`BackendError`] raised by a segment dispatch.
969pub 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
986/// Resident-resource variant of [`dispatch_with_grid_sync_split_timed`].
987///
988/// This keeps the same resource handles bound for every segment. Read-write
989/// buffers therefore refresh in place on the backend's device-resident storage
990/// between segment launches instead of downloading bytes to the host and
991/// re-uploading them as the next segment's inputs.
992///
993/// # Errors
994/// Propagates any [`BackendError`] raised by a segment resident dispatch.
995pub 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    // These are the explicit non-native grid-sync routes (host split /
1002    // resident fixpoint). They split unconditionally when the program carries a
1003    // grid-sync barrier: native cooperative launch has a residency ceiling, so
1004    // `supports_grid_sync()` no longer implies "this program runs natively".
1005    // The orchestrator (or the registry's `should_split_grid_sync`) decides
1006    // native-vs-split per program; once here, always split.
1007    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
1116/// Variant of [`dispatch_with_grid_sync_split`] that writes final outputs into
1117/// caller-owned storage.
1118///
1119/// # Errors
1120/// Propagates any `BackendError` raised by a segment dispatch.
1121fn 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    // These are the explicit non-native grid-sync routes (host split /
1132    // resident fixpoint). They split unconditionally when the program carries a
1133    // grid-sync barrier: native cooperative launch has a residency ceiling, so
1134    // `supports_grid_sync()` no longer implies "this program runs natively".
1135    // The orchestrator (or the registry's `should_split_grid_sync`) decides
1136    // native-vs-split per program; once here, always split.
1137    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    // Build a mutable input set we rotate between segments. ReadOnly
1151    // inputs stay borrowed from the caller for the whole split; only
1152    // ReadWrite buffers become owned after a segment produces updated
1153    // bytes. The previous implementation cloned every input before
1154    // the first launch, which turned large read-only buffers into a
1155    // host-memory copy on the slow path.
1156    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    // Honor the program's fixpoint contract across the split. The
1185    // non-split dispatch path (`dispatch_borrowed`) re-runs the WHOLE
1186    // program `fixpoint_iterations` times with persistent ReadWrite
1187    // buffers, so a program authored as a fixpoint closure converges
1188    // a multi-hop reachability/dataflow closure is exactly this shape: a
1189    // `seed (acc |= source) → hop (acc' = step(acc)) → merge (acc |= acc')`
1190    // body whose accumulator grows by ONE dataflow hop per whole-program
1191    // pass, relying on the dispatcher to iterate it to a fixpoint.
1192    //
1193    // GridSync barriers split that body across segments, so ONE pass over
1194    // the segment sequence advances the accumulator by exactly one hop.
1195    // Re-running an individual SEGMENT N times (the previous behavior:
1196    // `config` with its fixpoint count reached each segment) does NOT
1197    // converge, re-launching the isolated `hop` segment recomputes the
1198    // same frontier from an unchanged `acc`. The whole SEQUENCE must be
1199    // looped instead, with each segment run once per pass. Net device work
1200    // is identical (sequence_len × iterations launches either way); only
1201    // the nesting order changes, which is what makes the closure converge.
1202    // A flow that needs k hops through k-1 intermediate variables (the
1203    // dominant launch-rule shape: `q = src; sink(q)`) silently returned an
1204    // empty frontier under the old single-pass split (recall=0).
1205    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    // Adaptive convergence: `iterations` is an UPPER bound (the worst-case hop
1211    // depth, one hop per whole-sequence pass). The segment sequence is a
1212    // deterministic function of its live buffers, so once a full pass leaves
1213    // every evolving (Owned) accumulator unchanged the closure has reached a
1214    // fixpoint, every remaining pass would re-dispatch the entire segment
1215    // sequence (hundreds of launches on a large fused program) for zero new
1216    // dataflow. Stop as soon as two consecutive passes produce the same state.
1217    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, &current_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(&current_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
1241/// Split a grid-sync program at its barriers and dispatch every segment through
1242/// `backend`, looping the segment sequence to a fixpoint.
1243///
1244/// This is the `&dyn VyreBackend` entry; the split, refresh, and adaptive
1245/// convergence logic lives in [`dispatch_grid_sync_split_generic`], shared with
1246/// the closure entry [`dispatch_with_grid_sync_split_via_into`].
1247///
1248/// # Errors
1249/// Propagates any [`BackendError`] from splitting or a segment dispatch,
1250/// prefixed with the segment index.
1251pub 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
1263/// Closure-driven counterpart of [`dispatch_with_grid_sync_split_into`] for
1264/// callers that hold an opaque single-launch dispatch closure instead of a
1265/// `&dyn VyreBackend`.
1266///
1267/// This is the entry a host-loop fixpoint solver (an IFDS or dataflow solve)
1268/// uses to move its convergence loop onto the device without taking a backend
1269/// handle: it plugs any backend (CPU reference, CUDA, wgpu) as a
1270/// `Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(),
1271/// String>` closure. The closure receives each segment's program, its rotated
1272/// inputs, the whole-grid workgroup count (`config.grid_override`), and a
1273/// per-segment output slot to fill in the segment program's output order. The
1274/// split, refresh, and convergence logic is the SAME code as the backend entry
1275/// (both call [`dispatch_grid_sync_split_generic`]), so the two paths converge
1276/// to identical output.
1277///
1278/// # Errors
1279/// Propagates any error the closure returns (wrapped through
1280/// [`BackendError::new`]) and any structural split error, prefixed with the
1281/// segment index.
1282pub 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
1297/// Allocating wrapper over [`dispatch_with_grid_sync_split_via_into`].
1298///
1299/// # Errors
1300/// Propagates any error from [`dispatch_with_grid_sync_split_via_into`].
1301pub 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
1320/// Device-resident counterpart of [`dispatch_with_grid_sync_split_into`].
1321///
1322/// The host-split path round-trips every live buffer host↔device between each
1323/// split segment AND on every fixpoint pass. A fused multi-rule
1324/// `results_packed` accumulator is hundreds of MiB, so a program that splits
1325/// into hundreds of segments moves tens of GiB across PCIe per dispatch, that
1326/// transfer, not launch latency, is the host-split wall.
1327///
1328/// This variant uploads the program's inputs into backend-resident resources
1329/// ONCE, keeps them bound across every segment and every fixpoint pass, so a
1330/// multi-rule accumulator threads IN PLACE on device storage with no host copy
1331/// and no clobber (and reads back only the final output ranges a single time).
1332/// Net host↔device traffic drops from `O(segments × passes × live_bytes)` to
1333/// `O(inputs + outputs)`.
1334///
1335/// Every split segment from [`try_split_on_grid_sync`] carries the full program
1336/// buffer table (only the executable entry sequence differs), so one resident
1337/// resource slice binds to every segment. Resident dispatch never clears a
1338/// bound buffer between launches, so each rule's result-store accumulates into
1339/// the shared device `results_packed` exactly as the un-split program would.
1340///
1341/// `outputs` is shaped byte-identically to
1342/// [`dispatch_with_grid_sync_split_into`]: one `Vec<u8>` per original output
1343/// buffer, in declaration order, so a caller can swap paths without changing
1344/// readback.
1345///
1346/// Requires a backend implementing the resident half of the [`VyreBackend`]
1347/// contract (`allocate_resident` / `upload_resident` /
1348/// `dispatch_resident_repeated_sequence_read_ranges_into` / `free_resident`).
1349/// A backend without residency fails loudly with `UnsupportedFeature` at the
1350/// first resident call; callers route those to
1351/// [`dispatch_with_grid_sync_split_into`].
1352///
1353/// # Errors
1354/// Propagates any [`BackendError`] from splitting, resident allocation, upload,
1355/// segment dispatch, or readback. Resident resources allocated by this call are
1356/// always freed before returning, on success and on error.
1357pub 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    // These are the explicit non-native grid-sync routes (host split /
1365    // resident fixpoint). They split unconditionally when the program carries a
1366    // grid-sync barrier: native cooperative launch has a residency ceiling, so
1367    // `supports_grid_sync()` no longer implies "this program runs natively".
1368    // The orchestrator (or the registry's `should_split_grid_sync`) decides
1369    // native-vs-split per program; once here, always split.
1370    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    // Allocate one resident resource per non-shared binding (caller inputs
1385    // uploaded; output/scratch buffers zeroed so an accumulator's unfired
1386    // slots stay 0), then run the fixpoint and read back final outputs.
1387    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    // Free every resident resource before returning, success or error.
1391    let free_result = free_resident_program_resources(backend, resident);
1392    result.and(free_result)
1393}
1394
1395/// Resident resources backing one [`crate::grid_sync::dispatch_resident_grid_sync_fixpoint_into`]
1396/// call: the binding-ordered slice every segment dispatches against, plus a
1397/// name → (handle, byte-len) map for output readback.
1398struct ResidentProgramResources {
1399    /// One resource per non-shared binding, in [`BindingPlan`] order, the
1400    /// slice the backend's resident dispatch binds positionally.
1401    ordered: Vec<Resource>,
1402    /// Buffer-name → (resident handle clone, byte length) for output readback
1403    /// by name. The handle is a cheap id clone; freeing `ordered` frees it.
1404    by_name: HashMap<Ident, (Resource, usize)>,
1405}
1406
1407/// Allocate + initialize one resident resource per non-shared program binding.
1408///
1409/// Inputs are uploaded from the caller slice; output / write-only / scratch
1410/// buffers that consume no input are zeroed, mirroring the borrowed path's
1411/// memset of input-less buffers so a fused accumulator's unfired slots read 0.
1412fn 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        // Logical length is the caller input slice length (input bindings) or
1435        // the buffer's static size (outputs/scratch). The host path binds the
1436        // unused standard scanner buffers (counts/offsets/lengths/metadata) as
1437        // zero-length `&[]`; resident allocation rejects 0 bytes, so allocate
1438        // one element (element-aligned, so the backend's element-size
1439        // validation holds) for those, the kernel never reads a 0/1-element
1440        // unused buffer, so the placeholder is bound but inert (proven equal to
1441        // the host path by the resident/host differential gate).
1442        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        // Upload exactly `alloc_len` bytes so the backend's full-buffer upload
1446        // contract holds: the caller input when it is non-empty, else zeros
1447        // (output/scratch buffers, and the inert zero-length standard inputs).
1448        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
1467/// Byte length to allocate for a binding's resident resource: the caller input
1468/// slice length for input-consuming bindings, else the buffer's static size.
1469fn 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
1483/// Allocate a zero-filled host staging buffer of `byte_len` for initializing a
1484/// resident output/scratch resource.
1485fn 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
1498/// Run the fixpoint sequence resident: every segment dispatched against the
1499/// shared resident resource slice, the whole sequence repeated to the program's
1500/// fixpoint bound, then the final outputs read back by name into `outputs`.
1501fn 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    // Every split segment shares the full program buffer layout, so the same
1520    // resident resource slice binds positionally to each one.
1521    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            // Carry the workgroup too: `grid_override` is sized for this
1529            // workgroup, so dropping it would launch a grid that under-covers
1530            // the work and silently drops findings.
1531            workgroup_override: config.workgroup_override,
1532        });
1533    }
1534
1535    // Read back each original output buffer (declaration order) so the output
1536    // shape is byte-identical to the host-split path.
1537    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    // Size `outputs` to one slot per output buffer, reusing existing
1559    // allocations, then hand the readback mutable references in order.
1560    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
1578/// Free every resident resource allocated for a
1579/// [`dispatch_resident_grid_sync_fixpoint_into`] call. Attempts every free even
1580/// if one fails, returning the first error so a leak is surfaced loudly.
1581fn free_resident_program_resources(
1582    backend: &dyn VyreBackend,
1583    resident: ResidentProgramResources,
1584) -> Result<(), BackendError> {
1585    let ResidentProgramResources { ordered, by_name } = resident;
1586    // `by_name` holds handle clones of the same resources in `ordered`; drop
1587    // it first so each underlying handle is freed exactly once via `ordered`.
1588    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
1689/// Order-independent fingerprint of the EVOLVING accumulator state threaded
1690/// between grid-sync segments.
1691///
1692/// Only `Owned` entries are hashed: a `Borrowed` entry is a caller input that
1693/// is never written by any segment (constant for the whole split), so it cannot
1694/// change between passes and excluding it keeps the fingerprint cheap. Each
1695/// owned buffer mixes its NAME and its bytes (FNV-1a) so a value moving between
1696/// buffers is observed, and the per-buffer hashes are XOR-combined so map
1697/// iteration order does not affect the result. Two consecutive passes with an
1698/// identical fingerprint prove the deterministic segment sequence reached a
1699/// fixpoint (used to early-exit the outer iteration loop).
1700fn 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        // Separator so `name`+`bytes` cannot alias a different split.
1714        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
1832/// After each segment dispatch, overwrite every ReadWrite buffer's
1833/// slot in `inputs` with the freshly-read bytes from `outputs`. The
1834/// backend returns one Vec<u8> per ReadWrite buffer in declaration
1835/// order; this function locates each ReadWrite buffer's input-slot
1836/// index and overwrites it. ReadOnly buffers stay untouched between
1837/// segments.
1838fn 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    // Walk the segment's buffer table twice in lockstep  -  once for the
1845    // input slice, once for the output readback. Both paths must
1846    // mirror the convention `dispatch_borrowed` uses: input position
1847    // skips Workgroup AND `is_output` buffers; output position emits
1848    // one slot per ReadWrite buffer (whether or not is_output).
1849    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        // Refresh the input slot from the readback if this buffer
1859        // appears in BOTH input and output positions (i.e. ReadWrite
1860        // and NOT is_output  -  the rule scratch / `gets` case).
1861        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        // Advance the input cursor for every non-output buffer.
1870        if !is_output_buffer {
1871            input_idx += 1;
1872        }
1873        // Advance the output cursor for every ReadWrite buffer (output
1874        // or not  -  the backend includes them all in the readback).
1875        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    /// Get the inner-segment node count for a wrapped or unwrapped Program.
1934    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        // Original entry was [Region("a", ...)] so the inner sequence is 1.
1952        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        // Region("a"), Barrier(SeqCst), Region("b") = 3 inner nodes.
2025        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    /// Each `dispatch_borrowed_into` reads `inputs[0][0]`, writes `+1`. With the
2176    /// ReadWrite buffer rotating between segments, a single pass over a
2177    /// two-segment program advances the accumulator by 2. The multi-hop
2178    /// `flows_to` closure relies on the WHOLE sequence being re-run
2179    /// `fixpoint_iterations` times (one dataflow hop per pass); a single pass
2180    /// is one hop, which silently dropped every flow through an intermediate
2181    /// variable to recall=0.
2182    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            // Each segment must run exactly once per outer pass: the whole
2211            // sequence carries the fixpoint, not any single segment.
2212            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        // Two segments separated by a GridSync barrier.
2230        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        // Single pass (default): 2 segment launches, accumulator = 2.
2241        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        // Three fixpoint iterations: 3 passes × 2 segments = 6 launches, and
2257        // the accumulator advances one hop per pass to 6. This is the exact
2258        // property the multi-hop `flows_to` split depended on and the
2259        // single-pass implementation lacked.
2260        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    /// A backend-allocated atomic output starts from zero even when split
2289    /// liveness rewrites its first writer as a read-write segment input.
2290    #[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        // The `&dyn VyreBackend` entry and the closure entry both delegate to
2334        // `dispatch_grid_sync_split_generic`, so on the same grid-sync program,
2335        // config, and inputs they must drive the same segment dispatches and
2336        // produce byte-identical output. This is the ONE-PLACE contract that
2337        // lets a host-loop dataflow solver route its fixpoint through the closure
2338        // entry with no separate split implementation.
2339        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        // The closure delegates to an identical backend through the opaque
2368        // single-launch closure shape a host-loop solver supplies (grid override
2369        // only, per-segment fixpoint fixed at 1 by the shared core).
2370        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        // An OUTPUT buffer whose slots are written by DIFFERENT grid-sync
2604        // segments (the fused multi-rule `results_packed` shape: each rule's
2605        // result-store lands in its own segment) must ACCUMULATE across the host
2606        // split. The first writer establishes it (WriteOnly); every LATER writer
2607        // must read the forwarded value and merge its own slots (ReadWrite)
2608        // instead of overwriting it with a fresh write-only buffer, which would
2609        // silently zero the earlier segments' slots (recall=0 for every rule
2610        // whose store is not in the final segment).
2611        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    /// Emulates a backend that lacks native grid-sync: for the single output
2664    /// buffer `out`, it starts from the forwarded prior value (when the segment
2665    /// consumes it) or zeros, then applies that segment's literal `Store out[i]
2666    /// = v` writes, exactly the per-slot store shape a fused multi-rule program
2667    /// produces. Proves end-to-end that earlier segments' slots survive.
2668    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            // Locate `out`'s positional input/output slots using the SAME
2696            // role convention the host split planner uses.
2697            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        // Regression: a fused multi-arm program where arm A's result-store is in
2767        // segment 0 (slot at element 0) and arm B's in the final segment (slot
2768        // at element 2). Before the accumulator fix the final segment's
2769        // write-only `out` zeroed element 0, dropping arm A entirely (a co-fused
2770        // rule whose result-store does not land in the final grid-sync segment
2771        // returned recall=0). Both slots must now survive.
2772        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    /// Copies its input to its output and bumps byte 0 toward a saturation cap.
2812    /// Once the cap is reached the output equals the input, so a full pass over
2813    /// the split leaves the carried accumulator unchanged (a fixpoint).
2814    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        // Two segments (one GridSync barrier). With a generous iteration budget
2858        // of 10, byte 0 saturates at 3, after which a whole pass leaves the
2859        // accumulator unchanged. The outer loop must stop once two consecutive
2860        // passes match instead of burning all 10 iterations.
2861        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        // pass0 -> 2, pass1 -> 3 (saturates mid-pass), pass2 -> 3 (unchanged) =>
2888        // break after pass2. 3 passes x 2 segments = 6 launches, NOT 10x2=20.
2889        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        // The dual of the early-exit test: an accumulator that changes every
2904        // pass (never reaches a fixpoint within budget) must run all
2905        // iterations (early-exit must not fire on a still-advancing closure).
2906        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        // cap=255 so it never saturates within 4 passes (8 increments).
2916        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    /// In-memory device for the resident fixpoint path: holds one byte vector
3044    /// per resident handle, applies a segment's `out` stores IN PLACE to the
3045    /// bound device buffer (no clear between launches), and reads ranges back.
3046    /// `allocate_resident` fills fresh buffers with 0xFF so a test can prove the
3047    /// zero-init upload actually ran.
3048    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            // Fresh device memory is garbage (0xFF here) so the zero-init upload
3111            // path is actually exercised by the test assertions.
3112            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            // Find `out`'s index among the non-shared bindings  -  the same
3163            // order `allocate_resident_program_resources` builds `resources` in.
3164            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            // Apply the segment's `out` stores IN PLACE  -  never clearing the
3182            // buffer, so earlier segments' slots persist (the accumulator).
3183            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        // Same cross-anchor shape as the host-path regression: arm A stores slot
3221        // 0 in segment 0, arm B stores slot 2 in the final segment. The resident
3222        // path keeps ONE device `out` buffer bound across both segments, so both
3223        // slots must survive WITHOUT the host-path accumulator role-rewrite  -
3224        // the persistent device buffer is never cleared between launches.
3225        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        // Zero-init proof: every byte the kernel did not write is 0, not the
3259        // 0xFF garbage `allocate_resident` seeded  -  the output buffer was
3260        // zeroed before dispatch.
3261        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        // Every resident resource is freed exactly once.
3267        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        // With a fixpoint bound > 1, the whole segment sequence repeats that many
3281        // times against the same resident buffers (idempotent stores here, so the
3282        // result is unchanged, but the launch count proves the repeat wiring).
3283        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}