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 {}