vyre_foundation/ir_inner/model/node/impl_node.rs
1use super::{Node, NodeExtension};
2use crate::ir_inner::model::expr::{Expr, Ident};
3use crate::memory_model::MemoryOrdering;
4use std::sync::Arc;
5
6impl Node {
7 /// `let name = value;`
8 ///
9 /// # Examples
10 ///
11 /// ```
12 /// use vyre::ir::{Expr, Node};
13 /// let _ = Node::let_bind("x", Expr::u32(1));
14 /// ```
15 #[must_use]
16 #[inline]
17 pub fn let_bind(name: impl Into<Ident>, value: Expr) -> Self {
18 Self::Let {
19 name: name.into(),
20 value,
21 }
22 }
23
24 /// `name = value;`
25 ///
26 /// # Examples
27 ///
28 /// ```
29 /// use vyre::ir::{Expr, Node};
30 /// let _ = Node::assign("x", Expr::u32(2));
31 /// ```
32 #[must_use]
33 #[inline]
34 pub fn assign(name: impl Into<Ident>, value: Expr) -> Self {
35 Self::Assign {
36 name: name.into(),
37 value,
38 }
39 }
40
41 /// `buffer[index] = value;`
42 ///
43 /// # Examples
44 ///
45 /// ```
46 /// use vyre::ir::{Expr, Node};
47 /// let _ = Node::store("out", Expr::u32(0), Expr::u32(1));
48 /// ```
49 #[must_use]
50 #[inline]
51 pub fn store(buffer: impl Into<Ident>, index: Expr, value: Expr) -> Self {
52 Self::Store {
53 buffer: buffer.into(),
54 index,
55 value,
56 }
57 }
58
59 /// `if cond { then } else { otherwise }`
60 ///
61 /// # Examples
62 ///
63 /// ```
64 /// use vyre::ir::{Expr, Node};
65 /// let _ = Node::if_then_else(Expr::bool(true), vec![Node::Return], vec![]);
66 /// ```
67 #[must_use]
68 #[inline]
69 pub fn if_then_else(cond: Expr, then: Vec<Self>, otherwise: Vec<Self>) -> Self {
70 Self::If {
71 cond,
72 then,
73 otherwise,
74 }
75 }
76
77 /// `if cond { then }`
78 ///
79 /// # Examples
80 ///
81 /// ```
82 /// use vyre::ir::{Expr, Node};
83 /// let _ = Node::if_then(Expr::bool(true), vec![Node::Return]);
84 /// ```
85 #[must_use]
86 #[inline]
87 pub fn if_then(cond: Expr, then: Vec<Self>) -> Self {
88 Self::If {
89 cond,
90 then,
91 otherwise: Vec::new(),
92 }
93 }
94
95 /// `for var in from..to { body }`
96 ///
97 /// # Examples
98 ///
99 /// ```
100 /// use vyre::ir::{Expr, Node};
101 /// let _ = Node::loop_for("i", Expr::u32(0), Expr::u32(4), vec![]);
102 /// ```
103 #[must_use]
104 #[inline]
105 pub fn loop_for(var: impl Into<Ident>, from: Expr, to: Expr, body: Vec<Self>) -> Self {
106 Self::Loop {
107 var: var.into(),
108 from,
109 to,
110 body,
111 }
112 }
113
114 /// `for var in from..to { body }`
115 ///
116 /// # Examples
117 ///
118 /// ```
119 /// use vyre::ir::{Expr, Node};
120 ///
121 /// let node = Node::loop_("i", Expr::u32(0), Expr::u32(4), vec![Node::Return]);
122 /// assert!(matches!(node, Node::Loop { .. }));
123 /// ```
124 #[must_use]
125 #[inline]
126 pub fn loop_(var: impl Into<Ident>, from: Expr, to: Expr, body: Vec<Self>) -> Self {
127 Self::loop_for(var, from, to, body)
128 }
129
130 /// Effectively-infinite loop used by persistent kernels (megakernel,
131 /// event loops, streaming). Lowers to `Node::Loop` with
132 /// `from: 0, to: u32::MAX`. At 1 µs per iteration `u32::MAX` is ~68
133 /// years - for all practical purposes infinite. The inner body
134 /// drives termination via `Node::Return` or by observing an
135 /// atomic shutdown flag the host sets.
136 ///
137 /// Linus principle: one enum variant (`Node::Loop`) handles both
138 /// bounded and persistent cases. No cascade of match arms through
139 /// every pass; no new wire-format tag. An optimizer pass
140 /// that wants to distinguish "truly unbounded" from "large bound"
141 /// inspects the `to` expression.
142 ///
143 /// # Examples
144 ///
145 /// ```
146 /// use vyre::ir::Node;
147 ///
148 /// let persistent = Node::forever(vec![Node::Return]);
149 /// assert!(matches!(persistent, Node::Loop { .. }));
150 /// ```
151 #[must_use]
152 #[inline]
153 pub fn forever(body: Vec<Self>) -> Self {
154 Self::loop_for("__forever__", Expr::u32(0), Expr::u32(u32::MAX), body)
155 }
156
157 /// Sequence of statements.
158 ///
159 /// # Examples
160 ///
161 /// ```
162 /// use vyre::ir::Node;
163 ///
164 /// assert!(matches!(Node::block(vec![Node::Return]), Node::Block(_)));
165 /// ```
166 #[must_use]
167 #[inline]
168 pub fn block(nodes: Vec<Self>) -> Self {
169 Self::Block(nodes)
170 }
171
172 /// Early return from the entry point.
173 ///
174 /// # Examples
175 ///
176 /// ```
177 /// use vyre::ir::Node;
178 ///
179 /// assert!(matches!(Node::return_(), Node::Return));
180 /// ```
181 #[must_use]
182 pub const fn return_() -> Self {
183 Self::Return
184 }
185
186 /// Workgroup barrier statement.
187 ///
188 /// # Examples
189 ///
190 /// ```
191 /// use vyre::ir::Node;
192 ///
193 /// assert!(matches!(Node::barrier(), Node::Barrier { .. }));
194 /// ```
195 #[must_use]
196 pub const fn barrier() -> Self {
197 Self::Barrier {
198 ordering: MemoryOrdering::SeqCst,
199 }
200 }
201
202 /// Workgroup barrier statement with explicit memory ordering.
203 #[must_use]
204 pub const fn barrier_with_ordering(ordering: MemoryOrdering) -> Self {
205 Self::Barrier { ordering }
206 }
207
208 /// Statement-level invocation of another registered op by stable
209 /// op id.
210 ///
211 /// Represented as a named [`Node::Region`] whose `generator` is
212 /// the callee's op id and whose body is an internal sequence of
213 /// `Node::Let { name: "arg{i}", value: <arg_expr> }` bindings.
214 /// Every backend already handles `Node::Region` - the op-registry
215 /// inliner walks the arg binds, substitutes them into the callee's
216 /// fragment, and splices the result in place. No new IR variant
217 /// is introduced, and the arg values remain fully visible to CSE,
218 /// DCE, and constant folding through the let-chain.
219 #[must_use]
220 pub fn call(op_id: impl Into<Ident>, args: Vec<Expr>) -> Self {
221 let body: Vec<Node> = args
222 .into_iter()
223 .enumerate()
224 .map(|(idx, expr)| Node::let_bind(format!("arg{idx}"), expr))
225 .collect();
226 Self::Region {
227 generator: op_id.into(),
228 source_region: None,
229 body: Arc::new(body),
230 }
231 }
232
233 /// Command-level indirect dispatch metadata.
234 ///
235 /// # Examples
236 ///
237 /// ```
238 /// use vyre::ir::Node;
239 ///
240 /// let node = Node::indirect_dispatch("counts", 0);
241 /// assert!(matches!(node, Node::IndirectDispatch { .. }));
242 /// ```
243 #[must_use]
244 #[inline]
245 pub fn indirect_dispatch(count_buffer: impl Into<Ident>, count_offset: u64) -> Self {
246 Self::IndirectDispatch {
247 count_buffer: count_buffer.into(),
248 count_offset,
249 }
250 }
251
252 /// Begin an asynchronous transfer stream region that names its own
253 /// operands, so the GPU drives the transfer.
254 ///
255 /// # Examples
256 ///
257 /// ```
258 /// use vyre::ir::{Node, Expr};
259 ///
260 /// let node = Node::async_load_gpu_driven("ssd", "vram", Expr::u32(0), Expr::u32(1024), "tag-0");
261 /// assert!(matches!(node, Node::AsyncLoad { .. }));
262 /// ```
263 #[must_use]
264 #[inline]
265 pub fn async_load_gpu_driven(
266 source: impl Into<Ident>,
267 destination: impl Into<Ident>,
268 offset: Expr,
269 size: Expr,
270 tag: impl Into<Ident>,
271 ) -> Self {
272 Self::AsyncLoad {
273 source: source.into(),
274 destination: destination.into(),
275 offset: Box::new(offset),
276 size: Box::new(size),
277 tag: tag.into(),
278 }
279 }
280
281 /// Begin an asynchronous transfer stream region carrying only a tag.
282 ///
283 /// The host drives the transfer and knows the buffers out of band, so the
284 /// node records placeholder operands. Emit
285 /// [`Node::async_load_gpu_driven`] when the source, destination, offset,
286 /// and size are known in the IR.
287 #[must_use]
288 #[inline]
289 pub fn async_load(tag: impl Into<Ident>) -> Self {
290 Self::async_load_gpu_driven(
291 "__legacy_src__",
292 "__legacy_dst__",
293 Expr::u32(0),
294 Expr::u32(0),
295 tag,
296 )
297 }
298
299 /// Begin an asynchronous store transfer stream region (GPU-driven).
300 #[must_use]
301 #[inline]
302 pub fn async_store(
303 source: impl Into<Ident>,
304 destination: impl Into<Ident>,
305 offset: Expr,
306 size: Expr,
307 tag: impl Into<Ident>,
308 ) -> Self {
309 Self::AsyncStore {
310 source: source.into(),
311 destination: destination.into(),
312 offset: Box::new(offset),
313 size: Box::new(size),
314 tag: tag.into(),
315 }
316 }
317
318 /// Wait for an asynchronous transfer stream region.
319 ///
320 /// # Examples
321 ///
322 /// ```
323 /// use vyre::ir::Node;
324 ///
325 /// let node = Node::async_wait("stage-a");
326 /// assert!(matches!(node, Node::AsyncWait { .. }));
327 /// ```
328 #[must_use]
329 #[inline]
330 pub fn async_wait(tag: impl Into<Ident>) -> Self {
331 Self::AsyncWait { tag: tag.into() }
332 }
333
334 /// Trap the current execution lane (GPU-initiated page fault).
335 #[must_use]
336 #[inline]
337 pub fn trap(address: Expr, tag: impl Into<Ident>) -> Self {
338 Self::Trap {
339 address: Box::new(address),
340 tag: tag.into(),
341 }
342 }
343
344 /// Resume a previously trapped execution lane.
345 #[must_use]
346 #[inline]
347 pub fn resume(tag: impl Into<Ident>) -> Self {
348 Self::Resume { tag: tag.into() }
349 }
350
351 /// Wrap a downstream extension statement node.
352 #[must_use]
353 #[inline]
354 pub fn opaque(node: impl NodeExtension) -> Self {
355 Self::Opaque(Arc::new(node))
356 }
357
358 /// Wrap a shared downstream extension statement node.
359 #[must_use]
360 #[inline]
361 pub fn opaque_arc(node: Arc<dyn NodeExtension>) -> Self {
362 Self::Opaque(node)
363 }
364}