vyre_foundation/ir_inner/model/program/
meta.rs

1use std::hash::{Hash, Hasher as _};
2use std::sync::atomic::Ordering;
3use std::sync::Arc;
4
5use rustc_hash::FxHasher;
6use vyre_spec::bin_op::OpIntensity;
7
8use crate::ir::{Expr, Node};
9use crate::ir_inner::model::expr::Ident;
10use crate::ir_inner::model::spec_types::BufferAccess;
11use crate::transform::visit::{walk_nodes_and_exprs, ExprVisitor, NodeVisitor};
12
13use super::Program;
14
15/// Provenance for mutations that invalidate Program validation/cache state.
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17#[repr(u8)]
18pub enum ProgramMutationProvenance {
19    /// Program has not been mutated since construction or successful validation.
20    Clean = 0,
21    /// The non-composable dispatch flag changed.
22    NonComposableFlag = 1,
23    /// The target workgroup dimensions changed.
24    WorkgroupSize = 2,
25    /// The substrate-neutral parallel-region dimensions changed.
26    ParallelRegionSize = 3,
27    /// A caller borrowed the mutable entry vector.
28    EntryMutation = 4,
29    /// Internal builder or decode path rewrote Program shape.
30    InternalShapeMutation = 5,
31    /// Mutation provenance is unknown, so validation must fail closed.
32    Unknown = 255,
33}
34
35impl ProgramMutationProvenance {
36    #[inline]
37    const fn from_code(code: u8) -> Self {
38        match code {
39            0 => Self::Clean,
40            1 => Self::NonComposableFlag,
41            2 => Self::WorkgroupSize,
42            3 => Self::ParallelRegionSize,
43            4 => Self::EntryMutation,
44            5 => Self::InternalShapeMutation,
45            _ => Self::Unknown,
46        }
47    }
48}
49
50fn mix_wire_fallback_hashable<T: Hash>(hasher: &mut blake3::Hasher, value: &T) {
51    let mut state = FxHasher::default();
52    value.hash(&mut state);
53    hasher.update(&state.finish().to_le_bytes());
54}
55
56/// Bounded IR structure digest for wire-hash fallback (never formats full IR via `Debug`).
57struct FallbackWireHasher<'a>(&'a mut blake3::Hasher);
58
59impl NodeVisitor for FallbackWireHasher<'_> {
60    fn visit_node(&mut self, node: &Node) {
61        let h = &mut *self.0;
62        match node {
63            Node::Let { name, .. } => {
64                h.update(b"n:Let\0");
65                h.update(name.as_bytes());
66            }
67            Node::Assign { name, .. } => {
68                h.update(b"n:Assign\0");
69                h.update(name.as_bytes());
70            }
71            Node::Store { buffer, .. } => {
72                h.update(b"n:Store\0");
73                h.update(buffer.as_bytes());
74            }
75            Node::If { .. } => {
76                h.update(b"n:If\0");
77            }
78            Node::Loop { var, .. } => {
79                h.update(b"n:Loop\0");
80                h.update(var.as_bytes());
81            }
82            Node::IndirectDispatch {
83                count_buffer,
84                count_offset,
85            } => {
86                h.update(b"n:IndirectDispatch\0");
87                h.update(count_buffer.as_bytes());
88                h.update(&count_offset.to_le_bytes());
89            }
90            Node::AsyncLoad {
91                source,
92                destination,
93                tag,
94                ..
95            } => {
96                h.update(b"n:AsyncLoad\0");
97                h.update(source.as_bytes());
98                h.update(destination.as_bytes());
99                h.update(tag.as_bytes());
100            }
101            Node::AsyncStore {
102                source,
103                destination,
104                tag,
105                ..
106            } => {
107                h.update(b"n:AsyncStore\0");
108                h.update(source.as_bytes());
109                h.update(destination.as_bytes());
110                h.update(tag.as_bytes());
111            }
112            Node::AsyncWait { tag } => {
113                h.update(b"n:AsyncWait\0");
114                h.update(tag.as_bytes());
115            }
116            Node::Trap { tag, .. } => {
117                h.update(b"n:Trap\0");
118                h.update(tag.as_bytes());
119            }
120            Node::Resume { tag } => {
121                h.update(b"n:Resume\0");
122                h.update(tag.as_bytes());
123            }
124            Node::AllReduce { buffer, op, group } => {
125                h.update(b"n:AllReduce\0");
126                h.update(buffer.as_bytes());
127                h.update(&op.builtin_wire_tag().to_le_bytes());
128                h.update(&group.as_u32().to_le_bytes());
129            }
130            Node::AllGather {
131                input,
132                output,
133                group,
134            } => {
135                h.update(b"n:AllGather\0");
136                h.update(input.as_bytes());
137                h.update(output.as_bytes());
138                h.update(&group.as_u32().to_le_bytes());
139            }
140            Node::ReduceScatter {
141                input,
142                output,
143                op,
144                group,
145            } => {
146                h.update(b"n:ReduceScatter\0");
147                h.update(input.as_bytes());
148                h.update(output.as_bytes());
149                h.update(&op.builtin_wire_tag().to_le_bytes());
150                h.update(&group.as_u32().to_le_bytes());
151            }
152            Node::Broadcast {
153                buffer,
154                root,
155                group,
156            } => {
157                h.update(b"n:Broadcast\0");
158                h.update(buffer.as_bytes());
159                h.update(&root.to_le_bytes());
160                h.update(&group.as_u32().to_le_bytes());
161            }
162            Node::Return => {
163                h.update(b"n:Return\0");
164            }
165            Node::Barrier { ordering } => {
166                h.update(b"n:Barrier\0");
167                mix_wire_fallback_hashable(h, ordering);
168            }
169            Node::Block(_) => {
170                h.update(b"n:Block\0");
171            }
172            Node::Region {
173                generator,
174                source_region,
175                ..
176            } => {
177                h.update(b"n:Region\0");
178                h.update(generator.as_bytes());
179                if let Some(source_gen) = source_region {
180                    h.update(source_gen.name.as_bytes());
181                }
182            }
183            Node::Opaque(ext) => {
184                h.update(b"n:Opaque\0");
185                h.update(ext.extension_kind().as_bytes());
186            }
187        }
188    }
189}
190
191impl ExprVisitor for FallbackWireHasher<'_> {
192    fn visit_expr(&mut self, expr: &Expr) {
193        let h = &mut *self.0;
194        match expr {
195            Expr::LitU32(v) => {
196                h.update(b"e:LitU32\0");
197                h.update(&v.to_le_bytes());
198            }
199            Expr::LitI32(v) => {
200                h.update(b"e:LitI32\0");
201                h.update(&v.to_le_bytes());
202            }
203            Expr::LitF32(v) => {
204                h.update(b"e:LitF32\0");
205                h.update(&v.to_le_bytes());
206            }
207            Expr::LitBool(v) => {
208                h.update(b"e:LitBool\0");
209                h.update(&[u8::from(*v)]);
210            }
211            Expr::Var(name) => {
212                h.update(b"e:Var\0");
213                h.update(name.as_bytes());
214            }
215            Expr::Load { buffer, .. } => {
216                h.update(b"e:Load\0");
217                h.update(buffer.as_bytes());
218            }
219            Expr::BufLen { buffer } => {
220                h.update(b"e:BufLen\0");
221                h.update(buffer.as_bytes());
222            }
223            Expr::BufferRef { buffer } => {
224                h.update(b"e:BufferRef\0");
225                h.update(buffer.as_bytes());
226            }
227            Expr::InvocationId { axis } => {
228                h.update(b"e:InvocationId\0");
229                h.update(&[*axis]);
230            }
231            Expr::WorkgroupId { axis } => {
232                h.update(b"e:WorkgroupId\0");
233                h.update(&[*axis]);
234            }
235            Expr::LocalId { axis } => {
236                h.update(b"e:LocalId\0");
237                h.update(&[*axis]);
238            }
239            Expr::BinOp { op, .. } => {
240                h.update(b"e:BinOp\0");
241                mix_wire_fallback_hashable(h, op);
242            }
243            Expr::UnOp { op, .. } => {
244                h.update(b"e:UnOp\0");
245                mix_wire_fallback_hashable(h, op);
246            }
247            Expr::Call { op_id, .. } => {
248                h.update(b"e:Call\0");
249                h.update(op_id.as_bytes());
250            }
251            Expr::Select { .. } => {
252                h.update(b"e:Select\0");
253            }
254            Expr::Cast { target, .. } => {
255                h.update(b"e:Cast\0");
256                mix_wire_fallback_hashable(h, target);
257            }
258            Expr::Fma { .. } => {
259                h.update(b"e:Fma\0");
260            }
261            Expr::Atomic {
262                op,
263                buffer,
264                ordering,
265                ..
266            } => {
267                h.update(b"e:Atomic\0");
268                mix_wire_fallback_hashable(h, op);
269                h.update(buffer.as_bytes());
270                mix_wire_fallback_hashable(h, ordering);
271            }
272            Expr::SubgroupBallot { .. } => {
273                h.update(b"e:SubgroupBallot\0");
274            }
275            Expr::SubgroupShuffle { .. } => {
276                h.update(b"e:SubgroupShuffle\0");
277            }
278            Expr::SubgroupReduce { op, .. } => {
279                h.update(b"e:SubgroupReduce\0");
280                h.update(&[op.builtin_wire_tag()]);
281            }
282            Expr::SubgroupLocalId => {
283                h.update(b"e:SubgroupLocalId\0");
284            }
285            Expr::SubgroupSize => {
286                h.update(b"e:SubgroupSize\0");
287            }
288            Expr::Opaque(ext) => {
289                h.update(b"e:Opaque\0");
290                h.update(ext.extension_kind().as_bytes());
291            }
292        }
293    }
294}
295
296impl Program {
297    /// Re-apply the same top-level `Node::Region` contract as
298    /// [`Program::wrapped`].
299    ///
300    /// The [`region_inline_engine`](crate::optimizer::passes::cleanup::region_inline_engine)
301    /// pass flattens small Category-A regions so CSE/DCE can see a single
302    /// function-shaped body, which can leave a statement-shaped entry list. The
303    /// standard optimizer run ends with this helper so the program remains in
304    /// a runnable, validator/reference-interpreter–compatible form while
305    /// still benefiting from the inline pass.
306    #[must_use]
307    pub fn reconcile_runnable_top_level(self) -> Self {
308        if self.is_top_level_region_wrapped() {
309            return self;
310        }
311        // Move the entry Vec out via map_entry's Arc-aware path; one
312        // Program rebuild instead of two scaffold rebuilds.
313        self.map_entry(Self::wrap_entry)
314    }
315
316    /// Look up a buffer declaration by name.
317    #[must_use]
318    #[inline]
319    pub fn buffer(&self, name: &str) -> Option<&super::BufferDecl> {
320        self.buffer_index
321            .get(name)
322            .and_then(|&index| self.buffers.get(index))
323    }
324
325    /// Declared buffers.
326    #[must_use]
327    #[inline]
328    pub fn buffers(&self) -> &[super::BufferDecl] {
329        self.buffers.as_ref()
330    }
331
332    /// Access the buffer declaration Arc directly for identity checks.
333    #[must_use]
334    #[inline]
335    #[cfg(test)]
336    pub(crate) fn buffers_arc(&self) -> &Arc<[super::BufferDecl]> {
337        &self.buffers
338    }
339
340    /// Compare two programs by observable IR structure.
341    ///
342    /// This walk intentionally ignores buffer declaration order and never
343    /// consults arena-local allocation identity. Two programs are structurally
344    /// equal when they declare the same buffers, workgroup size, optional entry
345    /// op id, and entry body semantics.
346    #[must_use]
347    #[inline]
348    pub fn structural_eq(&self, other: &Self) -> bool {
349        // Identity short-circuit: Program::clone shares all the
350        // inner Arcs, so comparing a cloned program against its
351        // source (the common optimizer-pipeline pattern) is pure
352        // refcount comparison.
353        if std::ptr::eq(self, other)
354            || (Arc::ptr_eq(&self.buffers, &other.buffers)
355                && Arc::ptr_eq(&self.entry, &other.entry)
356                && self.entry_op_id == other.entry_op_id
357                && self.non_composable_with_self == other.non_composable_with_self
358                && self.workgroup_size == other.workgroup_size)
359        {
360            return true;
361        }
362        self.entry_op_id == other.entry_op_id
363            && self.non_composable_with_self == other.non_composable_with_self
364            && buffers_equal_ignoring_declaration_order(&self.buffers, &other.buffers)
365            && self.workgroup_size == other.workgroup_size
366            && self.entry == other.entry
367    }
368
369    /// Workgroup dimensions.
370    #[must_use]
371    #[inline]
372    pub fn workgroup_size(&self) -> [u32; 3] {
373        self.workgroup_size
374    }
375
376    /// Substrate-neutral alias for [`workgroup_size`](Self::workgroup_size).
377    ///
378    /// Naming: "parallel region" avoids picking a single target substrate's
379    /// word for one dispatch invocation grouping.
380    #[must_use]
381    #[inline]
382    pub fn parallel_region_size(&self) -> [u32; 3] {
383        self.workgroup_size
384    }
385
386    /// Return true when this program must not be fused with another copy
387    /// of itself in the same megakernel.
388    #[must_use]
389    #[inline]
390    pub fn is_non_composable_with_self(&self) -> bool {
391        self.non_composable_with_self
392    }
393
394    /// Mark this program as non-composable with itself.
395    #[must_use]
396    #[inline]
397    pub fn with_non_composable_with_self(mut self, flag: bool) -> Self {
398        self.non_composable_with_self = flag;
399        self.invalidate_caches_for(ProgramMutationProvenance::NonComposableFlag);
400        self
401    }
402
403    /// Set the workgroup dimensions in place. Used by harnesses that
404    /// need to clone-and-rewrite a program's workgroup size for fallback
405    /// dispatch  -  the alternative was to reconstruct the entire Program,
406    /// which is unnecessarily expensive when only one field changes.
407    #[inline]
408    pub fn set_workgroup_size(&mut self, workgroup_size: [u32; 3]) {
409        self.workgroup_size = workgroup_size;
410        self.invalidate_caches_for(ProgramMutationProvenance::WorkgroupSize);
411    }
412
413    /// Substrate-neutral alias for [`set_workgroup_size`](Self::set_workgroup_size).
414    #[inline]
415    pub fn set_parallel_region_size(&mut self, parallel_region_size: [u32; 3]) {
416        self.workgroup_size = parallel_region_size;
417        self.invalidate_caches_for(ProgramMutationProvenance::ParallelRegionSize);
418    }
419
420    /// Entry-point nodes.
421    #[must_use]
422    #[inline]
423    pub fn entry(&self) -> &[Node] {
424        self.entry.as_ref().as_slice()
425    }
426
427    /// Shared entry-point body Arc for identity checks.
428    #[must_use]
429    #[inline]
430    pub fn entry_arc(&self) -> &Arc<Vec<Node>> {
431        &self.entry
432    }
433
434    /// Return true when this Program is the canonical no-op shape produced by
435    /// [`Program::empty`]: no buffers and a single empty root Region.
436    #[must_use]
437    #[inline]
438    pub fn is_explicit_noop(&self) -> bool {
439        self.buffers().is_empty()
440            && matches!(self.entry(), [Node::Region { body, .. }] if body.is_empty())
441    }
442
443    /// Return true when the program satisfies the top-level region-chain
444    /// invariant: at least one top-level node, and every top-level node is a
445    /// `Node::Region`.
446    #[must_use]
447    #[inline]
448    pub fn is_top_level_region_wrapped(&self) -> bool {
449        !self.entry.is_empty()
450            && self
451                .entry()
452                .iter()
453                .all(|node| matches!(node, Node::Region { .. }))
454    }
455
456    /// Structural cause for a violated top-level Region invariant.
457    #[must_use]
458    pub fn top_level_region_violation_cause(&self) -> Option<String> {
459        if self.entry().is_empty() {
460            return Some("program entry has no top-level Region".to_string());
461        }
462
463        self.entry()
464            .iter()
465            .enumerate()
466            .find(|(_, node)| !matches!(node, Node::Region { .. }))
467            .map(|(index, node)| {
468                format!(
469                    "program entry node {index} is `{}` instead of `Node::Region`",
470                    crate::ir_inner::model::node::node_op_id(node)
471                )
472            })
473    }
474
475    /// Actionable error text describing why the top-level region invariant
476    /// failed, or `None` when the entry is valid.
477    #[must_use]
478    pub fn top_level_region_violation(&self) -> Option<String> {
479        self.top_level_region_violation_cause().map(|cause| {
480            format!(
481                "{cause}. Fix: construct runnable programs with Program::wrapped(...) or wrap the body in Node::Region before validation, interpretation, or dispatch."
482            )
483        })
484    }
485
486    /// Mutable entry-point nodes for transformation passes.
487    #[must_use]
488    #[inline]
489    pub fn entry_mut(&mut self) -> &mut Vec<Node> {
490        self.invalidate_caches_for(ProgramMutationProvenance::EntryMutation);
491        Arc::make_mut(&mut self.entry)
492    }
493
494    /// Stable BLAKE3 fingerprint of the canonical wire-format bytes.
495    #[must_use]
496    #[inline]
497    pub fn fingerprint(&self) -> [u8; 32] {
498        *self.fingerprint.get_or_init(|| {
499            let hash = self.compute_wire_hash();
500            let _ = self.hash.set(hash);
501            *hash.as_bytes()
502        })
503    }
504
505    /// VSA-style hypervector fingerprint of the canonical wire-format
506    /// bytes. Each `u32` lane is one segment of the program's blake3
507    /// hash; together they form an 8-lane hypervector suitable for
508    /// approximate similarity search via hamming distance.
509    ///
510    /// Use as the canonical cache key for approximate-match caches
511    /// (e.g. validation cache, AOT artifact dedup); use
512    /// [`Self::fingerprint`] for exact-match lookups.
513    ///
514    /// Wires the substrate's #29 hypervector primitive into Program
515    /// itself  -  every Program now carries its own VSA fingerprint
516    /// without callers having to reach into the substrate explicitly.
517    #[must_use]
518    pub fn vsa_fingerprint(&self) -> Vec<u32> {
519        self.fingerprint()
520            .chunks_exact(core::mem::size_of::<u32>())
521            .map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
522            .collect()
523    }
524
525    /// Indices of writable storage outputs in `buffers()` order.
526    #[must_use]
527    #[inline]
528    pub fn output_buffer_indices(&self) -> &[u32] {
529        self.output_buffer_index
530            .get_or_init(|| {
531                Arc::new(
532                    self.buffers()
533                        .iter()
534                        .enumerate()
535                        .filter_map(|(index, buffer)| {
536                            matches!(
537                                buffer.access(),
538                                BufferAccess::ReadWrite | BufferAccess::WriteOnly
539                            )
540                            .then(|| u32::try_from(index).ok())
541                            .flatten()
542                        })
543                        .collect(),
544                )
545            })
546            .as_slice()
547    }
548
549    /// True when the entry walk discovers any indirect dispatch node.
550    #[must_use]
551    #[inline]
552    pub fn has_indirect_dispatch(&self) -> bool {
553        *self.has_indirect_dispatch.get_or_init(|| {
554            // Fast-path: ProgramStats records every node kind seen during
555            // its single-pass walk. If the IndirectDispatch bit is unset,
556            // the tree definitely contains no IndirectDispatch nodes and
557            // the explicit traversal below would redundantly visit every
558            // node only to return false. Reading the bit is O(1).
559            if !self
560                .stats()
561                .has_any_node_kind(super::stats::NODE_KIND_INDIRECT_DISPATCH)
562            {
563                return false;
564            }
565            let mut stack: smallvec::SmallVec<[&Node; 32]> = self.entry().iter().rev().collect();
566            while let Some(node) = stack.pop() {
567                match node {
568                    Node::IndirectDispatch { .. } => return true,
569                    Node::If {
570                        then, otherwise, ..
571                    } => {
572                        stack.extend(otherwise.iter().rev());
573                        stack.extend(then.iter().rev());
574                    }
575                    Node::Loop { body, .. } | Node::Block(body) => {
576                        stack.extend(body.iter().rev());
577                    }
578                    Node::Region { body, .. } => {
579                        stack.extend(body.iter().rev());
580                    }
581                    Node::Let { .. }
582                    | Node::Assign { .. }
583                    | Node::Store { .. }
584                    | Node::AllReduce { .. }
585                    | Node::AllGather { .. }
586                    | Node::ReduceScatter { .. }
587                    | Node::Broadcast { .. }
588                    | Node::Return
589                    | Node::Barrier { .. }
590                    | Node::AsyncLoad { .. }
591                    | Node::AsyncStore { .. }
592                    | Node::AsyncWait { .. }
593                    | Node::Trap { .. }
594                    | Node::Resume { .. }
595                    | Node::Opaque(_) => {}
596                }
597            }
598            false
599        })
600    }
601
602    /// Check whether a named buffer exists.
603    #[must_use]
604    #[inline]
605    pub fn has_buffer(&self, name: &str) -> bool {
606        self.buffer_index.contains_key(name)
607    }
608
609    /// Number of declared buffers.
610    #[must_use]
611    #[inline]
612    pub fn buffer_count(&self) -> usize {
613        self.buffers.len()
614    }
615
616    #[inline]
617    pub(super) fn build_buffer_index(
618        buffers: &[super::BufferDecl],
619    ) -> rustc_hash::FxHashMap<Arc<str>, usize> {
620        let mut index = rustc_hash::FxHashMap::default();
621        index.reserve(buffers.len());
622        for (buffer_index, buffer) in buffers.iter().enumerate() {
623            index
624                .entry(Arc::clone(&buffer.name))
625                .or_insert(buffer_index);
626        }
627        index
628    }
629
630    /// Mark this program as successfully validated structurally.
631    #[inline]
632    pub fn mark_structurally_validated(&self) {
633        self.structural_validation_fingerprint.store(
634            self.current_validation_fingerprint_token(),
635            Ordering::Release,
636        );
637        self.mutation_provenance
638            .store(ProgramMutationProvenance::Clean as u8, Ordering::Release);
639        self.structural_validated.store(true, Ordering::Release);
640    }
641
642    /// Return true once structural validation has succeeded for this program shape.
643    #[must_use]
644    #[inline]
645    pub fn is_structurally_validated(&self) -> bool {
646        if !self.structural_validated.load(Ordering::Acquire) {
647            return false;
648        }
649        if self.validation_mutation_provenance() == ProgramMutationProvenance::Unknown {
650            self.structural_validated.store(false, Ordering::Release);
651            return false;
652        }
653        let recorded = self
654            .structural_validation_fingerprint
655            .load(Ordering::Acquire);
656        if recorded == 0 || recorded != self.current_validation_fingerprint_token() {
657            self.structural_validated.store(false, Ordering::Release);
658            return false;
659        }
660        true
661    }
662
663    /// Last mutation provenance recorded for validation/cache invalidation.
664    #[must_use]
665    #[inline]
666    pub fn validation_mutation_provenance(&self) -> ProgramMutationProvenance {
667        ProgramMutationProvenance::from_code(self.mutation_provenance.load(Ordering::Acquire))
668    }
669
670    /// Mark the Program as having been mutated by a boundary that cannot name a
671    /// concrete provenance. Validation fails closed until the Program is rebuilt
672    /// through a known constructor or known mutation API.
673    #[inline]
674    pub fn mark_unknown_mutation_provenance(&mut self) {
675        self.invalidate_caches_for(ProgramMutationProvenance::Unknown);
676    }
677
678    /// Mark this program as successfully validated for a specific backend.
679    #[inline]
680    pub fn mark_validated_on(&self, backend_id: &str) {
681        if self.validation_mutation_provenance() == ProgramMutationProvenance::Unknown {
682            return;
683        }
684        self.validation_set
685            .get_or_init(|| Arc::new(dashmap::DashSet::new()))
686            .insert(Arc::from(self.validation_cache_key(backend_id)));
687    }
688
689    /// Return true if this program has been validated for the given backend.
690    #[must_use]
691    #[inline]
692    pub fn is_validated_on(&self, backend_id: &str) -> bool {
693        self.validation_set
694            .get()
695            .is_some_and(|set| set.contains(self.validation_cache_key(backend_id).as_str()))
696    }
697
698    /// Validate the program and cache the successful result on the program.
699    ///
700    /// # Errors
701    ///
702    /// Returns [`crate::IrError::Validation`] with every structured issue when
703    /// the structural validator rejects the program.
704    pub fn validate(&self) -> crate::error::IrResult<()> {
705        if self.validation_mutation_provenance() == ProgramMutationProvenance::Unknown {
706            return Err(crate::error::IrError::WireFormatValidation {
707                message: "program validation cache was invalidated by unknown mutation provenance. Fix: rebuild the Program through Program::wrapped/from_wire or use a named Program mutation API before validating.".into(),
708            });
709        }
710        if self.is_structurally_validated() {
711            return Ok(());
712        }
713        let errors = crate::validate::validate(self);
714        if errors.is_empty() {
715            self.mark_structurally_validated();
716            return Ok(());
717        }
718        Err(crate::error::IrError::Validation { issues: errors })
719    }
720
721    #[inline]
722    /// Estimate the peak VRAM byte size of this Program.
723    ///
724    /// Innovation I.11: Static VRAM Pressure Analysis.
725    /// Returns the total bytes required by all storage and uniform buffers
726    /// declared in the Program. Optimizer passes use this to automatically
727    /// partition workloads if they would exceed a backend-specific safety
728    /// margin.
729    #[must_use]
730    pub fn estimate_peak_vram_bytes(&self) -> u64 {
731        self.buffers
732            .iter()
733            .map(|buffer| {
734                let Some(element_size) = buffer.element.size_bytes() else {
735                    return u64::MAX;
736                };
737                u64::from(buffer.count)
738                    .saturating_mul(u64::try_from(element_size).unwrap_or(u64::MAX))
739            })
740            .fold(0u64, u64::saturating_add)
741    }
742
743    /// Return the peak computational intensity found in any instruction.
744    #[must_use]
745    pub fn peak_intensity(&self) -> OpIntensity {
746        let mut peak = OpIntensity::Free;
747        for node in self.entry() {
748            peak = peak.max(Self::node_intensity(node));
749        }
750        peak
751    }
752
753    fn node_intensity(node: &crate::ir::Node) -> OpIntensity {
754        use crate::ir::Node;
755        match node {
756            Node::Let { value, .. } | Node::Assign { value, .. } => Self::expr_intensity(value),
757            Node::Store { index, value, .. } => {
758                Self::expr_intensity(index).max(Self::expr_intensity(value))
759            }
760            Node::If {
761                cond,
762                then,
763                otherwise,
764            } => {
765                let mut p = Self::expr_intensity(cond);
766                for n in then {
767                    p = p.max(Self::node_intensity(n));
768                }
769                for n in otherwise {
770                    p = p.max(Self::node_intensity(n));
771                }
772                p
773            }
774            Node::Loop { from, to, body, .. } => {
775                let mut p = Self::expr_intensity(from).max(Self::expr_intensity(to));
776                for n in body {
777                    p = p.max(Self::node_intensity(n));
778                }
779                p
780            }
781            Node::Block(nodes) => {
782                let mut p = OpIntensity::Free;
783                for n in nodes {
784                    p = p.max(Self::node_intensity(n));
785                }
786                p
787            }
788            Node::Region { body, .. } => {
789                let mut p = OpIntensity::Free;
790                for n in body.iter() {
791                    p = p.max(Self::node_intensity(n));
792                }
793                p
794            }
795            _ => OpIntensity::Free,
796        }
797    }
798
799    fn expr_intensity(expr: &crate::ir::Expr) -> OpIntensity {
800        use crate::ir::Expr;
801        match expr {
802            Expr::BinOp { op, left, right } => op
803                .intensity()
804                .max(Self::expr_intensity(left))
805                .max(Self::expr_intensity(right)),
806            Expr::UnOp { operand, .. } => Self::expr_intensity(operand),
807            Expr::Load { index, .. } => Self::expr_intensity(index),
808            Expr::Select {
809                cond,
810                true_val,
811                false_val,
812            } => Self::expr_intensity(cond)
813                .max(Self::expr_intensity(true_val))
814                .max(Self::expr_intensity(false_val)),
815            Expr::Cast { value, .. } => Self::expr_intensity(value),
816            Expr::Fma { a, b, c } => Self::expr_intensity(a)
817                .max(Self::expr_intensity(b))
818                .max(Self::expr_intensity(c)),
819            Expr::Atomic {
820                index,
821                value,
822                expected,
823                ..
824            } => {
825                let mut p = Self::expr_intensity(index).max(Self::expr_intensity(value));
826                if let Some(e) = expected {
827                    p = p.max(Self::expr_intensity(e));
828                }
829                p.max(OpIntensity::Heavy)
830            }
831            Expr::SubgroupBallot { cond } => Self::expr_intensity(cond).max(OpIntensity::Heavy),
832            Expr::SubgroupShuffle { value, lane } => Self::expr_intensity(value)
833                .max(Self::expr_intensity(lane))
834                .max(OpIntensity::Heavy),
835            Expr::SubgroupReduce { value, .. } => {
836                Self::expr_intensity(value).max(OpIntensity::Heavy)
837            }
838            _ => OpIntensity::Free,
839        }
840    }
841
842    fn compute_wire_hash(&self) -> blake3::Hash {
843        match self.canonical_wire_hash() {
844            Ok(hash) => hash,
845            Err(error) => {
846                let structural = self.structural_fingerprint_fallback();
847                let err_msg = error.to_string();
848                let mut fallback = Vec::with_capacity(96 + err_msg.len() + structural.len());
849                fallback.extend_from_slice(b"VYRE-PROGRAM-CANONICAL-WIRE-HASH-ERROR\0");
850                fallback.extend_from_slice(err_msg.as_bytes());
851                fallback.push(0);
852                fallback.extend_from_slice(structural.as_bytes());
853                blake3::hash(&fallback)
854            }
855        }
856    }
857
858    fn structural_fingerprint_fallback(&self) -> String {
859        let mut hasher = blake3::Hasher::new();
860        hasher.update(b"VYRE-WIRE-FALLBACK-V4\0");
861        if let Some(id) = self.entry_op_id.as_deref() {
862            hasher.update(id.as_bytes());
863        }
864        hasher.update(b"\0");
865        for axis in &self.workgroup_size {
866            hasher.update(&axis.to_le_bytes());
867        }
868        hasher.update(&[u8::from(self.non_composable_with_self)]);
869        let mut keys: Vec<Vec<u8>> = self
870            .buffers()
871            .iter()
872            .map(buffer_decl_canonical_key)
873            .collect();
874        keys.sort_unstable();
875        for key in keys {
876            hasher.update(&key);
877        }
878        let mut visitor = FallbackWireHasher(&mut hasher);
879        walk_nodes_and_exprs(self, &mut visitor);
880        hasher.finalize().to_hex().to_string()
881    }
882
883    fn validation_cache_key(&self, backend_id: &str) -> String {
884        const HEX: &[u8; 16] = b"0123456789abcdef";
885        let fingerprint = self.current_validation_fingerprint();
886        let mut key = String::with_capacity(backend_id.len() + 1 + 64);
887        key.push_str(backend_id);
888        key.push(':');
889        for &byte in &fingerprint {
890            key.push(HEX[(byte >> 4) as usize] as char);
891            key.push(HEX[(byte & 0x0f) as usize] as char);
892        }
893        key
894    }
895
896    #[inline]
897    pub(super) fn invalidate_caches(&mut self) {
898        self.invalidate_caches_for(ProgramMutationProvenance::InternalShapeMutation);
899    }
900
901    #[inline]
902    pub(super) fn invalidate_caches_for(&mut self, provenance: ProgramMutationProvenance) {
903        self.structural_validated.store(false, Ordering::Release);
904        self.structural_validation_fingerprint
905            .store(0, Ordering::Release);
906        self.mutation_provenance
907            .store(provenance as u8, Ordering::Release);
908        if let Some(set) = self.validation_set.get() {
909            set.clear();
910        }
911        let _ = self.hash.take();
912        let _ = self.fingerprint.take();
913        let _ = self.normalized_cache_digest.take();
914        drop(self.output_buffer_index.take());
915        let _ = self.has_indirect_dispatch.take();
916        drop(self.stats.take());
917    }
918
919    fn current_validation_fingerprint(&self) -> [u8; 32] {
920        *self.compute_wire_hash().as_bytes()
921    }
922
923    fn current_validation_fingerprint_token(&self) -> u64 {
924        let bytes = self.current_validation_fingerprint();
925        let token = u64::from_le_bytes([
926            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
927        ]);
928        token.max(1)
929    }
930
931    #[inline]
932    pub(super) fn wrap_entry(entry: Vec<Node>) -> Vec<Node> {
933        if !Self::entry_needs_root_region(&entry) {
934            return entry;
935        }
936        vec![Node::Region {
937            generator: Ident::from(Self::ROOT_REGION_GENERATOR),
938            source_region: None,
939            body: Arc::new(entry),
940        }]
941    }
942
943    #[inline]
944    fn entry_needs_root_region(entry: &[Node]) -> bool {
945        entry.is_empty()
946            || entry
947                .iter()
948                .any(|node| !matches!(node, Node::Region { .. }))
949    }
950}
951
952pub(crate) fn buffers_equal_ignoring_declaration_order(
953    left: &[super::BufferDecl],
954    right: &[super::BufferDecl],
955) -> bool {
956    if left.len() != right.len() {
957        return false;
958    }
959
960    // VYRE_IR_HOTSPOTS HIGH (meta.rs:360-379): previous impl allocated
961    // two Vec<Vec<u8>> then sorted on every equality call. Fast-path:
962    // if the slices compare equal in-place (declaration orders match)
963    // we skip the key-materialization entirely. This catches every
964    // Program::clone(prog) == prog and every `Arc::clone`-equivalent
965    // comparison, which dominate the call distribution.
966    if left == right {
967        return true;
968    }
969
970    let mut left_keys = Vec::with_capacity(left.len());
971    left_keys.extend(left.iter().map(buffer_decl_canonical_key));
972    let mut right_keys = Vec::with_capacity(right.len());
973    right_keys.extend(right.iter().map(buffer_decl_canonical_key));
974    left_keys.sort_unstable();
975    right_keys.sort_unstable();
976    left_keys == right_keys
977}
978
979/// Canonical bytes for one buffer declaration, used to compare buffer sets
980/// independently of declaration order.
981///
982/// THIS KEY DECIDES PROGRAM EQUALITY. `Program::eq` delegates to
983/// `structural_eq`, which compares buffers through
984/// `buffers_equal_ignoring_declaration_order`, which keys each declaration
985/// here. So any field this function omits is a field two programs may differ in
986/// while comparing EQUAL.
987///
988/// It previously omitted `linear_type` and `shape_predicate` while including
989/// `bytes_extraction`: one of three fields wired in, the other two forgotten.
990/// Both omitted fields DECIDE VALIDATION VERDICTS (`validate::linear_type` and
991/// `validate::shape_predicate::check_shape_predicates`), so a VALID program and
992/// an INVALID one compared equal, and any caller shaped as "a == b, therefore
993/// same program, therefore safe to reuse" was wrong for them.
994///
995/// COMPLETENESS IS ENFORCED BY THE COMPILER, not by review. The destructuring
996/// below binds every field by name, so adding a field to `BufferDecl` FAILS TO
997/// COMPILE here until somebody decides whether program equality depends on it.
998/// That is the only mechanism that catches the next omission at the moment it is
999/// introduced rather than years later through a wrong result. Do not replace it
1000/// with `buffer.field` accesses, and do not silence a new field with `..`.
1001///
1002/// `linear_type` and `shape_predicate` are encoded by the SAME functions the
1003/// wire encoder uses, so program equality and `Program::fingerprint` cannot
1004/// drift apart on these fields.
1005pub(super) fn buffer_decl_canonical_key(buffer: &super::BufferDecl) -> Vec<u8> {
1006    use crate::serial::wire::encode::to_wire::{linear_type_tag, put_shape_predicate};
1007    use crate::serial::wire::framing::{put_len_u32, put_u32, put_u8};
1008    use crate::serial::wire::tags::put_data_type;
1009
1010    // Exhaustive by field name: see the completeness note above. A new
1011    // `BufferDecl` field breaks this pattern and forces a decision.
1012    let super::BufferDecl {
1013        name,
1014        binding,
1015        access,
1016        kind,
1017        element,
1018        count,
1019        is_output,
1020        pipeline_live_out,
1021        output_byte_range,
1022        hints,
1023        bytes_extraction,
1024        linear_type,
1025        shape_predicate,
1026    } = buffer;
1027
1028    let mut key = Vec::with_capacity(96);
1029    if let Err(error) = put_len_u32(&mut key, name.len(), "buffer name length") {
1030        key.extend_from_slice(b"\0name-length-error\0");
1031        key.extend_from_slice(error.as_bytes());
1032    }
1033    key.extend_from_slice(name.as_bytes());
1034    put_u32(&mut key, *binding);
1035    match crate::serial::wire::tags::access_tag::access_tag(access) {
1036        Ok(tag) => put_u8(&mut key, tag),
1037        Err(error) => {
1038            put_u8(&mut key, u8::MAX);
1039            key.extend_from_slice(error.as_bytes());
1040        }
1041    }
1042    put_u8(&mut key, super::cache_digest::memory_kind_cache_tag(*kind));
1043    if let Err(error) = put_data_type(&mut key, element) {
1044        key.extend_from_slice(b"\0dtype-error\0");
1045        key.extend_from_slice(error.as_bytes());
1046    }
1047    put_u32(&mut key, *count);
1048    put_u8(&mut key, u8::from(*is_output));
1049    put_u8(&mut key, u8::from(*pipeline_live_out));
1050    match output_byte_range {
1051        Some(range) => {
1052            put_u8(&mut key, 1);
1053            match u32::try_from(range.start) {
1054                Ok(start) => put_u32(&mut key, start),
1055                Err(error) => {
1056                    put_u32(&mut key, u32::MAX);
1057                    key.extend_from_slice(error.to_string().as_bytes());
1058                }
1059            }
1060            match u32::try_from(range.end) {
1061                Ok(end) => put_u32(&mut key, end),
1062                Err(error) => {
1063                    put_u32(&mut key, u32::MAX);
1064                    key.extend_from_slice(error.to_string().as_bytes());
1065                }
1066            }
1067        }
1068        None => put_u8(&mut key, 0),
1069    }
1070    match hints.coalesce_axis {
1071        Some(axis) => {
1072            put_u8(&mut key, 1);
1073            put_u8(&mut key, axis);
1074        }
1075        None => put_u8(&mut key, 0),
1076    }
1077    put_u32(&mut key, hints.preferred_alignment);
1078    put_u8(
1079        &mut key,
1080        match hints.cache_locality {
1081            super::CacheLocality::Streaming => 0,
1082            super::CacheLocality::Temporal => 1,
1083            super::CacheLocality::Random => 2,
1084        },
1085    );
1086    put_u8(&mut key, u8::from(*bytes_extraction));
1087    put_u8(&mut key, linear_type_tag(*linear_type));
1088    if let Err(error) = put_shape_predicate(&mut key, shape_predicate.as_ref(), 0) {
1089        // The encoder refuses a predicate nested past the wire depth bound, and
1090        // it may already have emitted a prefix before refusing. A bare error
1091        // marker would let two DIFFERENT over-deep predicates share a key,
1092        // reintroducing the exact collision this field was added to close, so
1093        // fall back to the Debug form, which stays injective at any depth.
1094        key.extend_from_slice(b"\0shape-predicate-error\0");
1095        key.extend_from_slice(error.as_bytes());
1096        key.extend_from_slice(format!("{shape_predicate:?}").as_bytes());
1097    }
1098    key
1099}
vyre_foundation/ir_inner/model/program/meta.rs

vyre_foundation/ir_inner/model/program/
meta.rs
