Skip to main content

vyre_primitives/graph/csr_forward_or_changed/
layout.rs

1use crate::graph::program_graph::BINDING_PRIMITIVE_START;
2
3/// Canonical op id.
4pub(crate) const OP_ID: &str = "vyre-primitives::graph::csr_forward_or_changed";
5/// Canonical binding index for the frontier accumulator.
6pub(crate) const CSR_FORWARD_OR_CHANGED_FRONTIER_BUFFER: u32 = BINDING_PRIMITIVE_START;
7/// Canonical binding index for the changed flag/history buffer.
8pub(crate) const CSR_FORWARD_OR_CHANGED_CHANGED_BUFFER: u32 = BINDING_PRIMITIVE_START + 1;
9/// Canonical one-lane workgroup for CSR forward-or-changed programs.
10pub(crate) const CSR_FORWARD_OR_CHANGED_WORKGROUP_SIZE: [u32; 3] = [1, 1, 1];
11/// Source-lane workgroup for node-parallel CSR forward-or-changed programs.
12pub(crate) const CSR_FORWARD_OR_CHANGED_PARALLEL_WORKGROUP_SIZE: [u32; 3] = [256, 1, 1];
13/// Iteration ceiling where a changed-history buffer avoids per-iteration zeroing.
14pub(crate) const CSR_FORWARD_OR_CHANGED_HISTORY_FAST_PATH_MAX_ITERS: u32 = 64;
15
16/// Dispatch grid for a node-parallel CSR forward-or-changed pass.
17#[must_use]
18pub const fn csr_forward_or_changed_parallel_grid(node_count: u32) -> [u32; 3] {
19    [
20        ceil_div_u32(
21            at_least_one(node_count),
22            CSR_FORWARD_OR_CHANGED_PARALLEL_WORKGROUP_SIZE[0],
23        ),
24        1,
25        1,
26    ]
27}
28
29/// Dispatch grid for a batched node-parallel CSR forward-or-changed pass.
30#[must_use]
31pub const fn csr_forward_or_changed_parallel_batch_grid(
32    node_count: u32,
33    query_count: u32,
34) -> [u32; 3] {
35    [
36        ceil_div_u32(
37            at_least_one(node_count),
38            CSR_FORWARD_OR_CHANGED_PARALLEL_WORKGROUP_SIZE[0],
39        ),
40        at_least_one(query_count),
41        1,
42    ]
43}
44
45const fn at_least_one(value: u32) -> u32 {
46    if value == 0 {
47        1
48    } else {
49        value
50    }
51}
52
53const fn ceil_div_u32(value: u32, divisor: u32) -> u32 {
54    ((value - 1) / divisor) + 1
55}
56
57/// Validated dispatch layout for the forward-or-changed CSR primitive.
58///
59/// The primitive owns these derived counts so dispatch wrappers do not fork CSR
60/// offset, edge-array, frontier, or scratch sizing rules.
61#[derive(Clone, Copy, Debug, Eq, PartialEq)]
62pub struct CsrForwardOrChangedLayout {
63    /// Number of nodes accepted by the primitive.
64    pub node_count: u32,
65    /// Number of words required by node-indexed scratch buffers.
66    pub node_words: usize,
67    /// Number of words required by the edge-offset buffer.
68    pub edge_offset_words: usize,
69    /// Number of edge-array words supplied to the primitive.
70    pub edge_storage_words: usize,
71    /// Edge count used when constructing [`ProgramGraphShape`].
72    pub shape_edge_count: u32,
73    /// Number of frontier words used by the dispatch buffer.
74    pub frontier_words: usize,
75}
76
77/// Program identity for the forward-or-changed CSR primitive.
78///
79/// Dispatch consumers can cache generated programs by this key without
80/// re-implementing CSR validation, changed-history selection, or launch-grid
81/// policy outside `vyre-primitives`.
82#[derive(Clone, Copy, Debug, Eq, PartialEq)]
83pub struct CsrForwardOrChangedProgramKey {
84    layout: CsrForwardOrChangedLayout,
85    allow_mask: u32,
86    changed_slots: u32,
87    uses_changed_history: bool,
88}
89
90/// Primitive-owned identity for reusable CSR forward-or-changed static inputs.
91///
92/// Dispatch wrappers stage edge offsets, targets, masks, and changed-history
93/// buffers according to the primitive launch plan. This key keeps the content
94/// identity next to that plan so wrappers do not fork graph-fingerprint rules.
95#[derive(Clone, Copy, Debug, Eq, PartialEq)]
96pub struct CsrForwardOrChangedStaticInputKey {
97    /// Program identity selected by the primitive launch planner.
98    pub program_key: CsrForwardOrChangedProgramKey,
99    /// Words in the staged edge-offset input.
100    pub edge_offset_words: usize,
101    /// Words in each staged edge-indexed input.
102    pub edge_storage_words: usize,
103    /// Words in the changed readback/scratch buffer.
104    pub changed_words: usize,
105    /// Stable fingerprint of the padded edge-offset upload.
106    pub edge_offsets_hash: u64,
107    /// Stable fingerprint of the padded edge-target upload.
108    pub edge_targets_hash: u64,
109    /// Stable fingerprint of the padded edge-kind upload.
110    pub edge_kind_mask_hash: u64,
111}
112
113impl CsrForwardOrChangedProgramKey {
114    #[must_use]
115    pub(crate) const fn new(
116        layout: CsrForwardOrChangedLayout,
117        allow_mask: u32,
118        changed_slots: u32,
119        uses_changed_history: bool,
120    ) -> Self {
121        Self {
122            layout,
123            allow_mask,
124            changed_slots,
125            uses_changed_history,
126        }
127    }
128
129    /// Validated CSR/frontier layout represented by this program.
130    #[must_use]
131    pub const fn layout(&self) -> CsrForwardOrChangedLayout {
132        self.layout
133    }
134
135    /// Edge-kind mask accepted by this program.
136    #[must_use]
137    pub const fn allow_mask(&self) -> u32 {
138        self.allow_mask
139    }
140
141    /// Number of changed-buffer slots this program writes.
142    #[must_use]
143    pub const fn changed_slots(&self) -> u32 {
144        self.changed_slots
145    }
146
147    /// True when this program uses the dynamic changed-history fast path.
148    #[must_use]
149    pub const fn uses_changed_history(&self) -> bool {
150        self.uses_changed_history
151    }
152}