vyre_foundation/ir_inner/model/program/
builder.rs

1use std::sync::{Arc, OnceLock};
2
3use rustc_hash::FxHashMap;
4
5use crate::ir_inner::model::arena::{ArenaProgram, ExprArena};
6use crate::ir_inner::model::node::Node;
7
8use super::{BufferDecl, Program};
9
10macro_rules! define_raw_program_constructor {
11    () => {
12        /// Create a program that preserves raw top-level entry nodes.
13        ///
14        /// This constructor is reserved for wire decoding, reference adapters,
15        /// and negative validation tests. Use [`Program::wrapped`] for runnable
16        /// programs.
17        #[must_use]
18        #[inline]
19        pub fn from_raw_parts(
20            buffers: Vec<BufferDecl>,
21            workgroup_size: [u32; 3],
22            entry: Vec<Node>,
23        ) -> Self {
24            Self::new_raw(buffers, workgroup_size, entry)
25        }
26    };
27}
28
29impl Program {
30    /// Synthetic generator id used when callers submit a raw top-level body
31    /// instead of an explicit `Node::Region`.
32    pub const ROOT_REGION_GENERATOR: &'static str = "vyre.program.root";
33
34    /// Create a complete program from buffer declarations, workgroup size, and
35    /// entry-point nodes, auto-wrapping the top-level body in a root Region
36    /// when necessary.
37    ///
38    /// This is the default construction path for runnable Programs.
39    ///
40    /// # Examples
41    ///
42    /// ```
43    /// use vyre::ir::{BufferAccess, BufferDecl, DataType, Node, Program};
44    ///
45    /// let program = Program::wrapped(
46    ///     vec![BufferDecl::storage(
47    ///         "output",
48    ///         0,
49    ///         BufferAccess::ReadWrite,
50    ///         DataType::U32,
51    ///     )],
52    ///     [64, 1, 1],
53    ///     Vec::new(),
54    /// );
55    ///
56    /// assert_eq!(program.workgroup_size(), [64, 1, 1]);
57    /// assert_eq!(program.buffers().len(), 1);
58    /// assert!(matches!(program.entry(), [Node::Region { .. }]));
59    /// ```
60    ///
61    /// This constructs a NEW program. Do not use it to rebuild an existing one:
62    /// it starts the metadata from scratch, so `entry_op_id` and
63    /// `non_composable_with_self` are silently lost, and it deep-clones the
64    /// buffer table and re-interns every name. To change part of a program, use
65    /// [`Self::with_rewritten_buffers`], [`Self::with_rewritten_entry`],
66    /// [`Self::with_rewritten_wrapped_entry`], or [`Self::map_entry`], all of
67    /// which preserve the rest.
68    #[must_use]
69    #[inline]
70    pub fn wrapped(buffers: Vec<BufferDecl>, workgroup_size: [u32; 3], entry: Vec<Node>) -> Self {
71        Self::new_raw(buffers, workgroup_size, Self::wrap_entry(entry))
72    }
73
74    define_raw_program_constructor!();
75
76    #[must_use]
77    #[inline]
78    pub(crate) fn new_raw(
79        buffers: Vec<BufferDecl>,
80        workgroup_size: [u32; 3],
81        entry: Vec<Node>,
82    ) -> Self {
83        let mut interner = FxHashMap::<Arc<str>, Arc<str>>::default();
84        interner.reserve(buffers.len());
85        let buffers: Vec<BufferDecl> = buffers
86            .into_iter()
87            .map(|mut b| {
88                let arc = interner
89                    .entry(Arc::clone(&b.name))
90                    .or_insert_with(|| Arc::clone(&b.name))
91                    .clone();
92                b.name = arc;
93                b
94            })
95            .collect();
96        let buffer_index = Self::build_buffer_index(&buffers);
97        Self {
98            entry_op_id: None,
99            buffers: Arc::from(buffers),
100            buffer_index: Arc::new(buffer_index),
101            workgroup_size,
102            entry: Arc::new(entry),
103            hash: OnceLock::new(),
104            validation_set: OnceLock::new(),
105            structural_validated: std::sync::atomic::AtomicBool::new(false),
106            structural_validation_fingerprint: std::sync::atomic::AtomicU64::new(0),
107            mutation_provenance: std::sync::atomic::AtomicU8::new(0),
108            fingerprint: OnceLock::new(),
109            normalized_cache_digest: OnceLock::new(),
110            output_buffer_index: OnceLock::new(),
111            has_indirect_dispatch: OnceLock::new(),
112            stats: OnceLock::new(),
113            non_composable_with_self: false,
114        }
115    }
116
117    /// Same as [`Self::with_rewritten_entry`] but wraps the entry first via
118    /// the runnable-Region root contract (matches [`Self::wrapped`]). Use
119    /// from passes that produce a fully fresh entry body but want to reuse
120    /// the existing buffer Arc instead of paying for a full
121    /// [`Self::wrapped`] (which deep-clones buffers, re-interns names, and
122    /// rebuilds the buffer index).
123    #[must_use]
124    #[inline]
125    pub fn with_rewritten_wrapped_entry(&self, entry: Vec<Node>) -> Self {
126        self.with_rewritten_entry(Self::wrap_entry(entry))
127    }
128
129    /// Consume this program and rebuild it with `f` applied to the owned
130    /// entry vec. Reuses the entry Arc when uniquely owned (the common
131    /// case under the optimizer fixpoint)  -  no deep clone of the entry
132    /// body and no scaffold allocation. Equivalent to:
133    ///
134    /// ```ignore
135    /// let scaffold = program.with_rewritten_entry(Vec::new());
136    /// let entry = f(program.into_entry_vec());
137    /// scaffold.with_rewritten_entry(entry)
138    /// ```
139    ///
140    /// but produces only one new `Program` value instead of two.
141    #[must_use]
142    #[inline]
143    pub fn map_entry<F: FnOnce(Vec<Node>) -> Vec<Node>>(self, f: F) -> Self {
144        let entry_op_id = self.entry_op_id.clone();
145        let buffers = Arc::clone(&self.buffers);
146        let buffer_index = Arc::clone(&self.buffer_index);
147        let workgroup_size = self.workgroup_size;
148        let non_composable_with_self = self.non_composable_with_self;
149        let entry = f(self.into_entry_vec());
150        Self {
151            entry_op_id,
152            buffers,
153            buffer_index,
154            workgroup_size,
155            entry: Arc::new(entry),
156            hash: OnceLock::new(),
157            validation_set: OnceLock::new(),
158            structural_validated: std::sync::atomic::AtomicBool::new(false),
159            structural_validation_fingerprint: std::sync::atomic::AtomicU64::new(0),
160            mutation_provenance: std::sync::atomic::AtomicU8::new(0),
161            fingerprint: OnceLock::new(),
162            normalized_cache_digest: OnceLock::new(),
163            output_buffer_index: OnceLock::new(),
164            has_indirect_dispatch: OnceLock::new(),
165            stats: OnceLock::new(),
166            non_composable_with_self,
167        }
168    }
169
170    /// Clone this program with a replacement entry body while preserving the
171    /// existing buffer table, workgroup size, and optional certified op id.
172    #[must_use]
173    #[inline]
174    pub fn with_rewritten_entry(&self, entry: Vec<Node>) -> Self {
175        self.with_rewritten_workgroup_size_and_entry(self.workgroup_size, entry)
176    }
177
178    /// Clone this program with replacement buffer declarations while
179    /// preserving the entry body, workgroup size, and metadata flags.
180    #[must_use]
181    #[inline]
182    pub fn with_rewritten_buffers(&self, buffers: Vec<BufferDecl>) -> Self {
183        let buffer_index = Self::build_buffer_index(&buffers);
184        Self {
185            entry_op_id: self.entry_op_id.clone(),
186            buffers: Arc::from(buffers),
187            buffer_index: Arc::new(buffer_index),
188            workgroup_size: self.workgroup_size,
189            entry: Arc::clone(&self.entry),
190            hash: OnceLock::new(),
191            validation_set: OnceLock::new(),
192            structural_validated: std::sync::atomic::AtomicBool::new(false),
193            structural_validation_fingerprint: std::sync::atomic::AtomicU64::new(0),
194            mutation_provenance: std::sync::atomic::AtomicU8::new(0),
195            fingerprint: OnceLock::new(),
196            normalized_cache_digest: OnceLock::new(),
197            output_buffer_index: OnceLock::new(),
198            has_indirect_dispatch: OnceLock::new(),
199            stats: OnceLock::new(),
200            non_composable_with_self: self.non_composable_with_self,
201        }
202    }
203
204    /// Clone this program with replacement dispatch dimensions and entry body
205    /// while preserving the existing buffer table, indexes, and metadata flags.
206    #[must_use]
207    #[inline]
208    pub fn with_rewritten_workgroup_size_and_entry(
209        &self,
210        workgroup_size: [u32; 3],
211        entry: Vec<Node>,
212    ) -> Self {
213        Self {
214            entry_op_id: self.entry_op_id.clone(),
215            buffers: Arc::clone(&self.buffers),
216            buffer_index: Arc::clone(&self.buffer_index),
217            workgroup_size,
218            entry: Arc::new(entry),
219            hash: OnceLock::new(),
220            validation_set: OnceLock::new(),
221            structural_validated: std::sync::atomic::AtomicBool::new(false),
222            structural_validation_fingerprint: std::sync::atomic::AtomicU64::new(0),
223            mutation_provenance: std::sync::atomic::AtomicU8::new(0),
224            fingerprint: OnceLock::new(),
225            normalized_cache_digest: OnceLock::new(),
226            output_buffer_index: OnceLock::new(),
227            has_indirect_dispatch: OnceLock::new(),
228            stats: OnceLock::new(),
229            non_composable_with_self: self.non_composable_with_self,
230        }
231    }
232
233    /// Consume the program and return its entry nodes, reusing the
234    /// backing vector when this program owns the entry body uniquely.
235    #[must_use]
236    #[inline]
237    pub fn into_entry_vec(self) -> Vec<Node> {
238        Arc::try_unwrap(self.entry).unwrap_or_else(|entry| entry.as_ref().clone())
239    }
240
241    /// Create an arena-backed program scaffold.
242    ///
243    /// This constructor is the opt-in path for builders that want
244    /// [`ExprRef`](crate::ir_inner::model::arena::ExprRef) handles instead of boxed
245    /// expression trees. [`Program::wrapped`] remains the boxed-tree constructor.
246    #[must_use]
247    #[inline]
248    pub fn with_arena(
249        arena: &ExprArena,
250        buffers: Vec<BufferDecl>,
251        workgroup_size: [u32; 3],
252    ) -> ArenaProgram<'_> {
253        ArenaProgram::new(arena, buffers, workgroup_size)
254    }
255
256    /// Create a minimal program with no buffers and an empty body.
257    ///
258    /// # Examples
259    ///
260    /// ```
261    /// use vyre::ir::Program;
262    ///
263    /// let program = Program::empty();
264    ///
265    /// assert!(program.buffers().is_empty());
266    /// assert_eq!(program.workgroup_size(), [1, 1, 1]);
267    /// assert!(program.is_explicit_noop());
268    /// ```
269    #[must_use]
270    #[inline]
271    pub fn empty() -> Self {
272        Self::wrapped(Vec::new(), [1, 1, 1], Vec::new())
273    }
274
275    /// Attach the stable operation ID whose conform registry entry certifies
276    /// this program for runtime lowering.
277    #[must_use]
278    #[inline]
279    pub fn with_entry_op_id(mut self, op_id: impl Into<String>) -> Self {
280        self.entry_op_id = Some(op_id.into());
281        self.invalidate_caches();
282        self
283    }
284
285    /// Stable operation ID required by the conform gate.
286    #[must_use]
287    #[inline]
288    pub fn entry_op_id(&self) -> Option<&str> {
289        self.entry_op_id.as_deref()
290    }
291
292    /// Attach an optional operation ID while preserving anonymous test IR.
293    #[must_use]
294    #[inline]
295    pub(crate) fn with_optional_entry_op_id(mut self, op_id: Option<String>) -> Self {
296        self.entry_op_id = op_id;
297        self.invalidate_caches();
298        self
299    }
300}
vyre_foundation/ir_inner/model/program/builder.rs

vyre_foundation/ir_inner/model/program/
builder.rs
