Skip to main content

vyre_foundation/ir_inner/model/program/
definition.rs

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