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