vyre_driver/grid_sync/
segment_buffers.rs

1//! Per-segment buffer tables for the host split: which buffers a segment
2//! reads, writes, or accumulates, and the access roles that follow from that.
3
4use std::collections::{HashMap, HashSet};
5
6use vyre_foundation::ir::{BufferAccess, BufferDecl, Expr, Ident, MemoryKind, Node, Program};
7
8use super::barrier_split::{entry_sequence, try_split_on_grid_sync};
9use super::{reserve_grid_sync_hash_map, reserve_grid_sync_hash_set, reserve_grid_sync_vec};
10use crate::backend::BackendError;
11
12pub(super) struct PlannedGridSyncSegment {
13    pub(super) program: Program,
14    pub(super) input_names: Vec<Ident>,
15    pub(super) output_names: Vec<Ident>,
16}
17
18/// Diagnostics: the host-split segment **programs** (post buffer-rewrite) that
19/// the host-split dispatch path (`dispatch_with_grid_sync_split*`) validates and
20/// launches when the backend lacks native grid-sync. Exposed so tooling and
21/// tests can inspect or validate each segment without a live backend, the
22/// raw [`try_split_on_grid_sync`] output omits the per-segment buffer
23/// access/role rewrite, so it is not what the backend actually sees.
24///
25/// # Errors
26/// Propagates any [`BackendError`] from splitting or buffer rewriting.
27pub fn plan_host_grid_sync_segment_programs(
28    program: &Program,
29) -> Result<Vec<Program>, BackendError> {
30    Ok(plan_host_grid_sync_segments(program)?
31        .into_iter()
32        .map(|segment| segment.program)
33        .collect())
34}
35
36pub(super) fn plan_host_grid_sync_segments(
37    program: &Program,
38) -> Result<Vec<PlannedGridSyncSegment>, BackendError> {
39    let split = try_split_on_grid_sync(program)?;
40    let first_writer = first_writer_segment_per_buffer(&split, program)?;
41    let mut planned = Vec::new();
42    reserve_grid_sync_vec(&mut planned, split.len(), "grid-sync planned host segments")?;
43    for (segment_idx, segment) in split.into_iter().enumerate() {
44        let rewritten =
45            rewrite_segment_buffers_for_host_split(program, &segment, segment_idx, &first_writer)?;
46        let input_names = segment_input_names(&rewritten)?;
47        let output_names = segment_output_names(&rewritten)?;
48        planned.push(PlannedGridSyncSegment {
49            program: rewritten,
50            input_names,
51            output_names,
52        });
53    }
54    Ok(planned)
55}
56
57/// For each buffer name, the index of the FIRST split segment that writes it.
58///
59/// A source-output buffer written by more than one segment is an
60/// **accumulator**: each segment writes only its own slots (e.g. a fused
61/// multi-rule `results_packed`, where every rule's result-store lands in a
62/// different grid-sync segment). A LATER writer must therefore read+merge the
63/// value forwarded from earlier segments via `current_inputs`, never overwrite
64/// it with a fresh WriteOnly buffer, which would silently zero every earlier
65/// segment's slots (recall=0 for every rule whose store is not in the final
66/// segment). `rewrite_segment_buffers_for_host_split` uses this map to keep an
67/// already-produced output buffer as a `ReadWrite` accumulator in later
68/// segments instead of a write-only output.
69fn first_writer_segment_per_buffer(
70    split: &[Program],
71    program: &Program,
72) -> Result<HashMap<Ident, usize>, BackendError> {
73    let mut first_writer: HashMap<Ident, usize> = HashMap::new();
74    reserve_grid_sync_hash_map(
75        &mut first_writer,
76        program.buffers().len(),
77        "grid-sync first-writer map",
78    )?;
79    for (segment_idx, segment) in split.iter().enumerate() {
80        let mut reads = HashSet::new();
81        let mut writes = HashSet::new();
82        reserve_grid_sync_hash_set(
83            &mut reads,
84            program.buffers().len(),
85            "grid-sync first-writer read scan",
86        )?;
87        reserve_grid_sync_hash_set(
88            &mut writes,
89            program.buffers().len(),
90            "grid-sync first-writer write scan",
91        )?;
92        for node in entry_sequence(segment) {
93            collect_segment_buffer_targets(node, &mut reads, &mut writes);
94        }
95        for name in writes {
96            first_writer.entry(name).or_insert(segment_idx);
97        }
98    }
99    Ok(first_writer)
100}
101
102fn rewrite_segment_buffers_for_host_split(
103    source: &Program,
104    segment: &Program,
105    segment_idx: usize,
106    first_writer: &HashMap<Ident, usize>,
107) -> Result<Program, BackendError> {
108    let mut reads = HashSet::new();
109    let mut writes = HashSet::new();
110    reserve_grid_sync_hash_set(
111        &mut reads,
112        source.buffers().len(),
113        "grid-sync segment read set",
114    )?;
115    reserve_grid_sync_hash_set(
116        &mut writes,
117        source.buffers().len(),
118        "grid-sync segment write set",
119    )?;
120    for node in entry_sequence(segment) {
121        collect_segment_buffer_targets(node, &mut reads, &mut writes);
122    }
123
124    let mut buffers = Vec::new();
125    reserve_grid_sync_vec(
126        &mut buffers,
127        source.buffers().len(),
128        "grid-sync segment buffers",
129    )?;
130    for buffer in source.buffers() {
131        let name = Ident::from(buffer.name());
132        let reads_this = reads.contains(&name);
133        let writes_this = writes.contains(&name);
134        let readwrite_passthrough = matches!(buffer.access(), BufferAccess::ReadWrite)
135            && !buffer.is_output()
136            && !buffer.is_pipeline_live_out()
137            && !reads_this
138            && !writes_this;
139
140        if !reads_this && !writes_this && !readwrite_passthrough {
141            continue;
142        }
143
144        let mut rewritten = buffer.clone();
145        if matches!(rewritten.access(), BufferAccess::Workgroup) {
146            buffers.push(rewritten);
147            continue;
148        }
149
150        // A source-output buffer that an EARLIER segment already wrote is an
151        // accumulator across the split: this segment must read the value
152        // forwarded via `current_inputs` and merge its own slots, never
153        // overwrite it with a fresh WriteOnly buffer (which zeroes the earlier
154        // segments' slots, the silent recall=0 mode for any fused rule whose
155        // result-store does not land in the final segment).
156        let is_source_output = buffer.is_output() || buffer.is_pipeline_live_out();
157        let earlier_segment_wrote_output = is_source_output
158            && first_writer
159                .get(&name)
160                .is_some_and(|&first| first < segment_idx);
161
162        let access = if readwrite_passthrough {
163            BufferAccess::ReadWrite
164        } else if earlier_segment_wrote_output && writes_this {
165            // Later writer of a multi-segment output accumulator: read the
166            // accumulated prior value (uploaded as input) and merge this
167            // segment's slots in place.
168            BufferAccess::ReadWrite
169        } else {
170            match (reads_this, writes_this) {
171                (true, true) => BufferAccess::ReadWrite,
172                (true, false) => BufferAccess::ReadOnly,
173                (false, true) => BufferAccess::WriteOnly,
174                (false, false) => BufferAccess::ReadWrite,
175            }
176        };
177        rewrite_segment_buffer_access(&mut rewritten, access);
178        // Never mark a split segment's buffer as the program output: a
179        // multi-segment output accumulator must CONSUME its forwarded prior
180        // value as input in later segments, and `segment_buffer_consumes_input`
181        // refuses any `is_output` buffer. Each writing segment still produces
182        // the buffer (WriteOnly/ReadWrite both produce output), so its bytes
183        // are captured into `current_inputs`; the final host-visible values are
184        // reassembled by name from the SOURCE program's output set in
185        // `collect_final_named_outputs`, independent of any per-segment flag.
186        rewritten.is_output = false;
187        rewritten.pipeline_live_out = false;
188        buffers.push(rewritten);
189    }
190
191    Ok(segment.with_rewritten_buffers(buffers))
192}
193
194fn rewrite_segment_buffer_access(buffer: &mut BufferDecl, access: BufferAccess) {
195    buffer.kind = match &access {
196        BufferAccess::ReadOnly => MemoryKind::Readonly,
197        BufferAccess::Uniform => MemoryKind::Uniform,
198        BufferAccess::Workgroup => MemoryKind::Shared,
199        _ => MemoryKind::Global,
200    };
201    buffer.access = access;
202}
203
204pub(super) fn segment_input_names(segment: &Program) -> Result<Vec<Ident>, BackendError> {
205    let mut names = Vec::new();
206    reserve_grid_sync_vec(
207        &mut names,
208        segment.buffers().len(),
209        "grid-sync segment input names",
210    )?;
211    for buffer in segment.buffers() {
212        if matches!(buffer.access(), BufferAccess::Workgroup) {
213            continue;
214        }
215        if segment_buffer_consumes_input(buffer) {
216            names.push(Ident::from(buffer.name()));
217        }
218    }
219    Ok(names)
220}
221
222pub(super) fn segment_output_names(segment: &Program) -> Result<Vec<Ident>, BackendError> {
223    let mut names = Vec::new();
224    reserve_grid_sync_vec(
225        &mut names,
226        segment.buffers().len(),
227        "grid-sync segment output names",
228    )?;
229    for buffer in segment.buffers() {
230        if matches!(buffer.access(), BufferAccess::Workgroup) {
231            continue;
232        }
233        if segment_buffer_produces_output(buffer) {
234            names.push(Ident::from(buffer.name()));
235        }
236    }
237    Ok(names)
238}
239
240pub(super) fn original_input_names(program: &Program) -> Result<Vec<Ident>, BackendError> {
241    segment_input_names(program)
242}
243
244pub(super) fn original_output_names(program: &Program) -> Result<Vec<Ident>, BackendError> {
245    segment_output_names(program)
246}
247
248pub(super) fn segment_buffer_consumes_input(buffer: &BufferDecl) -> bool {
249    if buffer.is_output() || buffer.is_pipeline_live_out() {
250        return false;
251    }
252    matches!(
253        buffer.access(),
254        BufferAccess::ReadOnly | BufferAccess::ReadWrite | BufferAccess::Uniform
255    )
256}
257
258pub(super) fn segment_buffer_produces_output(buffer: &BufferDecl) -> bool {
259    buffer.is_output()
260        || buffer.is_pipeline_live_out()
261        || matches!(
262            buffer.access(),
263            BufferAccess::ReadWrite | BufferAccess::WriteOnly
264        )
265}
266
267fn collect_segment_buffer_targets(
268    node: &Node,
269    reads: &mut HashSet<Ident>,
270    writes: &mut HashSet<Ident>,
271) {
272    match node {
273        Node::Let { value, .. } | Node::Assign { value, .. } => {
274            collect_segment_expr_targets(value, reads, writes);
275        }
276        Node::Store {
277            buffer,
278            index,
279            value,
280        } => {
281            writes.insert(Ident::from(buffer));
282            collect_segment_expr_targets(index, reads, writes);
283            collect_segment_expr_targets(value, reads, writes);
284        }
285        Node::If {
286            cond,
287            then,
288            otherwise,
289        } => {
290            collect_segment_expr_targets(cond, reads, writes);
291            for child in then.iter().chain(otherwise.iter()) {
292                collect_segment_buffer_targets(child, reads, writes);
293            }
294        }
295        Node::Loop { from, to, body, .. } => {
296            collect_segment_expr_targets(from, reads, writes);
297            collect_segment_expr_targets(to, reads, writes);
298            for child in body {
299                collect_segment_buffer_targets(child, reads, writes);
300            }
301        }
302        Node::Block(body) => {
303            for child in body {
304                collect_segment_buffer_targets(child, reads, writes);
305            }
306        }
307        Node::Region { body, .. } => {
308            for child in body.iter() {
309                collect_segment_buffer_targets(child, reads, writes);
310            }
311        }
312        Node::AllReduce { buffer, .. } | Node::Broadcast { buffer, .. } => {
313            reads.insert(buffer.clone());
314            writes.insert(buffer.clone());
315        }
316        Node::AllGather { input, output, .. } | Node::ReduceScatter { input, output, .. } => {
317            reads.insert(input.clone());
318            writes.insert(output.clone());
319        }
320        Node::IndirectDispatch { .. }
321        | Node::Return
322        | Node::Barrier { .. }
323        | Node::AsyncLoad { .. }
324        | Node::AsyncStore { .. }
325        | Node::AsyncWait { .. }
326        | Node::Trap { .. }
327        | Node::Resume { .. }
328        | Node::Opaque(_) => {}
329        _ => {}
330    }
331}
332
333fn collect_segment_expr_targets(
334    expr: &Expr,
335    reads: &mut HashSet<Ident>,
336    writes: &mut HashSet<Ident>,
337) {
338    vyre_foundation::visit::visit_expr_buffer_accesses(expr, |access, buffer| {
339        reads.insert(buffer.clone());
340        if access == vyre_foundation::visit::ExprBufferAccess::Atomic {
341            writes.insert(buffer.clone());
342        }
343    });
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349    use crate::grid_sync::test_programs::region;
350    use vyre_foundation::ir::DataType;
351    use vyre_foundation::memory_model::MemoryOrdering;
352
353    #[test]
354    fn split_keeps_multi_segment_output_as_readwrite_accumulator() {
355        // An OUTPUT buffer whose slots are written by DIFFERENT grid-sync
356        // segments (the fused multi-rule `results_packed` shape: each rule's
357        // result-store lands in its own segment) must ACCUMULATE across the host
358        // split. The first writer establishes it (WriteOnly); every LATER writer
359        // must read the forwarded value and merge its own slots (ReadWrite)
360        // instead of overwriting it with a fresh write-only buffer, which would
361        // silently zero the earlier segments' slots (recall=0 for every rule
362        // whose store is not in the final segment).
363        let out = BufferDecl::output("out", 0, DataType::U32).with_count(4);
364        let program = Program::wrapped(
365            vec![out],
366            [1, 1, 1],
367            vec![
368                region("a", vec![Node::store("out", Expr::u32(0), Expr::u32(0xAA))]),
369                Node::barrier_with_ordering(MemoryOrdering::GridSync),
370                region("b", vec![Node::store("out", Expr::u32(2), Expr::u32(0xBB))]),
371            ],
372        );
373        let segments =
374            plan_host_grid_sync_segment_programs(&program).expect("plan host grid-sync segments");
375        assert_eq!(segments.len(), 2, "one GridSync barrier -> two segments");
376
377        let seg0_out = segments[0]
378            .buffers()
379            .iter()
380            .find(|b| b.name() == "out")
381            .expect("segment 0 must declare the output it writes");
382        assert_eq!(
383            seg0_out.access(),
384            BufferAccess::WriteOnly,
385            "the first writer establishes the accumulator as write-only"
386        );
387        assert!(
388            !seg0_out.is_output() && !seg0_out.is_pipeline_live_out(),
389            "split segment buffers must never be marked program-output; final values are reassembled by name"
390        );
391
392        let seg1_out = segments[1]
393            .buffers()
394            .iter()
395            .find(|b| b.name() == "out")
396            .expect("segment 1 must declare the output it writes");
397        assert_eq!(
398            seg1_out.access(),
399            BufferAccess::ReadWrite,
400            "a later writer of a multi-segment output must read+merge the accumulated value, not overwrite it"
401        );
402        assert!(
403            !seg1_out.is_output() && !seg1_out.is_pipeline_live_out(),
404            "the later writer must consume its forwarded prior value, which `segment_buffer_consumes_input` refuses for is_output buffers"
405        );
406        assert!(
407            segment_input_names(&segments[1])
408                .expect("segment 1 input names")
409                .iter()
410                .any(|n| n.as_str() == "out"),
411            "the accumulated output must be forwarded as an input to the later writing segment"
412        );
413    }
414}
vyre_driver/grid_sync/segment_buffers.rs

vyre_driver/grid_sync/
segment_buffers.rs
