1use 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
17pub const CSR_QUEUE_SPLIT_LOW_FORWARD_OP_ID: &str =
19 "vyre-primitives::graph::csr_queue_split_low_forward_traverse";
20
21pub const CSR_QUEUE_SPLIT_LOW_FORWARD_WORKGROUP_SIZE: [u32; 3] = [256, 1, 1];
23
24pub 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#[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#[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#[derive(Clone, Copy, Debug)]
49pub struct CsrQueueSplitLowForwardParams<'a> {
50 pub active_queue: &'a str,
52 pub queue_len: &'a str,
54 pub edge_offsets: &'a str,
56 pub edge_targets: &'a str,
58 pub edge_kind_mask: &'a str,
60 pub frontier_out: &'a str,
62 pub high_queue: &'a str,
64 pub high_len: &'a str,
66 pub node_count: u32,
68 pub edge_count: u32,
70 pub queue_capacity: u32,
72 pub high_queue_capacity: u32,
74 pub high_degree_threshold: u32,
76 pub allow_mask: u32,
78}
79
80#[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#[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#[cfg(any(test, feature = "cpu-parity"))]
185#[derive(Clone, Debug, Eq, PartialEq)]
186pub struct CsrQueueSplitLowForwardCpuResult {
187 pub frontier_out: Vec<u32>,
189 pub high_queue: Vec<u32>,
191 pub high_len: u32,
193}
194
195#[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;