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

vyre_foundation/ir_inner/model/program/
builder.rs
