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 (GPU-driven).
253 ///
254 /// # Examples
255 ///
256 /// ```
257 /// use vyre::ir::{Node, Expr};
258 ///
259 /// let node = Node::async_load_ext("ssd", "vram", Expr::u32(0), Expr::u32(1024), "tag-0");
260 /// assert!(matches!(node, Node::AsyncLoad { .. }));
261 /// ```
262 #[must_use]
263 #[inline]
264 pub fn async_load_ext(
265 source: impl Into<Ident>,
266 destination: impl Into<Ident>,
267 offset: Expr,
268 size: Expr,
269 tag: impl Into<Ident>,
270 ) -> Self {
271 Self::AsyncLoad {
272 source: source.into(),
273 destination: destination.into(),
274 offset: Box::new(offset),
275 size: Box::new(size),
276 tag: tag.into(),
277 }
278 }
279
280 /// Begin an asynchronous transfer stream region (legacy/host-driven).
281 #[must_use]
282 #[inline]
283 pub fn async_load(tag: impl Into<Ident>) -> Self {
284 Self::async_load_ext(
285 "__legacy_src__",
286 "__legacy_dst__",
287 Expr::u32(0),
288 Expr::u32(0),
289 tag,
290 )
291 }
292
293 /// Begin an asynchronous store transfer stream region (GPU-driven).
294 #[must_use]
295 #[inline]
296 pub fn async_store(
297 source: impl Into<Ident>,
298 destination: impl Into<Ident>,
299 offset: Expr,
300 size: Expr,
301 tag: impl Into<Ident>,
302 ) -> Self {
303 Self::AsyncStore {
304 source: source.into(),
305 destination: destination.into(),
306 offset: Box::new(offset),
307 size: Box::new(size),
308 tag: tag.into(),
309 }
310 }
311
312 /// Wait for an asynchronous transfer stream region.
313 ///
314 /// # Examples
315 ///
316 /// ```
317 /// use vyre::ir::Node;
318 ///
319 /// let node = Node::async_wait("stage-a");
320 /// assert!(matches!(node, Node::AsyncWait { .. }));
321 /// ```
322 #[must_use]
323 #[inline]
324 pub fn async_wait(tag: impl Into<Ident>) -> Self {
325 Self::AsyncWait { tag: tag.into() }
326 }
327
328 /// Trap the current execution lane (GPU-initiated page fault).
329 #[must_use]
330 #[inline]
331 pub fn trap(address: Expr, tag: impl Into<Ident>) -> Self {
332 Self::Trap {
333 address: Box::new(address),
334 tag: tag.into(),
335 }
336 }
337
338 /// Resume a previously trapped execution lane.
339 #[must_use]
340 #[inline]
341 pub fn resume(tag: impl Into<Ident>) -> Self {
342 Self::Resume { tag: tag.into() }
343 }
344
345 /// Wrap a downstream extension statement node.
346 #[must_use]
347 #[inline]
348 pub fn opaque(node: impl NodeExtension) -> Self {
349 Self::Opaque(Arc::new(node))
350 }
351
352 /// Wrap a shared downstream extension statement node.
353 #[must_use]
354 #[inline]
355 pub fn opaque_arc(node: Arc<dyn NodeExtension>) -> Self {
356 Self::Opaque(node)
357 }
358}