1use vyre_foundation::ir::{Expr, Node};
7
8pub fn load_packed_byte(haystack: &str, idx: Expr) -> (Node, Expr) {
16 let word_idx = Expr::div(idx.clone(), Expr::u32(4));
17 let byte_offset = Expr::mul(Expr::rem(idx, Expr::u32(4)), Expr::u32(8));
18
19 let node = Node::let_bind("_byte_word", Expr::load(haystack, word_idx));
20 let byte_expr = Expr::bitand(
21 Expr::shr(Expr::var("_byte_word"), byte_offset),
22 Expr::u32(0xFF),
23 );
24
25 (node, byte_expr)
26}
27
28pub fn load_packed_byte_expr(haystack: &str, idx: Expr) -> Expr {
44 Expr::bitand(
45 Expr::shr(
46 Expr::load(haystack, Expr::div(idx.clone(), Expr::u32(4))),
47 Expr::mul(Expr::rem(idx, Expr::u32(4)), Expr::u32(8)),
48 ),
49 Expr::u32(0xFF),
50 )
51}
52
53pub fn append_match(
58 hits_buffer: &str,
59 count_buffer: &str,
60 tag: impl Into<Expr>,
61 start: impl Into<Expr>,
62 end: impl Into<Expr>,
63) -> Node {
64 let max_hits = Expr::div(Expr::buf_len(hits_buffer), Expr::u32(3));
65
66 Node::Block(vec![
67 Node::let_bind(
68 "_vyre_match_slot",
69 Expr::atomic_add(count_buffer, Expr::u32(0), Expr::u32(1)),
70 ),
71 Node::if_then(
72 Expr::lt(Expr::var("_vyre_match_slot"), max_hits),
73 vec![
74 Node::store(
75 hits_buffer,
76 Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
77 tag.into(),
78 ),
79 Node::store(
80 hits_buffer,
81 Expr::add(
82 Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
83 Expr::u32(1),
84 ),
85 start.into(),
86 ),
87 Node::store(
88 hits_buffer,
89 Expr::add(
90 Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
91 Expr::u32(2),
92 ),
93 end.into(),
94 ),
95 ],
96 ),
97 ])
98}
99
100pub fn append_match_subgroup(
106 hits_buffer: &str,
107 count_buffer: &str,
108 tag: impl Into<Expr>,
109 start: impl Into<Expr>,
110 end: impl Into<Expr>,
111 cond: Expr,
112) -> Vec<Node> {
113 let tag = tag.into();
114 let start = start.into();
115 let end = end.into();
116 let max_hits = Expr::div(Expr::buf_len(hits_buffer), Expr::u32(3));
117 let lane_mask = Expr::sub(
118 Expr::shl(Expr::u32(1), Expr::subgroup_local_id()),
119 Expr::u32(1),
120 );
121 let rank = Expr::popcount(Expr::bitand(Expr::var("_vyre_match_ballot"), lane_mask));
122 let leader_pred = Expr::and(
123 cond.clone(),
124 Expr::eq(Expr::var("_vyre_match_rank"), Expr::u32(0)),
125 );
126 let slot = Expr::add(
127 Expr::subgroup_shuffle(
128 Expr::var("_vyre_match_leader_base"),
129 Expr::var("_vyre_match_leader"),
130 ),
131 Expr::var("_vyre_match_rank"),
132 );
133 let ballot_cond = cond.clone();
134 let bounded_hit = Expr::and(cond, Expr::lt(Expr::var("_vyre_match_slot"), max_hits));
135
136 vec![
137 Node::let_bind("_vyre_match_ballot", Expr::subgroup_ballot(ballot_cond)),
138 Node::let_bind("_vyre_match_rank", rank),
139 Node::let_bind(
140 "_vyre_match_count",
141 Expr::popcount(Expr::var("_vyre_match_ballot")),
142 ),
143 Node::let_bind(
144 "_vyre_match_leader",
145 Expr::select(
146 Expr::eq(Expr::var("_vyre_match_count"), Expr::u32(0)),
147 Expr::u32(0),
148 Expr::ctz(Expr::var("_vyre_match_ballot")), ),
150 ),
151 Node::let_bind("_vyre_match_leader_base", Expr::u32(0)),
152 Node::if_then(
153 leader_pred,
154 vec![Node::assign(
155 "_vyre_match_leader_base",
156 Expr::atomic_add(count_buffer, Expr::u32(0), Expr::var("_vyre_match_count")),
157 )],
158 ),
159 Node::let_bind("_vyre_match_slot", slot),
160 Node::if_then(
161 bounded_hit,
162 vec![
163 Node::store(
164 hits_buffer,
165 Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
166 tag,
167 ),
168 Node::store(
169 hits_buffer,
170 Expr::add(
171 Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
172 Expr::u32(1),
173 ),
174 start,
175 ),
176 Node::store(
177 hits_buffer,
178 Expr::add(
179 Expr::mul(Expr::var("_vyre_match_slot"), Expr::u32(3)),
180 Expr::u32(2),
181 ),
182 end,
183 ),
184 ],
185 ),
186 ]
187}
188
189#[cfg(test)]
190mod subgroup_append_shape_tests {
191 use super::*;
192
193 #[test]
194 fn subgroup_append_bounds_hits_with_bound_slot_variable() {
195 let nodes = append_match_subgroup(
196 "matches",
197 "match_count",
198 Expr::u32(7),
199 Expr::u32(11),
200 Expr::u32(13),
201 Expr::var("hit"),
202 );
203
204 let bounded_cond = match nodes.last() {
205 Some(Node::If { cond, .. }) => format!("{cond:?}"),
206 other => panic!("expected bounded-hit If as final subgroup append node, got {other:?}"),
207 };
208
209 assert!(
210 bounded_cond.contains("_vyre_match_slot"),
211 "bounded-hit predicate must use the already-bound slot variable: {bounded_cond}"
212 );
213 assert!(
214 !bounded_cond.contains("_vyre_match_leader"),
215 "bounded-hit predicate must not re-inline subgroup shuffle leader expressions: {bounded_cond}"
216 );
217 }
218}
219
220#[cfg(test)]
221mod packed_byte_dedup_lock {
222 use std::path::{Path, PathBuf};
241
242 fn vyre_libs_src() -> PathBuf {
243 let crate_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
244 crate_root.join("src")
245 }
246
247 fn walk(dir: &Path, out: &mut Vec<PathBuf>) {
248 let Ok(entries) = std::fs::read_dir(dir) else {
249 return;
250 };
251 for entry in entries.flatten() {
252 let path = entry.path();
253 if path.is_dir() {
254 walk(&path, out);
255 } else if path.extension().and_then(|e| e.to_str()) == Some("rs") {
256 out.push(path);
257 }
258 }
259 }
260}