1use crate::weight_matrix::QuantKind;
67use std::collections::{HashMap, HashSet};
68use std::panic::Location;
69use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
70use std::sync::{OnceLock, RwLock};
71
72macro_rules! define_backend_enum {
95 ([] $(($variant:ident, $name:literal, $ty:ident)),* $(,)?) => {
96 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
97 pub enum Backend {
98 Cpu,
99 $($variant,)*
100 }
101
102 impl Backend {
103 pub fn as_str(self) -> &'static str {
104 match self {
105 Backend::Cpu => "cpu",
106 $(Backend::$variant => $name,)*
107 }
108 }
109
110 pub const ALL: &'static [Backend] = &[Backend::Cpu, $(Backend::$variant,)*];
113 }
114 };
115}
116crate::weight_matrix::gpu_backend::gpu_backend_table!(define_backend_enum);
117
118impl Backend {
119 pub fn is_accelerator(self) -> bool {
122 !matches!(self, Backend::Cpu)
123 }
124}
125
126impl std::fmt::Display for Backend {
127 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128 f.write_str(self.as_str())
129 }
130}
131
132pub mod op {
135 pub const MATVEC: &str = "matvec";
137 pub const MATVEC_MULTI: &str = "matvec_multi";
139 pub const GEMM_PREFILL: &str = "gemm_prefill";
141 pub const FFN_SWIGLU: &str = "ffn_swiglu";
143 pub const ENGINE_PREFILL_BATCH: &str = "engine.prefill_batch";
146}
147
148#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
151pub struct Key {
152 pub backend: Backend,
153 pub op: &'static str,
155 pub role: &'static str,
159 pub kind: Option<QuantKind>,
161 pub file: &'static str,
162 pub line: u32,
163}
164
165impl Key {
166 fn shape(&self) -> (Backend, &'static str, Option<QuantKind>) {
170 (self.backend, self.op, self.kind)
171 }
172
173 fn kind_name(&self) -> &'static str {
174 self.kind.map_or("f32", QuantKind::name)
175 }
176}
177
178impl std::fmt::Display for Key {
179 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
180 write!(
181 f,
182 "{} {} {} ({}) at {}:{}",
183 self.backend,
184 self.op,
185 self.kind_name(),
186 self.role,
187 self.file,
188 self.line
189 )
190 }
191}
192
193#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
197pub enum Severity {
198 ByDesign,
202 SlowPath,
206}
207
208#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
210pub enum Outcome {
211 Hit,
213 Miss {
216 fallback: &'static str,
217 severity: Severity,
218 },
219}
220
221impl Outcome {
222 pub fn slow_path(fallback: &'static str) -> Self {
224 Outcome::Miss {
225 fallback,
226 severity: Severity::SlowPath,
227 }
228 }
229
230 pub fn by_design(fallback: &'static str) -> Self {
232 Outcome::Miss {
233 fallback,
234 severity: Severity::ByDesign,
235 }
236 }
237
238 pub fn is_miss(self) -> bool {
239 matches!(self, Outcome::Miss { .. })
240 }
241
242 pub fn is_slow_path(self) -> bool {
243 matches!(
244 self,
245 Outcome::Miss {
246 severity: Severity::SlowPath,
247 ..
248 }
249 )
250 }
251
252 pub fn fallback(self) -> Option<&'static str> {
253 match self {
254 Outcome::Miss { fallback, .. } => Some(fallback),
255 Outcome::Hit => None,
256 }
257 }
258}
259
260#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
263pub enum Phase {
264 Build,
265 Run,
266}
267
268#[derive(Clone, Copy, Debug)]
271pub struct Lookup {
272 pub backend: Backend,
273 pub op: &'static str,
274 pub role: &'static str,
275 pub kind: Option<QuantKind>,
276}
277
278impl Lookup {
279 pub fn new(backend: Backend, op: &'static str, kind: Option<QuantKind>) -> Self {
280 Lookup {
281 backend,
282 op,
283 role: "(dispatch)",
284 kind,
285 }
286 }
287
288 pub fn with_role(mut self, role: &'static str) -> Self {
289 self.role = role;
290 self
291 }
292}
293
294#[derive(Clone, Copy, Debug)]
296pub struct Entry {
297 pub key: Key,
298 pub outcome: Outcome,
299 pub phase: Phase,
300 pub count: u64,
302}
303
304impl std::fmt::Display for Entry {
305 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
306 match self.outcome {
307 Outcome::Hit => write!(f, "hit {} x{}", self.key, self.count),
308 Outcome::Miss {
309 fallback,
310 severity: Severity::ByDesign,
311 } => write!(f, "host {} -> {} x{}", self.key, fallback, self.count),
312 Outcome::Miss {
313 fallback,
314 severity: Severity::SlowPath,
315 } => write!(f, "MISS {} -> {} x{}", self.key, fallback, self.count),
316 }
317 }
318}
319
320#[derive(Clone, Debug, Default)]
322pub struct SealReport {
323 pub entries: Vec<Entry>,
325 pub misses: Vec<Entry>,
327 pub violations: Vec<Entry>,
333}
334
335impl SealReport {
336 pub fn render(&self) -> String {
338 let mut s = String::new();
339 for e in &self.entries {
340 s.push_str("ferrox kernels: ");
341 s.push_str(&e.to_string());
342 s.push('\n');
343 }
344 s
345 }
346
347 pub fn render_violations(&self) -> String {
350 let mut s = String::new();
351 for e in &self.violations {
352 let unit = if e.key.op.starts_with("engine.") {
355 String::new()
356 } else {
357 format!(", {} weights", e.count)
358 };
359 let kind = match e.key.kind {
360 Some(k) => format!(" {}", k.name()),
361 None => String::new(),
362 };
363 s.push_str(&format!(
364 "ferrox: NO KERNEL for {} {}{} ({}) -> falls back to {} [{}:{}{}]\n",
365 e.key.backend,
366 e.key.op,
367 kind,
368 e.key.role,
369 e.outcome.fallback().unwrap_or("(hit)"),
370 e.key.file,
371 e.key.line,
372 unit,
373 ));
374 }
375 s
376 }
377}
378
379#[derive(Clone, Debug)]
381pub struct StrictKernelError {
382 pub report: SealReport,
383}
384
385impl std::fmt::Display for StrictKernelError {
386 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
387 write!(
388 f,
389 "FERROX_STRICT_KERNELS=1 and {} kernel lookup(s) missed:\n{}",
390 self.report.violations.len(),
391 self.report.render_violations()
392 )
393 }
394}
395
396impl std::error::Error for StrictKernelError {}
397
398struct Row {
404 outcome: Outcome,
405 phase: Phase,
406 count: AtomicU64,
407}
408
409#[derive(Default)]
410struct State {
411 rows: HashMap<Key, Row>,
412 known: HashSet<(Backend, &'static str, Option<QuantKind>)>,
416 surprises: Vec<Entry>,
418}
419
420impl State {
421 fn snapshot(&self, build_only: bool) -> Vec<Entry> {
422 let mut entries: Vec<Entry> = self
423 .rows
424 .iter()
425 .filter(|(_, row)| !build_only || row.phase == Phase::Build)
426 .map(|(key, row)| Entry {
427 key: *key,
428 outcome: row.outcome,
429 phase: row.phase,
430 count: row.count.load(Ordering::Relaxed),
431 })
432 .collect();
433 entries.sort_by_key(|e| {
434 (
435 e.key.backend,
436 e.key.op,
437 e.key.kind.map(|k| k.name()).unwrap_or("f32"),
438 e.key.role,
439 e.key.line,
440 )
441 });
442 entries
443 }
444}
445
446pub struct Registry {
449 inner: RwLock<State>,
450 sealed: AtomicBool,
451}
452
453impl Default for Registry {
454 fn default() -> Self {
455 Self::new()
456 }
457}
458
459impl Registry {
460 pub fn new() -> Self {
461 Registry {
462 inner: RwLock::new(State::default()),
463 sealed: AtomicBool::new(false),
464 }
465 }
466
467 fn read(&self) -> std::sync::RwLockReadGuard<'_, State> {
468 self.inner.read().unwrap_or_else(|e| e.into_inner())
469 }
470
471 fn write(&self) -> std::sync::RwLockWriteGuard<'_, State> {
472 self.inner.write().unwrap_or_else(|e| e.into_inner())
473 }
474
475 fn bump_existing(&self, key: &Key) -> bool {
480 match self.read().rows.get(key) {
481 Some(row) => {
482 row.count.fetch_add(1, Ordering::Relaxed);
483 true
484 }
485 None => false,
486 }
487 }
488
489 pub fn is_sealed(&self) -> bool {
490 self.sealed.load(Ordering::Relaxed)
491 }
492
493 pub fn record_build_at(&self, loc: &'static Location<'static>, l: Lookup, outcome: Outcome) {
498 let key = Key {
499 backend: l.backend,
500 op: l.op,
501 role: l.role,
502 kind: l.kind,
503 file: loc.file(),
504 line: loc.line(),
505 };
506 if self.bump_existing(&key) {
507 return;
508 }
509 let mut st = self.write();
510 st.known.insert(key.shape());
511 st.rows
512 .entry(key)
513 .or_insert_with(|| Row {
514 outcome,
515 phase: Phase::Build,
516 count: AtomicU64::new(0),
517 })
518 .count
519 .fetch_add(1, Ordering::Relaxed);
520 }
521
522 pub fn record_at(&self, loc: &'static Location<'static>, l: Lookup, outcome: Outcome) {
527 let key = Key {
528 backend: l.backend,
529 op: l.op,
530 role: l.role,
531 kind: l.kind,
532 file: loc.file(),
533 line: loc.line(),
534 };
535 if self.bump_existing(&key) {
539 return;
540 }
541 let sealed = self.is_sealed();
542 let mut st = self.write();
543 if let Some(row) = st.rows.get(&key) {
546 row.count.fetch_add(1, Ordering::Relaxed);
547 return;
548 }
549 let phase = if sealed { Phase::Run } else { Phase::Build };
550 st.rows.insert(
551 key,
552 Row {
553 outcome,
554 phase,
555 count: AtomicU64::new(1),
556 },
557 );
558 if !sealed {
559 st.known.insert(key.shape());
560 return;
561 }
562 if !outcome.is_slow_path() || st.known.contains(&key.shape()) {
563 return;
564 }
565 st.surprises.push(Entry {
566 key,
567 outcome,
568 phase: Phase::Run,
569 count: 1,
570 });
571 drop(st);
572 let fallback = outcome
573 .fallback()
574 .unwrap_or("(unknown)" );
575 eprintln!(
576 "ferrox: SILENT SLOW PATH — kernel lookup missed after the model was sealed.\n\
577 ferrox: {} {} for {} has no kernel; falling back to {}.\n\
578 ferrox: call site {}:{} (role {}).\n\
579 ferrox: this was not predicted at load time, so no startup diagnostic covered it.\n\
580 ferrox: set FERROX_STRICT_KERNELS=1 to make this a hard error.",
581 key.backend,
582 key.op,
583 key.kind_name(),
584 fallback,
585 key.file,
586 key.line,
587 key.role,
588 );
589 }
590
591 pub fn seal(&self) -> SealReport {
595 self.sealed.store(true, Ordering::Relaxed);
596 let entries = self.read().snapshot(true);
597 let misses: Vec<Entry> = entries
598 .iter()
599 .copied()
600 .filter(|e| e.outcome.is_miss())
601 .collect();
602 let violations: Vec<Entry> = misses
603 .iter()
604 .copied()
605 .filter(|e| e.key.backend.is_accelerator() && e.outcome.is_slow_path())
606 .collect();
607 SealReport {
608 entries,
609 misses,
610 violations,
611 }
612 }
613
614 pub fn surprises(&self) -> Vec<Entry> {
616 self.read().surprises.clone()
617 }
618
619 pub fn entries(&self) -> Vec<Entry> {
621 self.read().snapshot(false)
622 }
623}
624
625static GLOBAL: OnceLock<Registry> = OnceLock::new();
626
627pub fn global() -> &'static Registry {
629 GLOBAL.get_or_init(Registry::new)
630}
631
632pub fn enabled() -> bool {
638 static V: OnceLock<bool> = OnceLock::new();
639 *V.get_or_init(|| {
640 !matches!(
641 std::env::var("FERROX_KERNEL_REGISTRY").ok().as_deref(),
642 Some("0") | Some("false") | Some("off")
643 )
644 })
645}
646
647pub fn verbose() -> bool {
649 static V: OnceLock<bool> = OnceLock::new();
650 *V.get_or_init(|| {
651 matches!(
652 std::env::var("FERROX_KERNEL_REGISTRY").ok().as_deref(),
653 Some("1") | Some("true") | Some("on") | Some("verbose")
654 )
655 })
656}
657
658pub fn strict() -> bool {
661 static V: OnceLock<bool> = OnceLock::new();
662 *V.get_or_init(|| {
663 matches!(
664 std::env::var("FERROX_STRICT_KERNELS").ok().as_deref(),
665 Some("1") | Some("true") | Some("on")
666 )
667 })
668}
669
670#[track_caller]
672pub fn record_build(l: Lookup, outcome: Outcome) {
673 if !enabled() {
674 return;
675 }
676 global().record_build_at(Location::caller(), l, outcome);
677}
678
679#[track_caller]
683pub fn hit(l: Lookup) {
684 if !enabled() {
685 return;
686 }
687 global().record_at(Location::caller(), l, Outcome::Hit);
688}
689
690#[track_caller]
696pub fn miss(l: Lookup, fallback: &'static str) {
697 if !enabled() {
698 return;
699 }
700 global().record_at(Location::caller(), l, Outcome::slow_path(fallback));
701}
702
703#[track_caller]
707pub fn miss_by_design(l: Lookup, fallback: &'static str) {
708 if !enabled() {
709 return;
710 }
711 global().record_at(Location::caller(), l, Outcome::by_design(fallback));
712}
713
714pub fn seal() -> SealReport {
718 let report = global().seal();
719 if !enabled() {
720 return report;
721 }
722 if verbose() {
723 eprint!("{}", report.render());
724 }
725 if !report.violations.is_empty() && !strict() {
726 eprint!("{}", report.render_violations());
727 eprintln!(
728 "ferrox: {} kernel lookup(s) above will run on a slower path than the \
729 selected backend. Set FERROX_STRICT_KERNELS=1 to refuse to run instead.",
730 report.violations.len()
731 );
732 }
733 report
734}
735
736pub fn seal_or_error() -> Result<SealReport, StrictKernelError> {
739 let report = seal();
740 if strict() && !report.violations.is_empty() {
741 return Err(StrictKernelError { report });
742 }
743 Ok(report)
744}
745
746#[cfg(test)]
747mod tests {
748 use super::*;
749
750 fn lookup(backend: Backend, kind: Option<QuantKind>) -> Lookup {
751 Lookup::new(backend, op::GEMM_PREFILL, kind).with_role("ffn_down")
752 }
753
754 #[test]
755 fn a_build_hit_is_recorded_once_per_shape_with_a_count() {
756 let r = Registry::new();
757 let loc = Location::caller();
758 for _ in 0..5 {
759 r.record_build_at(
760 loc,
761 lookup(Backend::Metal, Some(QuantKind::Q4K)),
762 Outcome::Hit,
763 );
764 }
765 let report = r.seal();
766 assert_eq!(report.entries.len(), 1);
767 assert_eq!(report.entries[0].count, 5);
768 assert!(report.misses.is_empty());
769 assert!(report.violations.is_empty());
770 }
771
772 #[test]
776 fn a_quantized_weight_with_no_accelerator_kernel_is_a_violation() {
777 let r = Registry::new();
778 r.record_build_at(
779 Location::caller(),
780 lookup(Backend::Metal, Some(QuantKind::IQ4XS)),
781 Outcome::slow_path("CPU apply_batch"),
782 );
783 let report = r.seal();
784 assert_eq!(report.violations.len(), 1);
785 assert_eq!(report.violations[0].key.kind, Some(QuantKind::IQ4XS));
786 assert!(report.render_violations().contains("IQ4_XS"));
787 }
788
789 #[test]
794 fn an_f32_miss_is_reported_but_is_not_a_violation() {
795 let r = Registry::new();
796 r.record_build_at(
797 Location::caller(),
798 lookup(Backend::Metal, None),
799 Outcome::by_design("host GEMV"),
800 );
801 let report = r.seal();
802 assert_eq!(report.misses.len(), 1);
803 assert!(report.violations.is_empty());
804 }
805
806 #[test]
809 fn a_cpu_backend_miss_is_not_a_violation() {
810 let r = Registry::new();
811 r.record_build_at(
812 Location::caller(),
813 lookup(Backend::Cpu, Some(QuantKind::IQ2XXS)),
814 Outcome::slow_path("f32 dequant-dot"),
815 );
816 assert!(r.seal().violations.is_empty());
817 }
818
819 #[test]
820 fn a_post_seal_miss_on_a_predicted_shape_is_not_a_surprise() {
821 let r = Registry::new();
822 let loc = Location::caller();
823 r.record_build_at(
824 loc,
825 lookup(Backend::Metal, Some(QuantKind::IQ4XS)),
826 Outcome::slow_path("CPU apply_batch"),
827 );
828 r.seal();
829 r.record_at(
830 loc,
831 lookup(Backend::Metal, Some(QuantKind::IQ4XS)),
832 Outcome::slow_path("CPU apply_batch"),
833 );
834 assert!(r.surprises().is_empty(), "seal already reported this shape");
835 }
836
837 #[test]
840 fn a_post_seal_miss_on_an_unpredicted_shape_is_a_surprise_reported_once() {
841 let r = Registry::new();
842 let loc = Location::caller();
843 r.record_build_at(
844 loc,
845 lookup(Backend::Metal, Some(QuantKind::Q4K)),
846 Outcome::Hit,
847 );
848 r.seal();
849 for _ in 0..3 {
850 r.record_at(
851 loc,
852 lookup(Backend::Metal, Some(QuantKind::Q2K)),
853 Outcome::slow_path("CPU apply_batch"),
854 );
855 }
856 let surprises = r.surprises();
857 assert_eq!(surprises.len(), 1, "warned once, not once per dispatch");
858 assert_eq!(surprises[0].key.kind, Some(QuantKind::Q2K));
859 assert_eq!(surprises[0].phase, Phase::Run);
860 }
861
862 #[test]
863 fn a_post_seal_hit_is_never_a_surprise() {
864 let r = Registry::new();
865 r.seal();
866 r.record_at(
867 Location::caller(),
868 lookup(Backend::Metal, Some(QuantKind::Q4K)),
869 Outcome::Hit,
870 );
871 assert!(r.surprises().is_empty());
872 }
873
874 #[test]
877 fn build_records_after_seal_extend_the_predicted_set() {
878 let r = Registry::new();
879 let loc = Location::caller();
880 r.seal();
881 r.record_build_at(
882 loc,
883 lookup(Backend::Metal, Some(QuantKind::Q6K)),
884 Outcome::slow_path("CPU apply_batch"),
885 );
886 r.record_at(
887 loc,
888 lookup(Backend::Metal, Some(QuantKind::Q6K)),
889 Outcome::slow_path("CPU apply_batch"),
890 );
891 assert!(r.surprises().is_empty());
892 }
893
894 #[test]
898 fn concurrent_dispatch_misses_warn_once_and_count_all() {
899 let r = std::sync::Arc::new(Registry::new());
900 let loc = Location::caller();
901 r.seal();
902 std::thread::scope(|s| {
903 for _ in 0..8 {
904 let r = std::sync::Arc::clone(&r);
905 s.spawn(move || {
906 for _ in 0..250 {
907 r.record_at(
908 loc,
909 lookup(Backend::Metal, Some(QuantKind::IQ1S)),
910 Outcome::slow_path("CPU apply_batch"),
911 );
912 }
913 });
914 }
915 });
916 assert_eq!(r.surprises().len(), 1, "warned once across 8 threads");
917 let counted: u64 = r
918 .entries()
919 .iter()
920 .filter(|e| e.key.kind == Some(QuantKind::IQ1S))
921 .map(|e| e.count)
922 .sum();
923 assert_eq!(counted, 2000, "every lookup counted exactly once");
924 }
925
926 #[test]
927 fn the_report_names_the_call_site_and_the_quant_kind() {
928 let r = Registry::new();
929 r.record_build_at(
930 Location::caller(),
931 lookup(Backend::Metal, Some(QuantKind::Q5K)),
932 Outcome::slow_path("CPU apply_batch"),
933 );
934 let rendered = r.seal().render();
935 assert!(rendered.contains("Q5_K"), "{rendered}");
936 assert!(rendered.contains("kernel_registry.rs"), "{rendered}");
937 assert!(rendered.contains("ffn_down"), "{rendered}");
938 assert!(rendered.contains("CPU apply_batch"), "{rendered}");
939 }
940}