1use 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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
54pub enum MemoryClass {
55 Global,
57 Shared,
59 Constant,
63 Uniform,
69 Scratch,
71}
72
73impl MemoryClass {
74 #[must_use]
77 pub fn is_global_visibility(self) -> bool {
78 matches!(self, Self::Global | Self::Constant)
79 }
80
81 #[must_use]
84 pub fn is_writable(self) -> bool {
85 !matches!(self, Self::Constant)
86 }
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
91pub enum BindingVisibility {
92 ReadOnly,
93 WriteOnly,
94 ReadWrite,
95}
96
97impl BindingVisibility {
98 #[must_use]
101 pub fn is_readable(self) -> bool {
102 matches!(self, Self::ReadOnly | Self::ReadWrite)
103 }
104
105 #[must_use]
108 pub fn is_writable(self) -> bool {
109 matches!(self, Self::WriteOnly | Self::ReadWrite)
110 }
111}
112
113#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
115pub struct BindingSlot {
116 pub slot: u32,
118 pub element_type: DataType,
122 pub element_count: Option<u32>,
124 pub memory_class: MemoryClass,
125 pub visibility: BindingVisibility,
126 pub name: String,
129}
130
131#[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#[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#[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#[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#[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 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#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
459pub struct IntentAnnotatedDescriptor {
460 pub descriptor: KernelDescriptor,
461 pub intents: DescriptorIntentSet,
462}
463
464impl IntentAnnotatedDescriptor {
465 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 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#[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 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
662pub type Name = Arc<str>;
666
667#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
673pub enum MatrixMmaShape {
674 M16N8K16,
676}
677
678#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
680pub enum MatrixMmaElement {
681 F16,
682 BF16,
683 TF32,
684 F32,
685}
686
687#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
689pub enum MatrixMmaLayout {
690 RowMajor,
691 ColMajor,
692}
693
694#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
699pub struct KernelOp {
700 pub kind: KernelOpKind,
701 pub operands: Vec<u32>,
704 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 #[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 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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
748pub enum KernelOpKind {
749 Literal,
753 Copy,
759 LocalInvocationId,
772 GlobalInvocationId,
774 WorkgroupId,
776 SubgroupLocalId,
778 SubgroupSize,
780 LoopIndex { loop_var: Name },
785
786 LoopCarrierInit { name: Name },
793
794 LoopCarrier { name: Name },
806
807 LoopCarrierEnd { name: Name },
813
814 LoadGlobal,
818 LoadShared,
820 LoadConstant,
822 BufferLength,
825 StoreGlobal,
828 StoreShared,
830
831 BinOpKind(BinOp),
835 UnOpKind(UnOp),
837
838 Fma,
841 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 Select,
860 Cast { target: DataType },
862 Atomic {
866 op: AtomicOp,
867 ordering: MemoryOrdering,
868 },
869
870 SubgroupBallot,
873 SubgroupShuffle,
875 SubgroupBroadcast,
880 SubgroupReduce { op: SubgroupReduceOp },
883
884 StructuredIfThen,
889 StructuredIfThenElse,
892 StructuredForLoop { loop_var: Name },
896 StructuredBlock,
899 Return,
901 Barrier { ordering: MemoryOrdering },
903 Region { generator: Name },
909
910 AsyncLoad { tag: Name },
915 AsyncStore { tag: Name },
918 AsyncWait { tag: Name },
920
921 Trap { tag: Name },
924 Resume { tag: Name },
926
927 IndirectDispatch { count_offset: u64 },
932
933 Call { op_id: Name },
938
939 OpaqueExpr(Box<OpaqueExprData>),
949 OpaqueNode(Box<OpaqueNodeData>),
951}
952
953#[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#[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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1028pub struct KernelBody {
1029 pub ops: Vec<KernelOp>,
1030 pub child_bodies: Vec<KernelBody>,
1033 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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1051pub struct KernelDescriptor {
1052 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 #[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 #[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 #[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 #[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 #[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 #[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 #[must_use]
1170 pub fn is_empty(&self) -> bool {
1171 self.total_ops() == 0
1172 }
1173
1174 #[must_use]
1180 pub fn is_pure(&self) -> bool {
1181 !self.has_side_effects()
1182 }
1183
1184 pub fn ops_iter(&self) -> KernelOpsIter<'_> {
1189 KernelOpsIter {
1190 stack: vec![(&self.body, 0)],
1191 }
1192 }
1193
1194 #[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 #[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 #[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 #[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 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 StructuredIfThen
1262 | StructuredIfThenElse
1263 | StructuredForLoop { .. }
1264 | StructuredBlock
1265 | Region { .. } => false,
1266 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
1302pub struct KernelOpsIter<'a> {
1304 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 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); assert_eq!(d.total_ops(), 3); }
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 let b = d.body_at(&[0]).unwrap();
1643 assert!(b.ops.is_empty());
1644 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 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 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 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); 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"); 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 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 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 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 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}