Skip to main content

vyre_lower/
descriptor.rs

1//! Substrate-neutral kernel descriptor.
2//!
3//! This is the type that lives BETWEEN the optimizer and the emitters.
4//! Every emitter takes a `KernelDescriptor` and produces a backend
5//! artifact.
6//!
7//! ## Design principles
8//!
9//! - **Faithful to vyre IR**: embeds the same `BinOp`, `UnOp`,
10//!   `AtomicOp`, `MemoryOrdering`, and `DataType` enums as the IR. No
11//!   re-enumeration that would force the lowering to map "vyre IR op X"
12//!   to "descriptor op Y" with a translation table; the descriptor
13//!   carries the same op identity.
14//! - **SSA-shaped**: every value-producing op assigns a unique 32-bit
15//!   `result` id. Operands reference earlier results by id. No named
16//!   variables at this layer  -  the lowering pass converts vyre IR's
17//!   named bindings (`Node::Let`, `Node::Assign`, `Expr::Var`) into
18//!   id references.
19//! - **Structured control flow only**: `StructuredIfThen`,
20//!   `StructuredIfThenElse`, `StructuredForLoop` carry indices into
21//!   `KernelBody::child_bodies`. There is no goto / arbitrary jump;
22//!   that's an explicit constraint required by structured compute
23//!   emitters and low-level instruction emitters alike.
24//! - **Substrate-neutral**: nothing in this module names any specific
25//!   backend. Substrate-specific assumptions live in emitter crates.
26//! - **Round-trippable**: serde-derived for every value; emitters can
27//!   cache descriptors on disk.
28
29use serde::{Deserialize, Serialize};
30use std::sync::Arc;
31use vyre_foundation::ir::{AtomicOp, BinOp, DataType, SubgroupReduceOp, UnOp};
32use vyre_foundation::runtime::memory_model::MemoryOrdering;
33
34pub const TRAP_SIDECAR_NAME: &str = "__vyre_descriptor_trap_sidecar";
35pub const TRAP_SIDECAR_WORDS: u32 = 4;
36
37/// Workgroup dispatch shape. `[x, y, z]` matches every modern
38/// compute backend. `(1, 1, 1)` is a single invocation per workgroup.
39#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
40pub struct Dispatch {
41    pub workgroup_size: [u32; 3],
42}
43
44impl Dispatch {
45    pub const fn new(x: u32, y: u32, z: u32) -> Self {
46        Self {
47            workgroup_size: [x, y, z],
48        }
49    }
50}
51
52/// Where a binding's storage lives.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
54pub enum MemoryClass {
55    /// Global / device memory; visible to every workgroup.
56    Global,
57    /// Workgroup-shared memory.
58    Shared,
59    /// Read-only constant memory backed by a storage buffer
60    /// (`BufferDecl::storage(.., ReadOnly, ..)`). Bind in group 0
61    /// alongside `Global` writers.
62    Constant,
63    /// True uniform-buffer memory backed by `BufferDecl::uniform`.
64    /// Maps to WGSL `var<uniform>` / Vulkan `uniform_buffer` descriptor
65    /// and binds in group 1 per `bind_group_for`. Distinct from
66    /// `Constant` so the emitter can pick `AddressSpace::Uniform` and
67    /// the layout builder can reserve the second bind group.
68    Uniform,
69    /// Backend-managed scratch storage.
70    Scratch,
71}
72
73impl MemoryClass {
74    /// True iff this memory class is visible across workgroups
75    /// (Global, Constant). Shared and Scratch are workgroup-local.
76    #[must_use]
77    pub fn is_global_visibility(self) -> bool {
78        matches!(self, Self::Global | Self::Constant)
79    }
80
81    /// True iff this memory class can be written by the kernel.
82    /// Constant is read-only; the rest are writable.
83    #[must_use]
84    pub fn is_writable(self) -> bool {
85        !matches!(self, Self::Constant)
86    }
87}
88
89/// Read/write visibility for a binding slot.
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
91pub enum BindingVisibility {
92    ReadOnly,
93    WriteOnly,
94    ReadWrite,
95}
96
97impl BindingVisibility {
98    /// True iff the binding can be read by the kernel
99    /// (`ReadOnly` or `ReadWrite`).
100    #[must_use]
101    pub fn is_readable(self) -> bool {
102        matches!(self, Self::ReadOnly | Self::ReadWrite)
103    }
104
105    /// True iff the binding can be written by the kernel
106    /// (`WriteOnly` or `ReadWrite`).
107    #[must_use]
108    pub fn is_writable(self) -> bool {
109        matches!(self, Self::WriteOnly | Self::ReadWrite)
110    }
111}
112
113/// One bound buffer at the kernel boundary.
114#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
115pub struct BindingSlot {
116    /// Bind-group-slot index, stable across emitters.
117    pub slot: u32,
118    /// Element type. Carries the full vyre IR DataType so emitters can
119    /// reproduce the exact type information (lane counts, sparse
120    /// layouts, etc.).
121    pub element_type: DataType,
122    /// Element count. `None` means runtime-sized.
123    pub element_count: Option<u32>,
124    pub memory_class: MemoryClass,
125    pub visibility: BindingVisibility,
126    /// Caller-friendly identifier (for debug; does NOT participate in
127    /// kernel hashing).
128    pub name: String,
129}
130
131/// Full binding layout for a kernel.
132#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
133pub struct BindingLayout {
134    pub slots: Vec<BindingSlot>,
135}
136
137pub const DESCRIPTOR_INTENT_SCHEMA_VERSION: u32 = 1;
138
139/// Backend-neutral scan intent attached beside a [`KernelDescriptor`].
140///
141/// These intents describe why descriptor regions exist, not how a backend must
142/// implement them. CUDA, WGPU, Metal, SPIR-V, and CPU emitters can route from
143/// these strategy classes without owning scan compiler policy.
144#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
145pub enum DescriptorIntentKind {
146    LiteralPrefilter,
147    AutomataTransition,
148    Verifier,
149    OutputCompaction,
150    RelationSeed,
151    StreamingState,
152}
153
154impl DescriptorIntentKind {
155    #[must_use]
156    pub const fn strategy(self) -> DescriptorIntentStrategy {
157        match self {
158            Self::LiteralPrefilter => DescriptorIntentStrategy::Prefilter,
159            Self::AutomataTransition => DescriptorIntentStrategy::Automata,
160            Self::Verifier => DescriptorIntentStrategy::Verifier,
161            Self::OutputCompaction => DescriptorIntentStrategy::Compaction,
162            Self::RelationSeed => DescriptorIntentStrategy::RelationSeed,
163            Self::StreamingState => DescriptorIntentStrategy::Streaming,
164        }
165    }
166
167    const fn digest_tag(self) -> u8 {
168        match self {
169            Self::LiteralPrefilter => 1,
170            Self::AutomataTransition => 2,
171            Self::Verifier => 3,
172            Self::OutputCompaction => 4,
173            Self::RelationSeed => 5,
174            Self::StreamingState => 6,
175        }
176    }
177}
178
179/// Routing class derived from [`DescriptorIntentKind`].
180#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
181pub enum DescriptorIntentStrategy {
182    Prefilter,
183    Automata,
184    Verifier,
185    Compaction,
186    RelationSeed,
187    Streaming,
188}
189
190#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
191pub enum ScanConstructIntentClass {
192    Literal,
193    Automata,
194    Verifier,
195    Derivative,
196    ExternalAccelerator,
197}
198
199#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)]
200pub struct ScanConstructIntentMapping {
201    pub construct_id: &'static str,
202    pub tier: &'static str,
203    pub classes: &'static [ScanConstructIntentClass],
204    pub required_intents: &'static [DescriptorIntentKind],
205    pub verifier_fragment_id: Option<&'static str>,
206}
207
208impl ScanConstructIntentMapping {
209    #[must_use]
210    pub fn requires_verifier_fragment(self) -> bool {
211        self.required_intents
212            .iter()
213            .any(|intent| *intent == DescriptorIntentKind::Verifier)
214    }
215
216    #[must_use]
217    pub fn includes_class(self, class: ScanConstructIntentClass) -> bool {
218        self.classes.iter().any(|candidate| *candidate == class)
219    }
220}
221
222pub const SCAN_CONSTRUCT_INTENT_MAPPINGS: &[ScanConstructIntentMapping] = &[
223    ScanConstructIntentMapping {
224        construct_id: "regular_exact_core",
225        tier: "supported",
226        classes: &[
227            ScanConstructIntentClass::Literal,
228            ScanConstructIntentClass::Automata,
229        ],
230        required_intents: &[
231            DescriptorIntentKind::LiteralPrefilter,
232            DescriptorIntentKind::AutomataTransition,
233            DescriptorIntentKind::OutputCompaction,
234        ],
235        verifier_fragment_id: None,
236    },
237    ScanConstructIntentMapping {
238        construct_id: "unsupported_backtracking_constructs",
239        tier: "rejected",
240        classes: &[ScanConstructIntentClass::Verifier],
241        required_intents: &[DescriptorIntentKind::Verifier],
242        verifier_fragment_id: Some("unsupported-backtracking-diagnostic"),
243    },
244    ScanConstructIntentMapping {
245        construct_id: "lookaround_prefilter_constructs",
246        tier: "approximated",
247        classes: &[
248            ScanConstructIntentClass::Literal,
249            ScanConstructIntentClass::Verifier,
250        ],
251        required_intents: &[
252            DescriptorIntentKind::LiteralPrefilter,
253            DescriptorIntentKind::Verifier,
254        ],
255        verifier_fragment_id: Some("lookaround-verifier"),
256    },
257    ScanConstructIntentMapping {
258        construct_id: "hardware_rule_database_constructs",
259        tier: "accelerator-only",
260        classes: &[
261            ScanConstructIntentClass::ExternalAccelerator,
262            ScanConstructIntentClass::Verifier,
263        ],
264        required_intents: &[
265            DescriptorIntentKind::Verifier,
266            DescriptorIntentKind::OutputCompaction,
267        ],
268        verifier_fragment_id: Some("external-rule-database-verifier"),
269    },
270    ScanConstructIntentMapping {
271        construct_id: "capture_extraction_constructs",
272        tier: "verifier-required",
273        classes: &[ScanConstructIntentClass::Verifier],
274        required_intents: &[
275            DescriptorIntentKind::Verifier,
276            DescriptorIntentKind::OutputCompaction,
277        ],
278        verifier_fragment_id: Some("capture-extraction-verifier"),
279    },
280    ScanConstructIntentMapping {
281        construct_id: "derivative_regex_constructs",
282        tier: "verifier-required",
283        classes: &[
284            ScanConstructIntentClass::Derivative,
285            ScanConstructIntentClass::Verifier,
286        ],
287        required_intents: &[
288            DescriptorIntentKind::AutomataTransition,
289            DescriptorIntentKind::Verifier,
290            DescriptorIntentKind::OutputCompaction,
291        ],
292        verifier_fragment_id: Some("derivative-regex-verifier"),
293    },
294];
295
296#[must_use]
297pub fn scan_construct_intent_mapping(
298    construct_id: &str,
299) -> Option<&'static ScanConstructIntentMapping> {
300    SCAN_CONSTRUCT_INTENT_MAPPINGS
301        .iter()
302        .find(|mapping| mapping.construct_id == construct_id)
303}
304
305/// One intent annotation for a descriptor region, binding, or result id.
306#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
307pub struct DescriptorIntent {
308    pub kind: DescriptorIntentKind,
309    pub binding_slot: Option<u32>,
310    pub op_result: Option<u32>,
311    pub stream_state_bytes: u32,
312    pub relation_arity: u16,
313    pub section_digest: u64,
314}
315
316impl DescriptorIntent {
317    #[must_use]
318    pub const fn new(kind: DescriptorIntentKind, section_digest: u64) -> Self {
319        Self {
320            kind,
321            binding_slot: None,
322            op_result: None,
323            stream_state_bytes: 0,
324            relation_arity: 0,
325            section_digest,
326        }
327    }
328
329    #[must_use]
330    pub const fn with_binding_slot(mut self, slot: u32) -> Self {
331        self.binding_slot = Some(slot);
332        self
333    }
334
335    #[must_use]
336    pub const fn with_op_result(mut self, result: u32) -> Self {
337        self.op_result = Some(result);
338        self
339    }
340
341    #[must_use]
342    pub const fn with_stream_state_bytes(mut self, bytes: u32) -> Self {
343        self.stream_state_bytes = bytes;
344        self
345    }
346
347    #[must_use]
348    pub const fn with_relation_arity(mut self, arity: u16) -> Self {
349        self.relation_arity = arity;
350        self
351    }
352}
353
354/// Intent sidecar for a descriptor.
355#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
356pub struct DescriptorIntentSet {
357    pub schema_version: u32,
358    pub intents: Vec<DescriptorIntent>,
359}
360
361impl DescriptorIntentSet {
362    #[must_use]
363    pub fn new(intents: Vec<DescriptorIntent>) -> Self {
364        Self {
365            schema_version: DESCRIPTOR_INTENT_SCHEMA_VERSION,
366            intents,
367        }
368    }
369
370    /// Validate intent references against a descriptor and return routing
371    /// evidence for emitters and benchmark artifacts.
372    ///
373    /// # Errors
374    ///
375    /// Returns [`DescriptorIntentError`] when the sidecar is empty,
376    /// version-skewed, references missing descriptor slots/results, or omits
377    /// required relation/streaming payload metadata.
378    pub fn validate_for_descriptor(
379        &self,
380        descriptor: &KernelDescriptor,
381    ) -> Result<DescriptorIntentEvidence, DescriptorIntentError> {
382        if self.schema_version != DESCRIPTOR_INTENT_SCHEMA_VERSION {
383            return Err(DescriptorIntentError::UnsupportedSchemaVersion {
384                expected: DESCRIPTOR_INTENT_SCHEMA_VERSION,
385                actual: self.schema_version,
386            });
387        }
388        if self.intents.is_empty() {
389            return Err(DescriptorIntentError::EmptyIntentSet);
390        }
391
392        let mut evidence = DescriptorIntentEvidence {
393            schema_version: DESCRIPTOR_INTENT_SCHEMA_VERSION,
394            descriptor_id: descriptor.id.clone(),
395            intent_count: self.intents.len(),
396            has_literal_prefilter: false,
397            has_automata_transition: false,
398            has_verifier: false,
399            has_output_compaction: false,
400            has_relation_seed: false,
401            has_streaming_state: false,
402            strategy_digest: stable_descriptor_intent_digest(&descriptor.id, &self.intents),
403        };
404
405        for intent in &self.intents {
406            if intent.section_digest == 0 {
407                return Err(DescriptorIntentError::MissingSectionDigest { kind: intent.kind });
408            }
409            if let Some(slot) = intent.binding_slot {
410                let known = descriptor
411                    .bindings
412                    .slots
413                    .iter()
414                    .any(|binding| binding.slot == slot);
415                if !known {
416                    return Err(DescriptorIntentError::UnknownBindingSlot {
417                        kind: intent.kind,
418                        slot,
419                    });
420                }
421            }
422            if let Some(result) = intent.op_result {
423                if descriptor.find_op_by_id(result).is_none() {
424                    return Err(DescriptorIntentError::UnknownOpResult {
425                        kind: intent.kind,
426                        result,
427                    });
428                }
429            }
430
431            match intent.kind {
432                DescriptorIntentKind::LiteralPrefilter => evidence.has_literal_prefilter = true,
433                DescriptorIntentKind::AutomataTransition => {
434                    evidence.has_automata_transition = true;
435                }
436                DescriptorIntentKind::Verifier => evidence.has_verifier = true,
437                DescriptorIntentKind::OutputCompaction => evidence.has_output_compaction = true,
438                DescriptorIntentKind::RelationSeed => {
439                    if intent.relation_arity == 0 {
440                        return Err(DescriptorIntentError::MissingRelationArity);
441                    }
442                    evidence.has_relation_seed = true;
443                }
444                DescriptorIntentKind::StreamingState => {
445                    if intent.stream_state_bytes == 0 {
446                        return Err(DescriptorIntentError::MissingStreamingStateBytes);
447                    }
448                    evidence.has_streaming_state = true;
449                }
450            }
451        }
452
453        Ok(evidence)
454    }
455}
456
457/// Descriptor plus validated intent sidecar.
458#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
459pub struct IntentAnnotatedDescriptor {
460    pub descriptor: KernelDescriptor,
461    pub intents: DescriptorIntentSet,
462}
463
464impl IntentAnnotatedDescriptor {
465    /// Construct an intent-annotated descriptor after validating intent
466    /// references against the descriptor body and binding layout.
467    ///
468    /// # Errors
469    ///
470    /// Returns [`DescriptorIntentError`] when the intent set is incomplete or
471    /// references descriptor state that does not exist.
472    pub fn try_new(
473        descriptor: KernelDescriptor,
474        intents: DescriptorIntentSet,
475    ) -> Result<Self, DescriptorIntentError> {
476        intents.validate_for_descriptor(&descriptor)?;
477        Ok(Self {
478            descriptor,
479            intents,
480        })
481    }
482
483    /// Preserve descriptor intent through a rewrite while revalidating any
484    /// referenced binding slots or result ids against the rewritten descriptor.
485    ///
486    /// This is the neutral seam VX-337 needs: descriptor rewrites transform
487    /// descriptor shape, while scan decomposition intent remains attached until
488    /// the rewrite actually invalidates one of its declared references.
489    ///
490    /// # Errors
491    ///
492    /// Returns [`DescriptorIntentError`] when the rewritten descriptor no
493    /// longer contains a binding slot or result id referenced by the sidecar.
494    pub fn preserve_intents_after_rewrite(
495        self,
496        descriptor: KernelDescriptor,
497    ) -> Result<Self, DescriptorIntentError> {
498        Self::try_new(descriptor, self.intents)
499    }
500
501    pub fn evidence(&self) -> Result<DescriptorIntentEvidence, DescriptorIntentError> {
502        self.intents.validate_for_descriptor(&self.descriptor)
503    }
504}
505
506#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
507pub struct DescriptorIntentEvidence {
508    pub schema_version: u32,
509    pub descriptor_id: String,
510    pub intent_count: usize,
511    pub has_literal_prefilter: bool,
512    pub has_automata_transition: bool,
513    pub has_verifier: bool,
514    pub has_output_compaction: bool,
515    pub has_relation_seed: bool,
516    pub has_streaming_state: bool,
517    pub strategy_digest: u64,
518}
519
520impl DescriptorIntentEvidence {
521    #[must_use]
522    pub const fn covers_full_scan_pipeline(&self) -> bool {
523        self.schema_version == DESCRIPTOR_INTENT_SCHEMA_VERSION
524            && self.intent_count > 0
525            && self.has_literal_prefilter
526            && self.has_automata_transition
527            && self.has_verifier
528            && self.has_output_compaction
529            && self.has_relation_seed
530            && self.has_streaming_state
531            && self.strategy_digest != 0
532    }
533}
534
535#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
536pub enum DescriptorIntentError {
537    UnsupportedSchemaVersion {
538        expected: u32,
539        actual: u32,
540    },
541    EmptyIntentSet,
542    MissingSectionDigest {
543        kind: DescriptorIntentKind,
544    },
545    UnknownBindingSlot {
546        kind: DescriptorIntentKind,
547        slot: u32,
548    },
549    UnknownOpResult {
550        kind: DescriptorIntentKind,
551        result: u32,
552    },
553    MissingRelationArity,
554    MissingStreamingStateBytes,
555}
556
557impl std::fmt::Display for DescriptorIntentError {
558    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
559        match self {
560            Self::UnsupportedSchemaVersion { expected, actual } => write!(
561                f,
562                "descriptor intent schema version {actual} is unsupported; expected {expected}"
563            ),
564            Self::EmptyIntentSet => f.write_str("descriptor intent set is empty"),
565            Self::MissingSectionDigest { kind } => {
566                write!(f, "descriptor intent {kind:?} is missing section_digest")
567            }
568            Self::UnknownBindingSlot { kind, slot } => write!(
569                f,
570                "descriptor intent {kind:?} references unknown binding slot {slot}"
571            ),
572            Self::UnknownOpResult { kind, result } => write!(
573                f,
574                "descriptor intent {kind:?} references unknown op result {result}"
575            ),
576            Self::MissingRelationArity => {
577                f.write_str("relation-seed descriptor intent is missing relation_arity")
578            }
579            Self::MissingStreamingStateBytes => {
580                f.write_str("streaming-state descriptor intent is missing stream_state_bytes")
581            }
582        }
583    }
584}
585
586impl std::error::Error for DescriptorIntentError {}
587
588fn stable_descriptor_intent_digest(descriptor_id: &str, intents: &[DescriptorIntent]) -> u64 {
589    const FNV_OFFSET: u64 = 0xcbf29ce484222325;
590    const FNV_PRIME: u64 = 0x100000001b3;
591
592    fn mix(mut digest: u64, byte: u8) -> u64 {
593        digest ^= u64::from(byte);
594        digest.wrapping_mul(FNV_PRIME)
595    }
596
597    fn mix_u64(mut digest: u64, value: u64) -> u64 {
598        for byte in value.to_le_bytes() {
599            digest = mix(digest, byte);
600        }
601        digest
602    }
603
604    let mut digest = FNV_OFFSET;
605    for byte in descriptor_id.as_bytes() {
606        digest = mix(digest, *byte);
607    }
608    for intent in intents {
609        digest = mix(digest, intent.kind.digest_tag());
610        digest = mix_u64(digest, u64::from(intent.binding_slot.unwrap_or(u32::MAX)));
611        digest = mix_u64(digest, u64::from(intent.op_result.unwrap_or(u32::MAX)));
612        digest = mix_u64(digest, u64::from(intent.stream_state_bytes));
613        digest = mix_u64(digest, u64::from(intent.relation_arity));
614        digest = mix_u64(digest, intent.section_digest);
615    }
616    if digest == 0 {
617        FNV_OFFSET
618    } else {
619        digest
620    }
621}
622
623/// A literal value that can sit in the literal pool.
624#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
625pub enum LiteralValue {
626    U32(u32),
627    I32(i32),
628    F32(f32),
629    Bool(bool),
630}
631
632impl Eq for LiteralValue {}
633
634impl std::hash::Hash for LiteralValue {
635    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
636        match self {
637            Self::U32(v) => {
638                0u8.hash(state);
639                v.hash(state);
640            }
641            Self::I32(v) => {
642                1u8.hash(state);
643                v.hash(state);
644            }
645            // Hash f32 by its bit pattern so NaN-with-different-payloads
646            // hash distinctly. Equality uses bit pattern too via PartialEq
647            // on the `==` of f32  -  note this means two NaNs are not equal,
648            // which is correct for caching purposes (they CAN be different
649            // NaNs).
650            Self::F32(v) => {
651                2u8.hash(state);
652                v.to_bits().hash(state);
653            }
654            Self::Bool(v) => {
655                3u8.hash(state);
656                v.hash(state);
657            }
658        }
659    }
660}
661
662/// Stable identifier for a named entity (variable, region label, async
663/// tag, trap tag). Mirrors vyre-foundation's `Ident` shape so the
664/// lowering can preserve names for diagnostics.
665pub type Name = Arc<str>;
666
667/// Matrix multiply-accumulate tile shape for descriptor-level MMA ops.
668///
669/// These are mathematical fragment shapes, not backend instruction names.
670/// Emitters map supported shapes to their native substrate and reject shapes
671/// they cannot lower.
672#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
673pub enum MatrixMmaShape {
674    /// 16 rows × 8 columns × 16 reduction lanes.
675    M16N8K16,
676}
677
678/// Element type used by a matrix MMA fragment.
679#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
680pub enum MatrixMmaElement {
681    F16,
682    BF16,
683    TF32,
684    F32,
685}
686
687/// Matrix fragment layout.
688#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
689pub enum MatrixMmaLayout {
690    RowMajor,
691    ColMajor,
692}
693
694/// One lowered op in the kernel body. Operands are referenced by
695/// 32-bit id; the id space is per-`KernelBody`. SoA-friendly: an
696/// emitter walks `body.ops` linearly and looks up operand ops by id
697/// when needed.
698#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
699pub struct KernelOp {
700    pub kind: KernelOpKind,
701    /// Operand ids into the same `KernelBody.ops` (or the literal pool
702    /// for `Literal*` kinds  -  see the per-kind documentation).
703    pub operands: Vec<u32>,
704    /// Result id this op assigns. `None` for ops with no value
705    /// (stores, barriers, returns, structured-control-flow markers).
706    pub result: Option<u32>,
707}
708
709impl Eq for KernelOp {}
710
711impl std::hash::Hash for KernelOp {
712    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
713        self.kind.hash(state);
714        self.operands.hash(state);
715        self.result.hash(state);
716    }
717}
718
719impl KernelOp {
720    /// Number of result ids this op defines.
721    #[must_use]
722    pub fn result_id_count(&self) -> u32 {
723        match self.kind {
724            KernelOpKind::MatrixMma { .. } => 4,
725            _ => u32::from(self.result.is_some()),
726        }
727    }
728
729    /// Every result id produced by this op.
730    ///
731    /// Most descriptor ops produce zero or one id. Matrix MMA produces a
732    /// compact four-id accumulator tuple starting at `result`.
733    pub fn result_ids(&self) -> impl Iterator<Item = u32> + '_ {
734        let base = self.result;
735        (0..self.result_id_count())
736            .filter_map(move |offset| base.and_then(|id| id.checked_add(offset)))
737    }
738}
739
740/// Lowered op kinds. Closed enum but covers the entire vyre IR
741/// surface. Adding a new vyre IR variant requires a matching variant
742/// here AND emit rules in every `vyre-emit-*` crate  -  that's the cost
743/// of substrate parity.
744///
745/// Operand semantics are documented per variant. Reading a kind without
746/// reading its operand contract gives wrong code.
747#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
748pub enum KernelOpKind {
749    // ---------- Literals ----------
750    /// Operand 0 = index into `KernelBody.literals`. Result is the
751    /// literal value typed per the literal pool entry.
752    Literal,
753    /// Snapshot a result value. Operand 0 = source result id. Result is
754    /// a fresh SSA value with the source value at this program point.
755    /// This is required when a source-level `Let` captures a mutable
756    /// loop carrier: aliasing the carrier result id would read the
757    /// later carrier value after a subsequent `Assign`.
758    Copy,
759    // ---------- Variable binding (lowered from Node::Let/Assign and Expr::Var) ----------
760    //
761    // The lowering pass converts vyre IR's named variables into SSA
762    // form. `Node::Let` becomes "the result id of the bound expression
763    // is now what `Var(name)` refers to in subsequent ops". `Node::Assign`
764    // becomes a fresh result id that supersedes the earlier one. Names
765    // are erased at this layer; the emitter never sees them.
766
767    // ---------- Builtins ----------
768    /// `LocalInvocationId.x/y/z`. Operand 0 = axis (0/1/2) as a small
769    /// inline literal (NOT a literal-pool reference  -  emit picks the
770    /// builtin directly).
771    LocalInvocationId,
772    /// `GlobalInvocationId.x/y/z`.
773    GlobalInvocationId,
774    /// `WorkgroupId.x/y/z`.
775    WorkgroupId,
776    /// Subgroup local invocation id (a.k.a. lane id).
777    SubgroupLocalId,
778    /// Subgroup size.
779    SubgroupSize,
780    /// Current induction value for the nearest structured loop that
781    /// declared this variable. Produced as the first op in that loop's
782    /// child body so uses of `Expr::Var(loop_var)` remain SSA-shaped
783    /// instead of resolving to the loop's lower bound.
784    LoopIndex { loop_var: Name },
785
786    /// Initialize the loop-carrier slot for `name` from the pre-loop
787    /// SSA value. Emitted ONCE in the PARENT body before the
788    /// `StructuredForLoop` op. Operands: `[seed_value_id]`. No result.
789    /// Emitters allocate a function-scope `LocalVariable` keyed by
790    /// `name` (if not already allocated) and `Store(local, seed_value)`
791    /// in the parent block.
792    LoopCarrierInit { name: Name },
793
794    /// Pure read of the carrier slot for `name`. Operands: `[]`.
795    /// Result: the SSA id that in-loop reads of the source-level
796    /// variable resolve to. Emit semantics: `Load` from the
797    /// function-local allocated by the matching `LoopCarrierInit`.
798    /// Used in three places per loop: (a) once at the top of each
799    /// iteration so per-iteration reads resolve to the latest stored
800    /// value; (b) in the parent body AFTER the loop so post-loop
801    /// readers observe the loop's final value. Without this op,
802    /// `Node::Assign` inside a loop body would have no observable
803    /// effect on subsequent iterations  -  name resolution would always
804    /// pick the pre-loop SSA, which is baked at lowering time.
805    LoopCarrier { name: Name },
806
807    /// Loop-carried-variable write at iteration end. Operands:
808    /// `[final_value_id]`. No result. Pairs with `LoopCarrier { name }`
809    /// to commit the iteration's final value of `name` back to the
810    /// carrier local so the next iteration (or the post-loop reader)
811    /// observes it.
812    LoopCarrierEnd { name: Name },
813
814    // ---------- Buffer access ----------
815    /// `load(buf, index)`. Operands: [binding_slot, index_op_id].
816    /// Result is the loaded value, dtype = binding's element type.
817    LoadGlobal,
818    /// `load(buf, index)` for a workgroup-shared binding.
819    LoadShared,
820    /// `load(buf, index)` for a constant/uniform binding.
821    LoadConstant,
822    /// Buffer length (number of elements). Operand 0 = binding_slot
823    /// inline. Result is u32.
824    BufferLength,
825    /// `store(buf, index, value)`. Operands: [binding_slot, index_op_id, value_op_id].
826    /// Result: None.
827    StoreGlobal,
828    /// `store(buf, index, value)` for a workgroup-shared binding.
829    StoreShared,
830
831    // ---------- Arithmetic / logic ----------
832    /// Binary op. Operands: [left_op_id, right_op_id]. Result has the
833    /// dtype dictated by the operand dtypes (per vyre-spec rules).
834    BinOpKind(BinOp),
835    /// Unary op. Operands: `operand_op_id`. Result dtype per spec.
836    UnOpKind(UnOp),
837
838    // ---------- Composite ops ----------
839    /// Fused multiply-add: `a * b + c`. Operands: [a_id, b_id, c_id].
840    Fma,
841    /// Matrix multiply-accumulate fragment op.
842    ///
843    /// Operand contract for `M16N8K16/F16/F16/F32`:
844    /// `[a0,a1,a2,a3, b0,b1, c0,c1,c2,c3]`, where `a*` and `b*` are
845    /// packed 16-bit fragment words and `c*` are f32 accumulators. `result`
846    /// is the first of four consecutive result ids (`result..result+4`).
847    /// This keeps the descriptor SSA-shaped without adding backend-specific
848    /// register-fragment objects to the neutral IR.
849    MatrixMma {
850        shape: MatrixMmaShape,
851        a_layout: MatrixMmaLayout,
852        b_layout: MatrixMmaLayout,
853        a_type: MatrixMmaElement,
854        b_type: MatrixMmaElement,
855        accum_type: MatrixMmaElement,
856    },
857    /// Conditional select: `if cond { true_val } else { false_val }`.
858    /// Operands: [cond_id, true_val_id, false_val_id].
859    Select,
860    /// Type cast. Operands: `value_id`. The target dtype is on the op.
861    Cast { target: DataType },
862    /// Atomic op. Operands: [binding_slot, index_op_id, value_op_id]
863    /// for most ops. CompareExchange variants prepend `expected_op_id`:
864    /// [binding_slot, index_op_id, expected_op_id, value_op_id].
865    Atomic {
866        op: AtomicOp,
867        ordering: MemoryOrdering,
868    },
869
870    // ---------- Subgroup ops ----------
871    /// Operand 0 = bool-typed cond_op_id. Result is u32 ballot mask.
872    SubgroupBallot,
873    /// Operands: [value_op_id, lane_op_id]. Result has the value's dtype.
874    SubgroupShuffle,
875    /// Operands: [value_op_id, lane_op_id]. Broadcasts `value` from the lane
876    /// named by `lane` (uniform) to every lane; result has the value's dtype.
877    /// Distinct from `SubgroupShuffle` (per-lane source), broadcast requires a
878    /// uniform source lane and emits `subgroupBroadcast`.
879    SubgroupBroadcast,
880    /// Operand 0 = value_op_id. Reduces across the subgroup with `op`;
881    /// result has the value's dtype.
882    SubgroupReduce { op: SubgroupReduceOp },
883
884    // ---------- Structured control flow ----------
885    /// `if (cond) { body }`. Operands: [cond_op_id, child_body_index].
886    /// `child_body_index` references `KernelBody.child_bodies`.
887    /// Result: None.
888    StructuredIfThen,
889    /// `if (cond) { then } else { otherwise }`. Operands:
890    /// [cond_op_id, then_body_index, otherwise_body_index].
891    StructuredIfThenElse,
892    /// `for (var = lo; var < hi; ++var) { body }`. Operands:
893    /// [lo_op_id, hi_op_id, body_index]. The loop variable name is
894    /// embedded on the op (preserved for debug, not for codegen).
895    StructuredForLoop { loop_var: Name },
896    /// Inline statement block  -  explicit grouping; semantically a
897    /// no-op (body is flattened during emit). Operand 0 = body_index.
898    StructuredBlock,
899    /// Function/kernel return. Operands: empty. Result: None.
900    Return,
901    /// Memory barrier with explicit ordering.
902    Barrier { ordering: MemoryOrdering },
903    /// Tracing/grouping marker (vyre IR `Node::Region`). Operand 0 =
904    /// body_index. Carries no execution semantics; emitters MAY pass
905    /// through as a comment or annotation. SEPARATION_AUDIT S5 plans
906    /// to move this to a sidecar; until then it's an op so the
907    /// descriptor preserves it round-trip.
908    Region { generator: Name },
909
910    // ---------- Async ----------
911    /// `cp.async`-style global-to-shared copy. Operands:
912    /// [src_binding, dst_binding, offset_op_id, size_op_id].
913    /// `tag` ties the load to a matching `AsyncWait`.
914    AsyncLoad { tag: Name },
915    /// Mirror of AsyncLoad for shared-to-global. Operands:
916    /// [src_binding, dst_binding, offset_op_id, size_op_id].
917    AsyncStore { tag: Name },
918    /// Wait on a previously-issued AsyncLoad/Store. Operands: empty.
919    AsyncWait { tag: Name },
920
921    // ---------- Effect handlers ----------
922    /// Trap into a host-side effect handler. Operands: `address_op_id`.
923    Trap { tag: Name },
924    /// Resume from a previously-trapped effect.
925    Resume { tag: Name },
926
927    // ---------- Indirect dispatch ----------
928    /// Indirect-dispatch hint. The dispatch shape comes from
929    /// `count_buffer[count_offset]`. Operand 0 = count_buffer
930    /// binding_slot. Result: None.
931    IndirectDispatch { count_offset: u64 },
932
933    // ---------- Calls ----------
934    /// Call into a known op-id (e.g., a vyre-primitives builder
935    /// surface). Operand list is the call's args. The op_id picks the
936    /// callee at emit time.
937    Call { op_id: Name },
938
939    // ---------- Extension escape hatches ----------
940    /// Opaque expression extension. The extension id resolves through
941    /// vyre-core's extension registry. Emitters that don't recognize
942    /// the extension MUST surface an error rather than silently emit
943    /// nothing.
944    ///
945    /// Boxed to keep the common-case `KernelOpKind` small: most ops
946    /// are Literal/BinOp/Load/Store at ≤16 bytes; without boxing,
947    /// every op in the `ops` Vec pays the 52-byte OpaqueExpr tax.
948    OpaqueExpr(Box<OpaqueExprData>),
949    /// Opaque statement-node extension.
950    OpaqueNode(Box<OpaqueNodeData>),
951}
952
953/// Heap-allocated payload for [`KernelOpKind::OpaqueExpr`].
954#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
955
956pub struct OpaqueExprData {
957    pub extension_id: u32,
958    pub extension_kind: String,
959    pub payload: Vec<u8>,
960}
961
962/// Heap-allocated payload for [`KernelOpKind::OpaqueNode`].
963#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
964pub struct OpaqueNodeData {
965    pub extension_kind: String,
966    pub payload: Vec<u8>,
967}
968
969impl Eq for KernelOpKind {}
970
971impl std::hash::Hash for KernelOpKind {
972    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
973        std::mem::discriminant(self).hash(state);
974        match self {
975            Self::BinOpKind(op) => op.hash(state),
976            Self::UnOpKind(op) => op.hash(state),
977            Self::MatrixMma {
978                shape,
979                a_layout,
980                b_layout,
981                a_type,
982                b_type,
983                accum_type,
984            } => {
985                shape.hash(state);
986                a_layout.hash(state);
987                b_layout.hash(state);
988                a_type.hash(state);
989                b_type.hash(state);
990                accum_type.hash(state);
991            }
992            Self::Cast { target } => target.hash(state),
993            Self::Atomic { op, ordering } => {
994                op.hash(state);
995                ordering.hash(state);
996            }
997            Self::StructuredForLoop { loop_var } => loop_var.hash(state),
998            Self::LoopIndex { loop_var } => loop_var.hash(state),
999            Self::Barrier { ordering } => ordering.hash(state),
1000            Self::Region { generator } => generator.hash(state),
1001            Self::AsyncLoad { tag }
1002            | Self::AsyncStore { tag }
1003            | Self::AsyncWait { tag }
1004            | Self::Trap { tag }
1005            | Self::Resume { tag } => tag.hash(state),
1006            Self::LoopCarrierInit { name }
1007            | Self::LoopCarrier { name }
1008            | Self::LoopCarrierEnd { name } => name.hash(state),
1009            Self::IndirectDispatch { count_offset } => count_offset.hash(state),
1010            Self::Call { op_id } => op_id.hash(state),
1011            Self::OpaqueExpr(data) => {
1012                data.extension_id.hash(state);
1013                data.extension_kind.hash(state);
1014                data.payload.hash(state);
1015            }
1016            Self::OpaqueNode(data) => {
1017                data.extension_kind.hash(state);
1018                data.payload.hash(state);
1019            }
1020            _ => {}
1021        }
1022    }
1023}
1024
1025/// One kernel body. Flat op stream + child bodies for nested
1026/// structured control flow. The entry point is `KernelDescriptor.body`.
1027#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1028pub struct KernelBody {
1029    pub ops: Vec<KernelOp>,
1030    /// Child bodies referenced by `StructuredIfThen` etc. operand
1031    /// indices. Indexed from 0 within this body's child_bodies vec.
1032    pub child_bodies: Vec<KernelBody>,
1033    /// Literal pool referenced by `KernelOpKind::Literal` ops.
1034    pub literals: Vec<LiteralValue>,
1035}
1036
1037impl Eq for KernelBody {}
1038
1039impl std::hash::Hash for KernelBody {
1040    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1041        self.ops.hash(state);
1042        self.child_bodies.hash(state);
1043        for lit in &self.literals {
1044            lit.hash(state);
1045        }
1046    }
1047}
1048
1049/// The full kernel descriptor.
1050#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1051pub struct KernelDescriptor {
1052    /// Stable kernel identifier (for caching). Computed from the
1053    /// content hash by `lower::lower`. Empty string until lowering
1054    /// assigns it.
1055    pub id: String,
1056    pub bindings: BindingLayout,
1057    pub dispatch: Dispatch,
1058    pub body: KernelBody,
1059}
1060
1061impl Eq for KernelDescriptor {}
1062
1063impl std::hash::Hash for KernelDescriptor {
1064    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1065        self.id.hash(state);
1066        self.bindings.hash(state);
1067        self.dispatch.hash(state);
1068        self.body.hash(state);
1069    }
1070}
1071
1072impl KernelDescriptor {
1073    /// One-line human-readable summary. Useful for diagnostic output.
1074    /// Format: `"<id>: N ops, M bindings, K child bodies, L literals,
1075    /// dispatch [x, y, z]"`.
1076    #[must_use]
1077    pub fn summary(&self) -> String {
1078        format!(
1079            "{}: {} ops, {} bindings, {} child bodies, {} literals, dispatch {:?}",
1080            if self.id.is_empty() {
1081                "<unnamed>"
1082            } else {
1083                &self.id
1084            },
1085            self.body.ops.len(),
1086            self.bindings.slots.len(),
1087            self.body.child_bodies.len(),
1088            self.body.literals.len(),
1089            self.dispatch.workgroup_size,
1090        )
1091    }
1092
1093    /// Terser alternative to [`Self::summary`]. Format: `"<id>(N ops, M bindings)"`.
1094    /// Useful for compact terminal output where the full summary is
1095    /// too noisy.
1096    #[must_use]
1097    pub fn summary_compact(&self) -> String {
1098        format!(
1099            "{}({} ops, {} bindings)",
1100            if self.id.is_empty() {
1101                "<unnamed>"
1102            } else {
1103                &self.id
1104            },
1105            self.body.ops.len(),
1106            self.bindings.slots.len(),
1107        )
1108    }
1109
1110    /// Total op count across the parent body AND every nested child
1111    /// body, recursively. The parent-only `body.ops.len()` is the
1112    /// flat count; this is the deep count.
1113    #[must_use]
1114    pub fn total_ops(&self) -> usize {
1115        fn walk(b: &KernelBody) -> usize {
1116            b.ops.len() + b.child_bodies.iter().map(walk).sum::<usize>()
1117        }
1118        walk(&self.body)
1119    }
1120
1121    /// Total number of bodies (the parent counts as 1, plus each
1122    /// nested child recursively). Useful for "how nested is this
1123    /// kernel?" telemetry  -  a kernel with one big flat body has
1124    /// `body_count() == 1`; one with deep control flow has more.
1125    #[must_use]
1126    pub fn body_count(&self) -> usize {
1127        fn walk(b: &KernelBody) -> usize {
1128            1 + b.child_bodies.iter().map(walk).sum::<usize>()
1129        }
1130        walk(&self.body)
1131    }
1132
1133    /// Maximum nesting depth of child bodies. A flat kernel returns
1134    /// `0`. A kernel with one If returns `1`. An If-inside-an-If
1135    /// returns `2`. Useful for routing decisions (deeply-nested
1136    /// kernels may need a different optimization strategy).
1137    #[must_use]
1138    pub fn max_body_depth(&self) -> usize {
1139        fn walk(b: &KernelBody) -> usize {
1140            b.child_bodies
1141                .iter()
1142                .map(|c| 1 + walk(c))
1143                .max()
1144                .unwrap_or(0)
1145        }
1146        walk(&self.body)
1147    }
1148
1149    /// Look up a body by its path (a Vec of child-body indices).
1150    /// Empty path returns the parent body. Each element of `path`
1151    /// indexes into the child_bodies of the body it descends into.
1152    /// Returns None if any index is out of range.
1153    ///
1154    /// Matches the `body_path` shape used by `verify::VerifyError`,
1155    /// so tooling can take a verify error and resolve it to the
1156    /// actual body the error refers to.
1157    #[must_use]
1158    pub fn body_at(&self, path: &[usize]) -> Option<&KernelBody> {
1159        let mut current = &self.body;
1160        for &idx in path {
1161            current = current.child_bodies.get(idx)?;
1162        }
1163        Some(current)
1164    }
1165
1166    /// True iff the descriptor has no ops at all (no parent ops AND
1167    /// no ops in any child body). The dispatch geometry and bindings
1168    /// can still be populated  -  this only asks about op content.
1169    #[must_use]
1170    pub fn is_empty(&self) -> bool {
1171        self.total_ops() == 0
1172    }
1173
1174    /// True iff the descriptor is pure  -  no side-effecting ops anywhere.
1175    /// Inverse of `has_side_effects`. Pure kernels can be safely
1176    /// cached by descriptor identity since they produce no observable
1177    /// output (the only "result" is whatever value-flow the consumer
1178    /// inspects, which is fully determined by the descriptor).
1179    #[must_use]
1180    pub fn is_pure(&self) -> bool {
1181        !self.has_side_effects()
1182    }
1183
1184    /// Iterator over every `KernelOp` in the descriptor (parent body
1185    /// + every nested child body, depth-first pre-order). Useful for
1186    /// tooling that wants to walk all ops without writing the
1187    /// recursion themselves.
1188    pub fn ops_iter(&self) -> KernelOpsIter<'_> {
1189        KernelOpsIter {
1190            stack: vec![(&self.body, 0)],
1191        }
1192    }
1193
1194    /// Find the first op anywhere in the descriptor whose `result`
1195    /// matches `id`. Per-body id space means an id may be reused
1196    /// across child bodies  -  this returns the FIRST match in DFS
1197    /// pre-order. For a given body's view, callers should iterate
1198    /// `body.ops` directly.
1199    #[must_use]
1200    pub fn find_op_by_id(&self, id: u32) -> Option<&KernelOp> {
1201        self.ops_iter().find(|op| op.result == Some(id))
1202    }
1203
1204    /// Total threads per workgroup (the product of `dispatch.workgroup_size`).
1205    /// Saturates on overflow rather than wrapping. Useful for
1206    /// per-dispatch resource calculations (shared memory budget,
1207    /// register pressure, etc.).
1208    #[must_use]
1209    pub fn dispatch_total_threads(&self) -> u32 {
1210        let wg = self.dispatch.workgroup_size;
1211        wg[0].saturating_mul(wg[1]).saturating_mul(wg[2])
1212    }
1213
1214    /// Return a clone of this descriptor with a new `id` field.
1215    /// Body, bindings, dispatch all unchanged. Useful for tooling
1216    /// that wants to fork a descriptor for ablation testing or
1217    /// versioning.
1218    #[must_use]
1219    pub fn with_id(&self, id: impl Into<String>) -> Self {
1220        let mut clone = self.clone();
1221        clone.id = id.into();
1222        clone
1223    }
1224
1225    /// True iff the descriptor has at least one side-effecting op (memory
1226    /// write, atomic, sync/async op, barrier, control exit, indirect dispatch,
1227    /// call, or opaque extension). A pure descriptor with no side effects
1228    /// produces no observable output, so a caller is free to drop it entirely.
1229    ///
1230    /// "Side effect" here means observable-or-cross-thread: `AsyncLoad` writes
1231    /// shared memory other threads read, `AsyncWait`/`Barrier` are sync points,
1232    /// and `IndirectDispatch` reconfigures the grid, all are unsafe to drop
1233    /// even though none produces a global-buffer write.
1234    #[must_use]
1235    pub fn has_side_effects(&self) -> bool {
1236        fn walk(b: &KernelBody) -> bool {
1237            for op in &b.ops {
1238                use KernelOpKind::*;
1239                // EXHAUSTIVE ON PURPOSE (no `_` wildcard): a future KernelOpKind
1240                // with an observable or cross-thread effect must be classified
1241                // here, never silently default to "droppable".
1242                let effecting = match op.kind {
1243                    StoreGlobal
1244                    | StoreShared
1245                    | LoopCarrierInit { .. }
1246                    | LoopCarrierEnd { .. }
1247                    | Atomic { .. }
1248                    | AsyncLoad { .. }
1249                    | AsyncStore { .. }
1250                    | AsyncWait { .. }
1251                    | Barrier { .. }
1252                    | Trap { .. }
1253                    | Resume { .. }
1254                    | Return
1255                    | IndirectDispatch { .. }
1256                    | Call { .. }
1257                    | OpaqueExpr(..)
1258                    | OpaqueNode(..) => true,
1259                    // Structured control flow / Region carry no direct effect:
1260                    // their child bodies are walked separately below.
1261                    StructuredIfThen
1262                    | StructuredIfThenElse
1263                    | StructuredForLoop { .. }
1264                    | StructuredBlock
1265                    | Region { .. } => false,
1266                    // Pure value/builtin/arith ops and READS (loads, carrier
1267                    // read, buffer length): no observable or cross-thread effect.
1268                    Literal
1269                    | Copy
1270                    | LocalInvocationId
1271                    | GlobalInvocationId
1272                    | WorkgroupId
1273                    | SubgroupLocalId
1274                    | SubgroupSize
1275                    | LoopIndex { .. }
1276                    | LoopCarrier { .. }
1277                    | LoadGlobal
1278                    | LoadShared
1279                    | LoadConstant
1280                    | BufferLength
1281                    | BinOpKind(_)
1282                    | UnOpKind(_)
1283                    | Fma
1284                    | MatrixMma { .. }
1285                    | Select
1286                    | Cast { .. }
1287                    | SubgroupBallot
1288                    | SubgroupShuffle
1289                    | SubgroupBroadcast
1290                    | SubgroupReduce { .. } => false,
1291                };
1292                if effecting {
1293                    return true;
1294                }
1295            }
1296            b.child_bodies.iter().any(walk)
1297        }
1298        walk(&self.body)
1299    }
1300}
1301
1302/// Iterator returned by [`KernelDescriptor::ops_iter`].
1303pub struct KernelOpsIter<'a> {
1304    /// Stack of (body, next_op_index) frames. Pushed as we descend
1305    /// into child bodies; popped when a body is exhausted.
1306    stack: Vec<(&'a KernelBody, usize)>,
1307}
1308
1309impl<'a> Iterator for KernelOpsIter<'a> {
1310    type Item = &'a KernelOp;
1311    fn next(&mut self) -> Option<Self::Item> {
1312        loop {
1313            let (body, idx) = self.stack.last_mut()?;
1314            if let Some(op) = body.ops.get(*idx) {
1315                *idx += 1;
1316                return Some(op);
1317            }
1318            // Body exhausted  -  push children and pop self.
1319            let body = *body;
1320            self.stack.pop();
1321            for child in body.child_bodies.iter().rev() {
1322                self.stack.push((child, 0));
1323            }
1324        }
1325    }
1326}
1327
1328#[cfg(test)]
1329mod desc_helper_tests {
1330    use super::*;
1331    use vyre_foundation::ir::DataType;
1332
1333    fn build(ops: Vec<KernelOp>, child_bodies: Vec<KernelBody>) -> KernelDescriptor {
1334        KernelDescriptor {
1335            id: "k".into(),
1336            bindings: BindingLayout {
1337                slots: vec![BindingSlot {
1338                    slot: 0,
1339                    element_type: DataType::U32,
1340                    element_count: None,
1341                    memory_class: MemoryClass::Global,
1342                    visibility: BindingVisibility::ReadWrite,
1343                    name: "buf".into(),
1344                }],
1345            },
1346            dispatch: Dispatch::new(64, 1, 1),
1347            body: KernelBody {
1348                ops,
1349                child_bodies,
1350                literals: vec![LiteralValue::U32(7)],
1351            },
1352        }
1353    }
1354
1355    fn literal_op() -> KernelOp {
1356        KernelOp {
1357            kind: KernelOpKind::Literal,
1358            operands: vec![0],
1359            result: Some(0),
1360        }
1361    }
1362
1363    fn full_scan_intents() -> DescriptorIntentSet {
1364        DescriptorIntentSet::new(vec![
1365            DescriptorIntent::new(DescriptorIntentKind::LiteralPrefilter, 11)
1366                .with_binding_slot(0)
1367                .with_op_result(0),
1368            DescriptorIntent::new(DescriptorIntentKind::AutomataTransition, 12)
1369                .with_binding_slot(0)
1370                .with_op_result(0),
1371            DescriptorIntent::new(DescriptorIntentKind::Verifier, 13)
1372                .with_binding_slot(0)
1373                .with_op_result(0),
1374            DescriptorIntent::new(DescriptorIntentKind::OutputCompaction, 14)
1375                .with_binding_slot(0)
1376                .with_op_result(0),
1377            DescriptorIntent::new(DescriptorIntentKind::RelationSeed, 15)
1378                .with_binding_slot(0)
1379                .with_relation_arity(2),
1380            DescriptorIntent::new(DescriptorIntentKind::StreamingState, 16)
1381                .with_binding_slot(0)
1382                .with_stream_state_bytes(64),
1383        ])
1384    }
1385
1386    #[test]
1387    fn descriptor_intents_cover_scan_pipeline_and_route_to_strategies() {
1388        let d = build(vec![literal_op()], vec![]);
1389        let intents = full_scan_intents();
1390        let evidence = intents.validate_for_descriptor(&d).unwrap();
1391
1392        assert!(evidence.covers_full_scan_pipeline());
1393        assert_ne!(evidence.strategy_digest, 0);
1394        assert_eq!(
1395            DescriptorIntentKind::LiteralPrefilter.strategy(),
1396            DescriptorIntentStrategy::Prefilter
1397        );
1398        assert_eq!(
1399            DescriptorIntentKind::AutomataTransition.strategy(),
1400            DescriptorIntentStrategy::Automata
1401        );
1402        assert_eq!(
1403            DescriptorIntentKind::Verifier.strategy(),
1404            DescriptorIntentStrategy::Verifier
1405        );
1406        assert_eq!(
1407            DescriptorIntentKind::OutputCompaction.strategy(),
1408            DescriptorIntentStrategy::Compaction
1409        );
1410        assert_eq!(
1411            DescriptorIntentKind::RelationSeed.strategy(),
1412            DescriptorIntentStrategy::RelationSeed
1413        );
1414        assert_eq!(
1415            DescriptorIntentKind::StreamingState.strategy(),
1416            DescriptorIntentStrategy::Streaming
1417        );
1418    }
1419
1420    #[test]
1421    fn scan_construct_intent_mappings_cover_tiers_and_verifier_fragments() {
1422        let mut tiers = std::collections::BTreeSet::new();
1423        let mut classes = std::collections::BTreeSet::new();
1424        for mapping in SCAN_CONSTRUCT_INTENT_MAPPINGS {
1425            assert!(
1426                !mapping.required_intents.is_empty(),
1427                "Fix: scan construct `{}` must map to at least one descriptor intent",
1428                mapping.construct_id
1429            );
1430            tiers.insert(mapping.tier);
1431            for class in mapping.classes {
1432                classes.insert(*class);
1433            }
1434            if mapping.requires_verifier_fragment() {
1435                assert!(
1436                    mapping.verifier_fragment_id.is_some(),
1437                    "Fix: verifier-required scan construct `{}` must name a verifier fragment",
1438                    mapping.construct_id
1439                );
1440                assert!(
1441                    mapping
1442                        .required_intents
1443                        .contains(&DescriptorIntentKind::Verifier),
1444                    "Fix: verifier-required scan construct `{}` must require DescriptorIntentKind::Verifier",
1445                    mapping.construct_id
1446                );
1447            }
1448        }
1449
1450        for tier in [
1451            "supported",
1452            "rejected",
1453            "approximated",
1454            "accelerator-only",
1455            "verifier-required",
1456        ] {
1457            assert!(
1458                tiers.contains(tier),
1459                "Fix: scan construct intent mappings must cover tier `{tier}`"
1460            );
1461        }
1462        for class in [
1463            ScanConstructIntentClass::Literal,
1464            ScanConstructIntentClass::Automata,
1465            ScanConstructIntentClass::Verifier,
1466            ScanConstructIntentClass::Derivative,
1467            ScanConstructIntentClass::ExternalAccelerator,
1468        ] {
1469            assert!(
1470                classes.contains(&class),
1471                "Fix: scan construct intent mappings must cover class `{class:?}`"
1472            );
1473        }
1474    }
1475
1476    #[test]
1477    fn scan_construct_intent_lookup_returns_shared_contract_rows() {
1478        let lookaround = scan_construct_intent_mapping("lookaround_prefilter_constructs")
1479            .expect("Fix: lookaround construct mapping must exist");
1480        assert!(lookaround.includes_class(ScanConstructIntentClass::Literal));
1481        assert!(lookaround.includes_class(ScanConstructIntentClass::Verifier));
1482        assert_eq!(lookaround.verifier_fragment_id, Some("lookaround-verifier"));
1483
1484        let derivative = scan_construct_intent_mapping("derivative_regex_constructs")
1485            .expect("Fix: derivative construct mapping must exist");
1486        assert!(derivative.includes_class(ScanConstructIntentClass::Derivative));
1487        assert!(derivative
1488            .required_intents
1489            .contains(&DescriptorIntentKind::AutomataTransition));
1490
1491        assert!(scan_construct_intent_mapping("unknown_construct").is_none());
1492    }
1493
1494    #[test]
1495    fn descriptor_intents_survive_rewrites_through_envelope() {
1496        let d = build(vec![literal_op()], vec![]);
1497        let intents = full_scan_intents();
1498        let envelope = IntentAnnotatedDescriptor::try_new(d.clone(), intents.clone()).unwrap();
1499        let rewritten = d.with_id("k.rewritten");
1500
1501        let preserved = envelope
1502            .preserve_intents_after_rewrite(rewritten)
1503            .unwrap();
1504
1505        assert_eq!(preserved.intents, intents);
1506        assert_eq!(preserved.descriptor.id, "k.rewritten");
1507        assert!(preserved.evidence().unwrap().covers_full_scan_pipeline());
1508    }
1509
1510    #[test]
1511    fn descriptor_intents_reject_invalid_references_and_payloads() {
1512        let d = build(vec![literal_op()], vec![]);
1513        let missing_slot = DescriptorIntentSet::new(vec![
1514            DescriptorIntent::new(DescriptorIntentKind::Verifier, 17).with_binding_slot(99),
1515        ]);
1516        assert_eq!(
1517            missing_slot.validate_for_descriptor(&d).unwrap_err(),
1518            DescriptorIntentError::UnknownBindingSlot {
1519                kind: DescriptorIntentKind::Verifier,
1520                slot: 99
1521            }
1522        );
1523
1524        let missing_relation_arity = DescriptorIntentSet::new(vec![DescriptorIntent::new(
1525            DescriptorIntentKind::RelationSeed,
1526            18,
1527        )
1528        .with_binding_slot(0)]);
1529        assert_eq!(
1530            missing_relation_arity
1531                .validate_for_descriptor(&d)
1532                .unwrap_err(),
1533            DescriptorIntentError::MissingRelationArity
1534        );
1535
1536        let missing_stream_bytes = DescriptorIntentSet::new(vec![DescriptorIntent::new(
1537            DescriptorIntentKind::StreamingState,
1538            19,
1539        )
1540        .with_binding_slot(0)]);
1541        assert_eq!(
1542            missing_stream_bytes.validate_for_descriptor(&d).unwrap_err(),
1543            DescriptorIntentError::MissingStreamingStateBytes
1544        );
1545    }
1546
1547    #[test]
1548    fn summary_includes_all_counts() {
1549        let d = build(vec![], vec![]);
1550        let s = d.summary();
1551        assert!(s.contains("k:"));
1552        assert!(s.contains("0 ops"));
1553        assert!(s.contains("1 bindings"));
1554        assert!(s.contains("0 child bodies"));
1555        assert!(s.contains("1 literals"));
1556        assert!(s.contains("[64, 1, 1]"));
1557    }
1558
1559    #[test]
1560    fn summary_compact_terser_form() {
1561        let d = build(
1562            vec![KernelOp {
1563                kind: KernelOpKind::Literal,
1564                operands: vec![0],
1565                result: Some(0),
1566            }],
1567            vec![],
1568        );
1569        let s = d.summary_compact();
1570        assert_eq!(s, "k(1 ops, 1 bindings)");
1571    }
1572
1573    #[test]
1574    fn unnamed_descriptor_uses_unnamed_label() {
1575        let mut d = build(vec![], vec![]);
1576        d.id = String::new();
1577        let s = d.summary();
1578        assert!(s.contains("<unnamed>"));
1579    }
1580
1581    #[test]
1582    fn total_ops_recurses_into_child_bodies() {
1583        let child = KernelBody {
1584            ops: vec![
1585                KernelOp {
1586                    kind: KernelOpKind::Literal,
1587                    operands: vec![0],
1588                    result: Some(0),
1589                },
1590                KernelOp {
1591                    kind: KernelOpKind::Literal,
1592                    operands: vec![0],
1593                    result: Some(1),
1594                },
1595            ],
1596            child_bodies: vec![],
1597            literals: vec![LiteralValue::U32(5)],
1598        };
1599        let parent_ops = vec![KernelOp {
1600            kind: KernelOpKind::Literal,
1601            operands: vec![0],
1602            result: Some(0),
1603        }];
1604        let d = build(parent_ops, vec![child]);
1605        assert_eq!(d.body.ops.len(), 1); // shallow
1606        assert_eq!(d.total_ops(), 3); // 1 parent + 2 child
1607    }
1608
1609    #[test]
1610    fn body_at_empty_path_returns_parent() {
1611        let d = build(
1612            vec![KernelOp {
1613                kind: KernelOpKind::Literal,
1614                operands: vec![0],
1615                result: Some(7),
1616            }],
1617            vec![],
1618        );
1619        let body = d.body_at(&[]).unwrap();
1620        assert_eq!(body.ops.len(), 1);
1621        assert_eq!(body.ops[0].result, Some(7));
1622    }
1623
1624    #[test]
1625    fn body_at_descends_into_children() {
1626        let grandchild = KernelBody {
1627            ops: vec![KernelOp {
1628                kind: KernelOpKind::Literal,
1629                operands: vec![0],
1630                result: Some(99),
1631            }],
1632            child_bodies: vec![],
1633            literals: vec![LiteralValue::U32(7)],
1634        };
1635        let child = KernelBody {
1636            ops: vec![],
1637            child_bodies: vec![grandchild],
1638            literals: vec![],
1639        };
1640        let d = build(vec![], vec![child]);
1641        // Path [0]: first child of parent  -  empty body with one grandchild.
1642        let b = d.body_at(&[0]).unwrap();
1643        assert!(b.ops.is_empty());
1644        // Path [0, 0]: grandchild  -  has the Literal with result 99.
1645        let b = d.body_at(&[0, 0]).unwrap();
1646        assert_eq!(b.ops[0].result, Some(99));
1647    }
1648
1649    #[test]
1650    fn body_at_out_of_range_returns_none() {
1651        let d = build(vec![], vec![]);
1652        assert!(d.body_at(&[5]).is_none());
1653        assert!(d.body_at(&[0, 0]).is_none());
1654    }
1655
1656    #[test]
1657    fn body_count_includes_parent_plus_recursive_children() {
1658        let nested = KernelBody {
1659            ops: vec![],
1660            child_bodies: vec![KernelBody {
1661                ops: vec![],
1662                child_bodies: vec![],
1663                literals: vec![],
1664            }],
1665            literals: vec![],
1666        };
1667        let d = build(vec![], vec![nested]);
1668        // Parent (1) + first child (1) + grandchild (1) = 3.
1669        assert_eq!(d.body_count(), 3);
1670    }
1671
1672    #[test]
1673    fn body_count_flat_kernel_is_one() {
1674        let d = build(vec![], vec![]);
1675        assert_eq!(d.body_count(), 1);
1676    }
1677
1678    #[test]
1679    fn max_body_depth_flat_is_zero() {
1680        let d = build(vec![], vec![]);
1681        assert_eq!(d.max_body_depth(), 0);
1682    }
1683
1684    #[test]
1685    fn max_body_depth_one_if_is_one() {
1686        let child = KernelBody {
1687            ops: vec![],
1688            child_bodies: vec![],
1689            literals: vec![],
1690        };
1691        let d = build(vec![], vec![child]);
1692        assert_eq!(d.max_body_depth(), 1);
1693    }
1694
1695    #[test]
1696    fn max_body_depth_two_levels() {
1697        let grandchild = KernelBody {
1698            ops: vec![],
1699            child_bodies: vec![],
1700            literals: vec![],
1701        };
1702        let child = KernelBody {
1703            ops: vec![],
1704            child_bodies: vec![grandchild],
1705            literals: vec![],
1706        };
1707        let d = build(vec![], vec![child]);
1708        assert_eq!(d.max_body_depth(), 2);
1709    }
1710
1711    #[test]
1712    fn total_ops_zero_for_empty_kernel() {
1713        let d = build(vec![], vec![]);
1714        assert_eq!(d.total_ops(), 0);
1715    }
1716
1717    #[test]
1718    fn is_empty_true_when_no_ops() {
1719        let d = build(vec![], vec![]);
1720        assert!(d.is_empty());
1721    }
1722
1723    #[test]
1724    fn is_empty_false_when_parent_has_ops() {
1725        let d = build(
1726            vec![KernelOp {
1727                kind: KernelOpKind::Literal,
1728                operands: vec![0],
1729                result: Some(0),
1730            }],
1731            vec![],
1732        );
1733        assert!(!d.is_empty());
1734        assert_eq!(d.total_ops(), 1);
1735    }
1736
1737    #[test]
1738    fn is_empty_false_when_child_has_ops() {
1739        let child = KernelBody {
1740            ops: vec![KernelOp {
1741                kind: KernelOpKind::Literal,
1742                operands: vec![0],
1743                result: Some(0),
1744            }],
1745            child_bodies: vec![],
1746            literals: vec![LiteralValue::U32(1)],
1747        };
1748        let d = build(vec![], vec![child]);
1749        assert!(!d.is_empty());
1750        assert_eq!(d.total_ops(), 1);
1751    }
1752
1753    #[test]
1754    fn has_side_effects_true_with_store() {
1755        let d = build(
1756            vec![
1757                KernelOp {
1758                    kind: KernelOpKind::Literal,
1759                    operands: vec![0],
1760                    result: Some(0),
1761                },
1762                KernelOp {
1763                    kind: KernelOpKind::StoreGlobal,
1764                    operands: vec![0, 0, 0],
1765                    result: None,
1766                },
1767            ],
1768            vec![],
1769        );
1770        assert!(d.has_side_effects());
1771    }
1772
1773    #[test]
1774    fn has_side_effects_false_with_only_arithmetic() {
1775        let d = build(
1776            vec![
1777                KernelOp {
1778                    kind: KernelOpKind::Literal,
1779                    operands: vec![0],
1780                    result: Some(0),
1781                },
1782                KernelOp {
1783                    kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Add),
1784                    operands: vec![0, 0],
1785                    result: Some(1),
1786                },
1787            ],
1788            vec![],
1789        );
1790        assert!(!d.has_side_effects());
1791    }
1792
1793    #[test]
1794    fn has_side_effects_true_for_async_and_indirect_dispatch_ops() {
1795        // Regression: AsyncLoad writes shared memory other threads read,
1796        // AsyncWait is a sync point, and IndirectDispatch reconfigures the grid
1797        //: all cross-thread/dispatch effects (like the already-listed Barrier /
1798        // AsyncStore), so a descriptor containing one is NOT droppable. They
1799        // were omitted from the side-effecting set before the exhaustive-match
1800        // change, which would have let a "drop pure descriptor" caller drop one.
1801        for kind in [
1802            KernelOpKind::AsyncLoad { tag: "t".into() },
1803            KernelOpKind::AsyncWait { tag: "t".into() },
1804            KernelOpKind::IndirectDispatch { count_offset: 0 },
1805        ] {
1806            let d = build(
1807                vec![KernelOp {
1808                    kind: kind.clone(),
1809                    operands: vec![0],
1810                    result: None,
1811                }],
1812                vec![],
1813            );
1814            assert!(
1815                d.has_side_effects(),
1816                "{kind:?} is a cross-thread/dispatch effect and must not be droppable"
1817            );
1818            assert!(!d.is_pure(), "{kind:?} must not be classified pure");
1819        }
1820    }
1821
1822    #[test]
1823    fn ops_iter_visits_parent_then_children_in_order() {
1824        let child0 = KernelBody {
1825            ops: vec![
1826                KernelOp {
1827                    kind: KernelOpKind::Literal,
1828                    operands: vec![0],
1829                    result: Some(10),
1830                },
1831                KernelOp {
1832                    kind: KernelOpKind::Literal,
1833                    operands: vec![0],
1834                    result: Some(11),
1835                },
1836            ],
1837            child_bodies: vec![],
1838            literals: vec![LiteralValue::U32(1)],
1839        };
1840        let child1 = KernelBody {
1841            ops: vec![KernelOp {
1842                kind: KernelOpKind::Literal,
1843                operands: vec![0],
1844                result: Some(20),
1845            }],
1846            child_bodies: vec![],
1847            literals: vec![LiteralValue::U32(2)],
1848        };
1849        let d = build(
1850            vec![
1851                KernelOp {
1852                    kind: KernelOpKind::Literal,
1853                    operands: vec![0],
1854                    result: Some(0),
1855                },
1856                KernelOp {
1857                    kind: KernelOpKind::Literal,
1858                    operands: vec![0],
1859                    result: Some(1),
1860                },
1861            ],
1862            vec![child0, child1],
1863        );
1864        let visited: Vec<u32> = d.ops_iter().map(|o| o.result.unwrap()).collect();
1865        // Parent ops (0, 1) first, then child0 (10, 11), then child1 (20).
1866        assert_eq!(visited, vec![0, 1, 10, 11, 20]);
1867    }
1868
1869    #[test]
1870    fn ops_iter_count_matches_total_ops() {
1871        let child = KernelBody {
1872            ops: vec![
1873                KernelOp {
1874                    kind: KernelOpKind::Literal,
1875                    operands: vec![0],
1876                    result: Some(0),
1877                },
1878                KernelOp {
1879                    kind: KernelOpKind::Literal,
1880                    operands: vec![0],
1881                    result: Some(1),
1882                },
1883            ],
1884            child_bodies: vec![],
1885            literals: vec![LiteralValue::U32(7)],
1886        };
1887        let d = build(
1888            vec![KernelOp {
1889                kind: KernelOpKind::Literal,
1890                operands: vec![0],
1891                result: Some(0),
1892            }],
1893            vec![child],
1894        );
1895        assert_eq!(d.ops_iter().count(), d.total_ops());
1896    }
1897
1898    #[test]
1899    fn memory_class_predicates() {
1900        assert!(MemoryClass::Global.is_global_visibility());
1901        assert!(MemoryClass::Constant.is_global_visibility());
1902        assert!(!MemoryClass::Shared.is_global_visibility());
1903        assert!(!MemoryClass::Scratch.is_global_visibility());
1904
1905        assert!(MemoryClass::Global.is_writable());
1906        assert!(MemoryClass::Shared.is_writable());
1907        assert!(MemoryClass::Scratch.is_writable());
1908        assert!(!MemoryClass::Constant.is_writable());
1909    }
1910
1911    #[test]
1912    fn binding_visibility_readable_writable() {
1913        assert!(BindingVisibility::ReadOnly.is_readable());
1914        assert!(!BindingVisibility::ReadOnly.is_writable());
1915        assert!(!BindingVisibility::WriteOnly.is_readable());
1916        assert!(BindingVisibility::WriteOnly.is_writable());
1917        assert!(BindingVisibility::ReadWrite.is_readable());
1918        assert!(BindingVisibility::ReadWrite.is_writable());
1919    }
1920
1921    #[test]
1922    fn dispatch_total_threads_multiplies_dims() {
1923        let d = build(vec![], vec![]);
1924        assert_eq!(d.dispatch_total_threads(), 64); // build() uses Dispatch::new(64, 1, 1)
1925
1926        let mut d2 = build(vec![], vec![]);
1927        d2.dispatch = Dispatch::new(8, 8, 4);
1928        assert_eq!(d2.dispatch_total_threads(), 256);
1929    }
1930
1931    #[test]
1932    fn with_id_preserves_everything_else() {
1933        let d = build(
1934            vec![KernelOp {
1935                kind: KernelOpKind::Literal,
1936                operands: vec![0],
1937                result: Some(0),
1938            }],
1939            vec![],
1940        );
1941        let renamed = d.with_id("renamed");
1942        assert_eq!(renamed.id, "renamed");
1943        assert_eq!(d.id, "k"); // original unchanged
1944        assert_eq!(renamed.body.ops.len(), d.body.ops.len());
1945        assert_eq!(renamed.bindings, d.bindings);
1946        assert_eq!(renamed.dispatch, d.dispatch);
1947    }
1948
1949    #[test]
1950    fn dispatch_total_threads_saturates_on_overflow() {
1951        let mut d = build(vec![], vec![]);
1952        d.dispatch = Dispatch::new(u32::MAX, u32::MAX, u32::MAX);
1953        // Saturating multiplication means we get u32::MAX rather than wrap.
1954        assert_eq!(d.dispatch_total_threads(), u32::MAX);
1955    }
1956
1957    #[test]
1958    fn find_op_by_id_in_parent() {
1959        let d = build(
1960            vec![
1961                KernelOp {
1962                    kind: KernelOpKind::Literal,
1963                    operands: vec![0],
1964                    result: Some(7),
1965                },
1966                KernelOp {
1967                    kind: KernelOpKind::Literal,
1968                    operands: vec![0],
1969                    result: Some(42),
1970                },
1971            ],
1972            vec![],
1973        );
1974        let op = d.find_op_by_id(42).expect("Fix: found");
1975        assert_eq!(op.result, Some(42));
1976        assert!(d.find_op_by_id(99).is_none());
1977    }
1978
1979    #[test]
1980    fn find_op_by_id_finds_in_child() {
1981        let child = KernelBody {
1982            ops: vec![KernelOp {
1983                kind: KernelOpKind::Literal,
1984                operands: vec![0],
1985                result: Some(100),
1986            }],
1987            child_bodies: vec![],
1988            literals: vec![LiteralValue::U32(7)],
1989        };
1990        let d = build(vec![], vec![child]);
1991        assert!(d.find_op_by_id(100).is_some());
1992    }
1993
1994    #[test]
1995    fn ops_iter_empty_descriptor_yields_none() {
1996        let d = build(vec![], vec![]);
1997        assert!(d.ops_iter().next().is_none());
1998    }
1999
2000    #[test]
2001    fn is_pure_inverse_of_has_side_effects() {
2002        let pure_kernel = build(
2003            vec![KernelOp {
2004                kind: KernelOpKind::Literal,
2005                operands: vec![0],
2006                result: Some(0),
2007            }],
2008            vec![],
2009        );
2010        assert!(pure_kernel.is_pure());
2011        assert!(!pure_kernel.has_side_effects());
2012
2013        let impure = build(
2014            vec![
2015                KernelOp {
2016                    kind: KernelOpKind::Literal,
2017                    operands: vec![0],
2018                    result: Some(0),
2019                },
2020                KernelOp {
2021                    kind: KernelOpKind::StoreGlobal,
2022                    operands: vec![0, 0, 0],
2023                    result: None,
2024                },
2025            ],
2026            vec![],
2027        );
2028        assert!(!impure.is_pure());
2029        assert!(impure.has_side_effects());
2030    }
2031}
2032
2033#[cfg(test)]
2034
2035mod tests {
2036    use super::*;
2037
2038    fn binding(slot: u32, element: DataType, mc: MemoryClass) -> BindingSlot {
2039        BindingSlot {
2040            slot,
2041            element_type: element,
2042            element_count: None,
2043            memory_class: mc,
2044            visibility: BindingVisibility::ReadWrite,
2045            name: format!("b{slot}"),
2046        }
2047    }
2048
2049    #[test]
2050    fn empty_descriptor_round_trips_serde_byte_stable() {
2051        let k = KernelDescriptor {
2052            id: "test".into(),
2053            bindings: BindingLayout { slots: vec![] },
2054            dispatch: Dispatch::new(1, 1, 1),
2055            body: KernelBody {
2056                ops: vec![],
2057                child_bodies: vec![],
2058                literals: vec![],
2059            },
2060        };
2061        let json1 = serde_json::to_string(&k).unwrap();
2062        let parsed: KernelDescriptor = serde_json::from_str(&json1).unwrap();
2063        let json2 = serde_json::to_string(&parsed).unwrap();
2064        assert_eq!(json1, json2);
2065        assert_eq!(k, parsed);
2066    }
2067
2068    #[test]
2069    fn one_store_kernel_round_trips_byte_stable() {
2070        let k = KernelDescriptor {
2071            id: "store_one".into(),
2072            bindings: BindingLayout {
2073                slots: vec![binding(0, DataType::U32, MemoryClass::Global)],
2074            },
2075            dispatch: Dispatch::new(1, 1, 1),
2076            body: KernelBody {
2077                ops: vec![
2078                    KernelOp {
2079                        kind: KernelOpKind::Literal,
2080                        operands: vec![0],
2081                        result: Some(0),
2082                    },
2083                    KernelOp {
2084                        kind: KernelOpKind::Literal,
2085                        operands: vec![1],
2086                        result: Some(1),
2087                    },
2088                    KernelOp {
2089                        kind: KernelOpKind::StoreGlobal,
2090                        operands: vec![0, 0, 1],
2091                        result: None,
2092                    },
2093                ],
2094                child_bodies: vec![],
2095                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
2096            },
2097        };
2098        let json1 = serde_json::to_string(&k).unwrap();
2099        let parsed: KernelDescriptor = serde_json::from_str(&json1).unwrap();
2100        let json2 = serde_json::to_string(&parsed).unwrap();
2101        assert_eq!(json1, json2);
2102    }
2103
2104    #[test]
2105    fn binop_kind_carries_full_vyre_spec_op() {
2106        let op = KernelOp {
2107            kind: KernelOpKind::BinOpKind(BinOp::SaturatingAdd),
2108            operands: vec![0, 1],
2109            result: Some(2),
2110        };
2111        let json = serde_json::to_string(&op).unwrap();
2112        let parsed: KernelOp = serde_json::from_str(&json).unwrap();
2113        assert_eq!(op, parsed);
2114        // Confirm the variant survives  -  serde_json round-trip preserves it.
2115        match parsed.kind {
2116            KernelOpKind::BinOpKind(BinOp::SaturatingAdd) => {}
2117            other => panic!("lost BinOp variant: {other:?}"),
2118        }
2119    }
2120
2121    #[test]
2122    fn unop_kind_carries_full_vyre_spec_op() {
2123        let op = KernelOp {
2124            kind: KernelOpKind::UnOpKind(UnOp::InverseSqrt),
2125            operands: vec![5],
2126            result: Some(6),
2127        };
2128        let json = serde_json::to_string(&op).unwrap();
2129        let parsed: KernelOp = serde_json::from_str(&json).unwrap();
2130        assert_eq!(op, parsed);
2131        match parsed.kind {
2132            KernelOpKind::UnOpKind(UnOp::InverseSqrt) => {}
2133            other => panic!("lost UnOp variant: {other:?}"),
2134        }
2135    }
2136
2137    #[test]
2138    fn atomic_carries_op_and_ordering() {
2139        let op = KernelOp {
2140            kind: KernelOpKind::Atomic {
2141                op: AtomicOp::CompareExchange,
2142                ordering: MemoryOrdering::AcqRel,
2143            },
2144            operands: vec![0, 1, 2, 3],
2145            result: Some(4),
2146        };
2147        let json = serde_json::to_string(&op).unwrap();
2148        let parsed: KernelOp = serde_json::from_str(&json).unwrap();
2149        assert_eq!(op, parsed);
2150    }
2151
2152    #[test]
2153    fn nested_if_then_body_round_trips() {
2154        let inner = KernelBody {
2155            ops: vec![KernelOp {
2156                kind: KernelOpKind::Barrier {
2157                    ordering: MemoryOrdering::SeqCst,
2158                },
2159                operands: vec![],
2160                result: None,
2161            }],
2162            child_bodies: vec![],
2163            literals: vec![],
2164        };
2165        let outer = KernelBody {
2166            ops: vec![
2167                KernelOp {
2168                    kind: KernelOpKind::Literal,
2169                    operands: vec![0],
2170                    result: Some(0),
2171                },
2172                KernelOp {
2173                    kind: KernelOpKind::StructuredIfThen,
2174                    operands: vec![0, 0],
2175                    result: None,
2176                },
2177            ],
2178            child_bodies: vec![inner],
2179            literals: vec![LiteralValue::Bool(true)],
2180        };
2181        let k = KernelDescriptor {
2182            id: "if_then".into(),
2183            bindings: BindingLayout { slots: vec![] },
2184            dispatch: Dispatch::new(1, 1, 1),
2185            body: outer,
2186        };
2187        let json1 = serde_json::to_string(&k).unwrap();
2188        let parsed: KernelDescriptor = serde_json::from_str(&json1).unwrap();
2189        let json2 = serde_json::to_string(&parsed).unwrap();
2190        assert_eq!(json1, json2);
2191    }
2192
2193    #[test]
2194    fn for_loop_with_var_name_round_trips() {
2195        let body = KernelBody {
2196            ops: vec![],
2197            child_bodies: vec![],
2198            literals: vec![],
2199        };
2200        let outer = KernelBody {
2201            ops: vec![
2202                KernelOp {
2203                    kind: KernelOpKind::Literal,
2204                    operands: vec![0],
2205                    result: Some(0),
2206                },
2207                KernelOp {
2208                    kind: KernelOpKind::Literal,
2209                    operands: vec![1],
2210                    result: Some(1),
2211                },
2212                KernelOp {
2213                    kind: KernelOpKind::StructuredForLoop {
2214                        loop_var: "i".into(),
2215                    },
2216                    operands: vec![0, 1, 0],
2217                    result: None,
2218                },
2219            ],
2220            child_bodies: vec![body],
2221            literals: vec![LiteralValue::U32(0), LiteralValue::U32(64)],
2222        };
2223        let k = KernelDescriptor {
2224            id: "for_i".into(),
2225            bindings: BindingLayout { slots: vec![] },
2226            dispatch: Dispatch::new(64, 1, 1),
2227            body: outer,
2228        };
2229        let json = serde_json::to_string(&k).unwrap();
2230        let parsed: KernelDescriptor = serde_json::from_str(&json).unwrap();
2231        assert_eq!(k, parsed);
2232    }
2233
2234    #[test]
2235    fn async_load_wait_carry_tag() {
2236        let body = KernelBody {
2237            ops: vec![
2238                KernelOp {
2239                    kind: KernelOpKind::Literal,
2240                    operands: vec![0],
2241                    result: Some(0),
2242                },
2243                KernelOp {
2244                    kind: KernelOpKind::Literal,
2245                    operands: vec![1],
2246                    result: Some(1),
2247                },
2248                KernelOp {
2249                    kind: KernelOpKind::AsyncLoad {
2250                        tag: "chunk-0".into(),
2251                    },
2252                    operands: vec![0, 1, 0, 1],
2253                    result: None,
2254                },
2255                KernelOp {
2256                    kind: KernelOpKind::AsyncWait {
2257                        tag: "chunk-0".into(),
2258                    },
2259                    operands: vec![],
2260                    result: None,
2261                },
2262            ],
2263            child_bodies: vec![],
2264            literals: vec![LiteralValue::U32(0), LiteralValue::U32(16)],
2265        };
2266        let k = KernelDescriptor {
2267            id: "async".into(),
2268            bindings: BindingLayout {
2269                slots: vec![
2270                    binding(0, DataType::U32, MemoryClass::Global),
2271                    binding(1, DataType::U32, MemoryClass::Shared),
2272                ],
2273            },
2274            dispatch: Dispatch::new(64, 1, 1),
2275            body,
2276        };
2277        let json = serde_json::to_string(&k).unwrap();
2278        let parsed: KernelDescriptor = serde_json::from_str(&json).unwrap();
2279        assert_eq!(k, parsed);
2280    }
2281
2282    #[test]
2283    fn cast_op_preserves_target_dtype() {
2284        let op = KernelOp {
2285            kind: KernelOpKind::Cast {
2286                target: DataType::F16,
2287            },
2288            operands: vec![3],
2289            result: Some(4),
2290        };
2291        let json = serde_json::to_string(&op).unwrap();
2292        let parsed: KernelOp = serde_json::from_str(&json).unwrap();
2293        match parsed.kind {
2294            KernelOpKind::Cast {
2295                target: DataType::F16,
2296            } => {}
2297            other => panic!("lost cast target: {other:?}"),
2298        }
2299    }
2300
2301    #[test]
2302    fn binding_carries_full_data_type() {
2303        // Confirm a parametric DataType (Vec) round-trips through binding.
2304        let b = BindingSlot {
2305            slot: 5,
2306            element_type: DataType::Vec {
2307                element: Box::new(DataType::F32),
2308                count: 4,
2309            },
2310            element_count: Some(64),
2311            memory_class: MemoryClass::Global,
2312            visibility: BindingVisibility::ReadWrite,
2313            name: "v4f32".into(),
2314        };
2315        let json = serde_json::to_string(&b).unwrap();
2316        let parsed: BindingSlot = serde_json::from_str(&json).unwrap();
2317        assert_eq!(b, parsed);
2318    }
2319
2320    #[test]
2321    fn literal_value_eq_treats_nan_as_distinct_via_bits() {
2322        let nan1 = LiteralValue::F32(f32::NAN);
2323        let nan2 = LiteralValue::F32(f32::NAN);
2324        // PartialEq for f32 treats NaN as not equal to itself; our derive
2325        // inherits that, so two NaNs are never equal.
2326        assert_ne!(nan1, nan2);
2327    }
2328
2329    #[test]
2330    fn region_op_round_trips_with_generator_name() {
2331        let op = KernelOp {
2332            kind: KernelOpKind::Region {
2333                generator: "vyre.libs.nn.gqa_attention".into(),
2334            },
2335            operands: vec![0],
2336            result: None,
2337        };
2338        let json = serde_json::to_string(&op).unwrap();
2339        let parsed: KernelOp = serde_json::from_str(&json).unwrap();
2340        assert_eq!(op, parsed);
2341    }
2342
2343    #[test]
2344    fn dispatch_constructor_preserves_axes() {
2345        let d = Dispatch::new(64, 4, 2);
2346        assert_eq!(d.workgroup_size, [64, 4, 2]);
2347    }
2348}