Skip to main content

vyre_primitives/graph/
csr_queue_split.rs

1//! Mixed sparse CSR queue traversal for active sets with a small number of hubs.
2//!
3//! A global row-strided pass is excellent for true hub rows, but wastes lanes on
4//! the many one-edge and three-edge rows that usually travel in the same active
5//! queue. This primitive keeps low-degree rows in a scalar queue pass and
6//! compacts only high-degree sources into a second queue for row-strided
7//! traversal.
8
9use vyre_foundation::ir::{DataType, Program};
10
11use crate::graph::csr_frontier_step::{
12    csr_queue_step_program, CsrQueueEmit, CsrQueueInputs, CsrQueueLanes, CsrQueueRowPlan,
13    CsrQueueStepSpec,
14};
15use crate::graph::csr_queue_strided::CSR_QUEUE_STRIDED_FORWARD_LANES_PER_SOURCE;
16
17/// Canonical op id for mixed low-row traversal and high-row compaction.
18pub const CSR_QUEUE_SPLIT_LOW_FORWARD_OP_ID: &str =
19    "vyre-primitives::graph::csr_queue_split_low_forward_traverse";
20
21/// Workgroup shape for the low-row split pass.
22pub const CSR_QUEUE_SPLIT_LOW_FORWARD_WORKGROUP_SIZE: [u32; 3] = [256, 1, 1];
23
24/// Degree at which a queued row has enough work to amortize a 32-lane team.
25pub const CSR_QUEUE_SPLIT_HIGH_DEGREE_THRESHOLD: u32 =
26    CSR_QUEUE_STRIDED_FORWARD_LANES_PER_SOURCE * CSR_QUEUE_STRIDED_FORWARD_LANES_PER_SOURCE;
27
28/// Dispatch grid for the one-lane-per-active-source low split pass.
29#[must_use]
30pub const fn csr_queue_split_low_dispatch_grid(queue_capacity: u32) -> [u32; 3] {
31    let blocks = queue_capacity.div_ceil(CSR_QUEUE_SPLIT_LOW_FORWARD_WORKGROUP_SIZE[0]);
32    [if blocks == 0 { 1 } else { blocks }, 1, 1]
33}
34
35/// Logical lanes consumed by low split plus a high row-strided follow-up pass.
36#[must_use]
37pub const fn csr_queue_split_mixed_logical_lanes(
38    queue_capacity: u32,
39    high_queue_capacity: u32,
40) -> u64 {
41    (queue_capacity as u64).saturating_add(
42        (high_queue_capacity as u64)
43            .saturating_mul(CSR_QUEUE_STRIDED_FORWARD_LANES_PER_SOURCE as u64),
44    )
45}
46
47/// Positional inputs for [`csr_queue_split_low_forward_traverse`].
48#[derive(Clone, Copy, Debug)]
49pub struct CsrQueueSplitLowForwardParams<'a> {
50    /// Compacted queue of active source nodes.
51    pub active_queue: &'a str,
52    /// Single-element resident length of `active_queue`.
53    pub queue_len: &'a str,
54    /// CSR row pointers, `node_count + 1` entries.
55    pub edge_offsets: &'a str,
56    /// CSR edge destinations.
57    pub edge_targets: &'a str,
58    /// Per-edge kind bits tested against `allow_mask`.
59    pub edge_kind_mask: &'a str,
60    /// Packed bitset the reached destinations are ORed into.
61    pub frontier_out: &'a str,
62    /// Compact queue collecting hub sources for a later row-strided pass.
63    pub high_queue: &'a str,
64    /// Single-element observed hub count, which may exceed the capacity.
65    pub high_len: &'a str,
66    /// Node count the CSR row pointers and destination bounds are sized by.
67    pub node_count: u32,
68    /// Logical edge count the edge-slot bound check uses.
69    pub edge_count: u32,
70    /// Static capacity of `active_queue`.
71    pub queue_capacity: u32,
72    /// Static capacity of `high_queue`.
73    pub high_queue_capacity: u32,
74    /// Row degree at which a source is worth a 32-lane team.
75    pub high_degree_threshold: u32,
76    /// Edge kinds this traversal is allowed to follow.
77    pub allow_mask: u32,
78}
79
80/// Build the low-row half of a mixed queue traversal.
81///
82/// Low-degree rows are expanded directly into `frontier_out`. High-degree rows
83/// are appended to `high_queue` and counted in `high_len`; callers then run
84/// `csr_queue_strided_forward_traverse` over that compact high queue. If
85/// `high_queue` is undersized, overflow high rows are expanded by the scalar
86/// lane in this pass so correctness does not depend on perfect sizing.
87#[must_use]
88#[allow(clippy::too_many_arguments)]
89pub fn csr_queue_split_low_forward_traverse(
90    active_queue: &str,
91    queue_len: &str,
92    edge_offsets: &str,
93    edge_targets: &str,
94    edge_kind_mask: &str,
95    frontier_out: &str,
96    high_queue: &str,
97    high_len: &str,
98    node_count: u32,
99    edge_count: u32,
100    queue_capacity: u32,
101    high_queue_capacity: u32,
102    high_degree_threshold: u32,
103    allow_mask: u32,
104) -> Program {
105    csr_queue_split_low_forward_traverse_with(CsrQueueSplitLowForwardParams {
106        active_queue,
107        queue_len,
108        edge_offsets,
109        edge_targets,
110        edge_kind_mask,
111        frontier_out,
112        high_queue,
113        high_len,
114        node_count,
115        edge_count,
116        queue_capacity,
117        high_queue_capacity,
118        high_degree_threshold,
119        allow_mask,
120    })
121}
122
123/// Build the low-row half of a mixed queue traversal.
124#[must_use]
125pub fn csr_queue_split_low_forward_traverse_with(
126    params: CsrQueueSplitLowForwardParams<'_>,
127) -> Program {
128    let CsrQueueSplitLowForwardParams {
129        active_queue,
130        queue_len,
131        edge_offsets,
132        edge_targets,
133        edge_kind_mask,
134        frontier_out,
135        high_queue,
136        high_len,
137        node_count,
138        edge_count,
139        queue_capacity,
140        high_queue_capacity,
141        high_degree_threshold,
142        allow_mask,
143    } = params;
144    if node_count == 0
145        || queue_capacity == 0
146        || high_queue_capacity == 0
147        || high_degree_threshold == 0
148    {
149        return crate::invalid_output_program(CSR_QUEUE_SPLIT_LOW_FORWARD_OP_ID,
150        frontier_out,
151        DataType::U32,
152        format!(
153            "Fix: csr_queue_split_low_forward_traverse requires node_count > 0, non-zero queue capacities, and high_degree_threshold > 0; got node_count={node_count} queue_capacity={queue_capacity} high_queue_capacity={high_queue_capacity} high_degree_threshold={high_degree_threshold}."
154        ),);
155    }
156    csr_queue_step_program(&CsrQueueStepSpec {
157        op_id: CSR_QUEUE_SPLIT_LOW_FORWARD_OP_ID,
158        builder_name: "csr_queue_split_low_forward_traverse",
159        prefix: "qsl",
160        workgroup_size: CSR_QUEUE_SPLIT_LOW_FORWARD_WORKGROUP_SIZE,
161        inputs: CsrQueueInputs {
162            active_queue,
163            queue_len,
164            edge_offsets,
165            edge_targets,
166            edge_kind_mask,
167        },
168        lanes: CsrQueueLanes::Scalar,
169        row_plan: CsrQueueRowPlan::CompactHighDegree {
170            high_queue,
171            high_len,
172            high_queue_capacity,
173            high_degree_threshold,
174        },
175        emit: CsrQueueEmit::Frontier { frontier_out },
176        node_count,
177        edge_count,
178        queue_capacity,
179        allow_mask,
180    })
181}
182
183/// CPU result for the low split pass.
184#[cfg(any(test, feature = "cpu-parity"))]
185#[derive(Clone, Debug, Eq, PartialEq)]
186pub struct CsrQueueSplitLowForwardCpuResult {
187    /// Frontier bitset after low-degree rows and overflow high rows were emitted.
188    pub frontier_out: Vec<u32>,
189    /// Compact queue of high-degree sources that fit in the high queue capacity.
190    pub high_queue: Vec<u32>,
191    /// Total high-degree source count observed, including entries beyond capacity.
192    pub high_len: u32,
193}
194
195/// Fallible CPU reference for the low split pass.
196#[cfg(any(test, feature = "cpu-parity"))]
197#[allow(clippy::too_many_arguments)]
198pub fn try_csr_queue_split_low_forward_traverse_cpu(
199    active_queue: &[u32],
200    queue_len: u32,
201    edge_offsets: &[u32],
202    edge_targets: &[u32],
203    edge_kind_mask: &[u32],
204    frontier_out_seed: &[u32],
205    node_count: u32,
206    high_queue_capacity: usize,
207    high_degree_threshold: u32,
208    allow_mask: u32,
209) -> Result<CsrQueueSplitLowForwardCpuResult, String> {
210    let layout = super::csr_frontier_queue::validate_csr_queue_graph(
211        node_count,
212        edge_offsets,
213        edge_targets,
214        edge_kind_mask,
215    )?;
216    if frontier_out_seed.len() != layout.words {
217        return Err(format!(
218            "Fix: csr_queue_split_low_forward_traverse requires frontier_out_seed.len() == bitset_words(node_count), got len={} but expected {} for node_count={node_count}.",
219            frontier_out_seed.len(),
220            layout.words
221        ));
222    }
223    let mut high_queue_probe: Vec<u32> = Vec::new();
224    crate::graph::scratch::reserve_graph_items(
225        &mut high_queue_probe,
226        high_queue_capacity,
227        "CSR queue split CPU oracle",
228        "high-degree active queue",
229    )?;
230
231    let mut frontier_out = frontier_out_seed.to_vec();
232    let mut high_queue = Vec::with_capacity(high_queue_capacity);
233    let mut high_len = 0_u32;
234    let take = (queue_len as usize).min(active_queue.len());
235
236    for &src in &active_queue[..take] {
237        if src >= node_count {
238            continue;
239        }
240        let start = edge_offsets[src as usize] as usize;
241        let end = edge_offsets[src as usize + 1] as usize;
242        if end.saturating_sub(start) as u32 >= high_degree_threshold {
243            high_len = high_len.saturating_add(1);
244            if high_queue.len() < high_queue_capacity {
245                high_queue.push(src);
246                continue;
247            }
248        }
249        emit_scalar_row_cpu(
250            start,
251            end,
252            edge_targets,
253            edge_kind_mask,
254            node_count,
255            allow_mask,
256            &mut frontier_out,
257        );
258    }
259
260    Ok(CsrQueueSplitLowForwardCpuResult {
261        frontier_out,
262        high_queue,
263        high_len,
264    })
265}
266
267#[cfg(any(test, feature = "cpu-parity"))]
268fn emit_scalar_row_cpu(
269    start: usize,
270    end: usize,
271    edge_targets: &[u32],
272    edge_kind_mask: &[u32],
273    node_count: u32,
274    allow_mask: u32,
275    frontier_out: &mut [u32],
276) {
277    for edge in start..end {
278        if edge_kind_mask[edge] & allow_mask == 0 {
279            continue;
280        }
281        let dst = edge_targets[edge];
282        if dst >= node_count {
283            continue;
284        }
285        frontier_out[dst as usize / 32] |= 1_u32 << (dst % 32);
286    }
287}
288
289#[cfg(test)]
290mod tests;