Skip to main content

vyre_foundation/ir_inner/model/program/
core.rs

1use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering};
2use std::sync::{Arc, OnceLock};
3
4use rustc_hash::FxHashMap;
5
6use crate::ir_inner::model::node::Node;
7
8use super::BufferDecl;
9
10/// A complete vyre program.
11///
12/// Contains everything needed to execute a GPU compute dispatch:
13/// buffer declarations, workgroup configuration, and the entry point body.
14///
15/// # Example
16///
17/// A program that XORs two input buffers element-wise:
18///
19/// ```rust
20/// use vyre::ir::{Program, BufferDecl, BufferAccess, DataType, Node, Expr, BinOp};
21///
22/// let program = Program::wrapped(
23///     vec![
24///         BufferDecl::storage("a", 0, BufferAccess::ReadOnly, DataType::U32),
25///         BufferDecl::storage("b", 1, BufferAccess::ReadOnly, DataType::U32),
26///         BufferDecl::storage("out", 2, BufferAccess::ReadWrite, DataType::U32),
27///     ],
28///     [64, 1, 1],
29///     vec![
30///         Node::let_bind("idx", Expr::gid_x()),
31///         Node::if_then(
32///             Expr::lt(Expr::var("idx"), Expr::buf_len("out")),
33///             vec![
34///                 Node::store("out", Expr::var("idx"),
35///                     Expr::bitxor(
36///                         Expr::load("a", Expr::var("idx")),
37///                         Expr::load("b", Expr::var("idx")),
38///                     ),
39///                 ),
40///             ],
41///         ),
42///     ],
43/// );
44/// assert_eq!(program.buffers().len(), 3);
45/// ```
46#[derive(Debug)]
47pub struct Program {
48    /// Stable ID of the certified operation this program implements.
49    ///
50    /// Runtime lowering must reject programs without an ID because anonymous IR
51    /// cannot be tied back to a conform registry entry.
52    pub entry_op_id: Option<String>,
53    /// Buffer declarations. Each declares a named, typed, bound memory region.
54    pub buffers: Arc<[BufferDecl]>,
55    /// Sidecar index for O(1) buffer lookup by name.
56    pub(crate) buffer_index: Arc<FxHashMap<Arc<str>, usize>>,
57    /// Workgroup size: `[x, y, z]`. Controls `@workgroup_size` in target-text.
58    pub workgroup_size: [u32; 3],
59    /// Entry point body. Executes once per invocation.
60    pub entry: Arc<Vec<Node>>,
61    /// Cached blake3 hash of the program for fast equality and cache lookups.
62    pub(crate) hash: OnceLock<blake3::Hash>,
63    /// Per-backend validation cache (lazily initialized  -  most intermediate
64    /// programs created during fixpoint iteration are never validated, so the
65    /// sharded DashSet is only allocated on first `mark_validated_on` call).
66    #[doc(hidden)]
67    pub(crate) validation_set: OnceLock<Arc<dashmap::DashSet<Arc<str>>>>,
68    pub(crate) structural_validated: AtomicBool,
69    pub(crate) structural_validation_fingerprint: AtomicU64,
70    pub(crate) mutation_provenance: AtomicU8,
71    pub(crate) fingerprint: OnceLock<[u8; 32]>,
72    /// Memoized normalized compiled-pipeline cache digest.
73    ///
74    /// Same shape and lifecycle as `fingerprint`: a pure function of the
75    /// program value, so it lives on the value and is keyed by nothing, which
76    /// is why it structurally cannot serve another program's digest.
77    pub(crate) normalized_cache_digest: OnceLock<[u8; 32]>,
78    // VYRE_IR_HOTSPOTS HIGH (core.rs:100-117): both caches were
79    // plain values, so `Program::clone` copied the whole Vec / whole
80    // ProgramStats by value. Wrapping them in Arc turns the clone
81    // into a refcount bump, keeping Program::clone O(1) on every
82    // field.
83    pub(crate) output_buffer_index: OnceLock<Arc<Vec<u32>>>,
84    pub(crate) has_indirect_dispatch: OnceLock<bool>,
85    /// Cached statistics computed from a single walk of the program.
86    ///
87    /// This is a transient cache: it is not serialized to wire format and is
88    /// invalidated whenever the program shape mutates.
89    pub(crate) stats: OnceLock<Arc<super::ProgramStats>>,
90    /// When true, this program must not be fused with another copy of itself
91    /// in the same megakernel. Parser programs that use workgroup-local scratch
92    /// buffers set this to avoid state corruption when two invocations share
93    /// the same workgroup memory.
94    pub non_composable_with_self: bool,
95}
96
97impl Default for Program {
98    #[inline]
99    fn default() -> Self {
100        Self::empty()
101    }
102}
103
104impl Clone for Program {
105    fn clone(&self) -> Self {
106        let cloned = Self {
107            entry_op_id: self.entry_op_id.clone(),
108            buffers: Arc::clone(&self.buffers),
109            buffer_index: Arc::clone(&self.buffer_index),
110            workgroup_size: self.workgroup_size,
111            entry: Arc::clone(&self.entry),
112            hash: OnceLock::new(),
113            validation_set: {
114                let cell = OnceLock::new();
115                if let Some(set) = self.validation_set.get() {
116                    let _ = cell.set(Arc::clone(set));
117                }
118                cell
119            },
120            structural_validated: AtomicBool::new(self.is_structurally_validated()),
121            structural_validation_fingerprint: AtomicU64::new(
122                self.structural_validation_fingerprint
123                    .load(Ordering::Acquire),
124            ),
125            mutation_provenance: AtomicU8::new(self.mutation_provenance.load(Ordering::Acquire)),
126            fingerprint: OnceLock::new(),
127            normalized_cache_digest: OnceLock::new(),
128            output_buffer_index: OnceLock::new(),
129            has_indirect_dispatch: OnceLock::new(),
130            stats: OnceLock::new(),
131            non_composable_with_self: self.non_composable_with_self,
132        };
133        // Each OnceLock above was just initialised with `OnceLock::new()`,
134        // so the matching `.set(...)` is infallible by construction  -
135        // `let _ = ...set(...)` matches the same pattern used for the
136        // validation_set initialiser above and avoids the raw-expect
137        // CI gate without introducing fake error paths.
138        if let Some(hash) = self.hash.get() {
139            let _ = cloned.hash.set(*hash);
140        }
141        if let Some(fingerprint) = self.fingerprint.get() {
142            let _ = cloned.fingerprint.set(*fingerprint);
143        }
144        if let Some(normalized_cache_digest) = self.normalized_cache_digest.get() {
145            let _ = cloned.normalized_cache_digest.set(*normalized_cache_digest);
146        }
147        if let Some(output_buffer_index) = self.output_buffer_index.get() {
148            // Arc::clone = refcount bump, no Vec<u32> copy.
149            let _ = cloned
150                .output_buffer_index
151                .set(Arc::clone(output_buffer_index));
152        }
153        if let Some(has_indirect_dispatch) = self.has_indirect_dispatch.get() {
154            let _ = cloned.has_indirect_dispatch.set(*has_indirect_dispatch);
155        }
156        if let Some(stats) = self.stats.get() {
157            // Arc::clone = refcount bump, no ProgramStats copy.
158            let _ = cloned.stats.set(Arc::clone(stats));
159        }
160        cloned
161    }
162}
163
164impl PartialEq for Program {
165    fn eq(&self, other: &Self) -> bool {
166        self.structural_eq(other)
167    }
168}
169
170impl Eq for Program {}