ipfrs_tensorlogic/
kernel_registry.rs

1//! `TensorKernelRegistry` — a registry of computational kernels (named tensor
2//! operations with metadata), enabling dynamic kernel lookup, versioning, and
3//! capability-based selection.
4//!
5//! # Overview
6//!
7//! The registry stores [`KernelDescriptor`] entries keyed by a monotonically
8//! increasing `u64` identifier. Callers may:
9//!
10//! - [`TensorKernelRegistry::register`] — add a new kernel and obtain its id.
11//! - [`TensorKernelRegistry::lookup`] — query by name/precision/target/tag.
12//! - [`TensorKernelRegistry::best_for`] — select the most capable kernel for a
13//!   given name+precision combination (Simd > Cpu > Gpu > Generic, then highest
14//!   version).
15//! - [`TensorKernelRegistry::remove`] / [`TensorKernelRegistry::get`] — remove
16//!   or retrieve individual entries.
17//! - [`TensorKernelRegistry::stats`] — observe aggregate counters.
18
19use std::collections::HashMap;
20
21// ---------------------------------------------------------------------------
22// KernelPrecision
23// ---------------------------------------------------------------------------
24
25/// Numeric precision of a kernel's computation.
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
27pub enum KernelPrecision {
28    /// 16-bit IEEE 754 half precision float.
29    F16,
30    /// 32-bit IEEE 754 single precision float.
31    F32,
32    /// 64-bit IEEE 754 double precision float.
33    F64,
34    /// 8-bit signed integer.
35    I8,
36    /// 32-bit signed integer.
37    I32,
38}
39
40// ---------------------------------------------------------------------------
41// KernelTarget
42// ---------------------------------------------------------------------------
43
44/// Execution target for a kernel.
45#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
46pub enum KernelTarget {
47    /// Scalar CPU implementation.
48    Cpu,
49    /// GPU (CUDA / OpenCL / Vulkan) implementation.
50    Gpu,
51    /// Vectorised CPU implementation (SIMD/AVX/NEON/…).
52    Simd,
53    /// Portable fallback — no architecture-specific instructions.
54    Generic,
55}
56
57// ---------------------------------------------------------------------------
58// KernelDescriptor
59// ---------------------------------------------------------------------------
60
61/// Metadata record for a single computational kernel.
62#[derive(Clone, Debug)]
63pub struct KernelDescriptor {
64    /// Unique stable identifier assigned by the registry.
65    pub kernel_id: u64,
66    /// Human-readable operation name, e.g. `"matmul"`, `"relu"`, `"softmax"`.
67    pub name: String,
68    /// Monotonically increasing version counter for the same operation.
69    pub version: u32,
70    /// Numeric precision used by this kernel.
71    pub precision: KernelPrecision,
72    /// Execution target for this kernel.
73    pub target: KernelTarget,
74    /// Approximate number of floating-point operations per output element.
75    pub flops_per_element: f64,
76    /// Searchable tags (e.g. `["fused"`, `"blas"`, `"experimental"]`).
77    pub tags: Vec<String>,
78}
79
80// ---------------------------------------------------------------------------
81// KernelQuery
82// ---------------------------------------------------------------------------
83
84/// Filter criteria for [`TensorKernelRegistry::lookup`].
85///
86/// All `Some` fields must match; `None` fields are ignored (wildcard).
87#[derive(Clone, Debug, Default)]
88pub struct KernelQuery {
89    /// Case-insensitive substring match against [`KernelDescriptor::name`].
90    pub name: Option<String>,
91    /// Exact match against [`KernelDescriptor::precision`].
92    pub precision: Option<KernelPrecision>,
93    /// Exact match against [`KernelDescriptor::target`].
94    pub target: Option<KernelTarget>,
95    /// Exact match against any element of [`KernelDescriptor::tags`].
96    pub tag: Option<String>,
97}
98
99// ---------------------------------------------------------------------------
100// KernelRegistryStats
101// ---------------------------------------------------------------------------
102
103/// Aggregate statistics for a [`TensorKernelRegistry`].
104#[derive(Clone, Debug, Default)]
105pub struct KernelRegistryStats {
106    /// Total number of kernels currently registered.
107    pub total_kernels: usize,
108    /// Number of registered kernels per [`KernelTarget`].
109    pub by_target: HashMap<KernelTarget, usize>,
110    /// Number of registered kernels per [`KernelPrecision`].
111    pub by_precision: HashMap<KernelPrecision, usize>,
112    /// Running count of [`lookup`](TensorKernelRegistry::lookup) and
113    /// [`best_for`](TensorKernelRegistry::best_for) calls.
114    pub total_lookups: u64,
115}
116
117// ---------------------------------------------------------------------------
118// TensorKernelRegistry
119// ---------------------------------------------------------------------------
120
121/// Registry of computational kernels with dynamic lookup and versioning.
122pub struct TensorKernelRegistry {
123    kernels: HashMap<u64, KernelDescriptor>,
124    next_id: u64,
125    stats: KernelRegistryStats,
126}
127
128impl Default for TensorKernelRegistry {
129    fn default() -> Self {
130        Self::new()
131    }
132}
133
134impl TensorKernelRegistry {
135    /// Create an empty registry.
136    pub fn new() -> Self {
137        Self {
138            kernels: HashMap::new(),
139            next_id: 1,
140            stats: KernelRegistryStats::default(),
141        }
142    }
143
144    // -----------------------------------------------------------------------
145    // Mutation helpers
146    // -----------------------------------------------------------------------
147
148    /// Rebuild the `by_target` and `by_precision` histogram maps from scratch.
149    ///
150    /// Called after every register/remove to keep stats consistent.
151    fn rebuild_histograms(&mut self) {
152        let mut by_target: HashMap<KernelTarget, usize> = HashMap::new();
153        let mut by_precision: HashMap<KernelPrecision, usize> = HashMap::new();
154        for desc in self.kernels.values() {
155            *by_target.entry(desc.target).or_insert(0) += 1;
156            *by_precision.entry(desc.precision).or_insert(0) += 1;
157        }
158        self.stats.by_target = by_target;
159        self.stats.by_precision = by_precision;
160        self.stats.total_kernels = self.kernels.len();
161    }
162
163    // -----------------------------------------------------------------------
164    // Public API
165    // -----------------------------------------------------------------------
166
167    /// Register a new kernel and return its assigned `kernel_id`.
168    ///
169    /// # Arguments
170    ///
171    /// * `name`             — Operation name (e.g. `"matmul"`).
172    /// * `version`          — Caller-supplied version number.
173    /// * `precision`        — Numeric precision.
174    /// * `target`           — Execution target.
175    /// * `flops_per_element`— Approximate FLOPs per output element.
176    /// * `tags`             — Searchable tag strings.
177    pub fn register(
178        &mut self,
179        name: String,
180        version: u32,
181        precision: KernelPrecision,
182        target: KernelTarget,
183        flops_per_element: f64,
184        tags: Vec<String>,
185    ) -> u64 {
186        let kernel_id = self.next_id;
187        self.next_id += 1;
188
189        let descriptor = KernelDescriptor {
190            kernel_id,
191            name,
192            version,
193            precision,
194            target,
195            flops_per_element,
196            tags,
197        };
198        self.kernels.insert(kernel_id, descriptor);
199        self.rebuild_histograms();
200        kernel_id
201    }
202
203    /// Query the registry and return matching [`KernelDescriptor`] references.
204    ///
205    /// All non-`None` query fields must match. Results are sorted by
206    /// `(name asc, version desc, kernel_id asc)`.
207    ///
208    /// Increments [`KernelRegistryStats::total_lookups`].
209    pub fn lookup(&mut self, query: &KernelQuery) -> Vec<&KernelDescriptor> {
210        self.stats.total_lookups += 1;
211
212        let mut matches: Vec<&KernelDescriptor> = self
213            .kernels
214            .values()
215            .filter(|desc| {
216                // Name: case-insensitive substring
217                if let Some(ref name_filter) = query.name {
218                    let lc_name = desc.name.to_lowercase();
219                    let lc_filter = name_filter.to_lowercase();
220                    if !lc_name.contains(lc_filter.as_str()) {
221                        return false;
222                    }
223                }
224                // Precision: exact
225                if let Some(prec) = query.precision {
226                    if desc.precision != prec {
227                        return false;
228                    }
229                }
230                // Target: exact
231                if let Some(tgt) = query.target {
232                    if desc.target != tgt {
233                        return false;
234                    }
235                }
236                // Tag: exact match against any tag in the descriptor
237                if let Some(ref tag_filter) = query.tag {
238                    if !desc.tags.iter().any(|t| t == tag_filter) {
239                        return false;
240                    }
241                }
242                true
243            })
244            .collect();
245
246        // Sort: name asc, then version desc, then kernel_id asc
247        matches.sort_by(|a, b| {
248            a.name
249                .cmp(&b.name)
250                .then_with(|| b.version.cmp(&a.version))
251                .then_with(|| a.kernel_id.cmp(&b.kernel_id))
252        });
253
254        matches
255    }
256
257    /// Select the best kernel for `name` + `precision`.
258    ///
259    /// Target preference: `Simd > Cpu > Gpu > Generic`. Within the same target,
260    /// the highest `version` wins.
261    ///
262    /// Increments [`KernelRegistryStats::total_lookups`].
263    pub fn best_for(
264        &mut self,
265        name: &str,
266        precision: KernelPrecision,
267    ) -> Option<&KernelDescriptor> {
268        self.stats.total_lookups += 1;
269
270        // Collect IDs of candidates first to avoid borrow issues.
271        let candidate_ids: Vec<u64> = self
272            .kernels
273            .values()
274            .filter(|desc| desc.name.eq_ignore_ascii_case(name) && desc.precision == precision)
275            .map(|desc| desc.kernel_id)
276            .collect();
277
278        if candidate_ids.is_empty() {
279            return None;
280        }
281
282        // Map target to priority (lower = better).
283        let target_priority = |t: KernelTarget| match t {
284            KernelTarget::Simd => 0u8,
285            KernelTarget::Cpu => 1,
286            KernelTarget::Gpu => 2,
287            KernelTarget::Generic => 3,
288        };
289
290        // Pick the best by (target priority asc, version desc).
291        let best_id = candidate_ids.into_iter().min_by(|&a_id, &b_id| {
292            let a = &self.kernels[&a_id];
293            let b = &self.kernels[&b_id];
294            target_priority(a.target)
295                .cmp(&target_priority(b.target))
296                .then_with(|| b.version.cmp(&a.version)) // higher version first
297        });
298
299        best_id.and_then(|id| self.kernels.get(&id))
300    }
301
302    /// Remove a kernel by id.
303    ///
304    /// Returns `true` if the kernel existed and was removed, `false` otherwise.
305    pub fn remove(&mut self, kernel_id: u64) -> bool {
306        let removed = self.kernels.remove(&kernel_id).is_some();
307        if removed {
308            self.rebuild_histograms();
309        }
310        removed
311    }
312
313    /// Retrieve a kernel by id without modifying stats.
314    pub fn get(&self, kernel_id: u64) -> Option<&KernelDescriptor> {
315        self.kernels.get(&kernel_id)
316    }
317
318    /// Return a reference to the current registry statistics.
319    pub fn stats(&self) -> &KernelRegistryStats {
320        &self.stats
321    }
322}
323
324// ---------------------------------------------------------------------------
325// Tests
326// ---------------------------------------------------------------------------
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    /// Helper: register a kernel with default flops and no tags.
333    fn reg_simple(
334        r: &mut TensorKernelRegistry,
335        name: &str,
336        version: u32,
337        precision: KernelPrecision,
338        target: KernelTarget,
339    ) -> u64 {
340        r.register(name.to_string(), version, precision, target, 1.0, vec![])
341    }
342
343    /// Helper: register a kernel with tags.
344    fn reg_tagged(
345        r: &mut TensorKernelRegistry,
346        name: &str,
347        version: u32,
348        precision: KernelPrecision,
349        target: KernelTarget,
350        tags: Vec<&str>,
351    ) -> u64 {
352        r.register(
353            name.to_string(),
354            version,
355            precision,
356            target,
357            1.0,
358            tags.into_iter().map(String::from).collect(),
359        )
360    }
361
362    // -----------------------------------------------------------------------
363    // register
364    // -----------------------------------------------------------------------
365
366    #[test]
367    fn test_register_creates_kernel() {
368        let mut r = TensorKernelRegistry::new();
369        let id = reg_simple(&mut r, "matmul", 1, KernelPrecision::F32, KernelTarget::Cpu);
370        assert!(r.get(id).is_some());
371    }
372
373    #[test]
374    fn test_register_assigns_unique_ids() {
375        let mut r = TensorKernelRegistry::new();
376        let id1 = reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
377        let id2 = reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Gpu);
378        assert_ne!(id1, id2);
379    }
380
381    #[test]
382    fn test_register_increments_total_kernels() {
383        let mut r = TensorKernelRegistry::new();
384        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
385        reg_simple(
386            &mut r,
387            "sigmoid",
388            1,
389            KernelPrecision::F64,
390            KernelTarget::Gpu,
391        );
392        assert_eq!(r.stats().total_kernels, 2);
393    }
394
395    #[test]
396    fn test_register_stores_correct_metadata() {
397        let mut r = TensorKernelRegistry::new();
398        let id = r.register(
399            "softmax".to_string(),
400            3,
401            KernelPrecision::F16,
402            KernelTarget::Simd,
403            4.0,
404            vec!["fused".to_string(), "stable".to_string()],
405        );
406        let desc = r.get(id).expect("kernel must exist");
407        assert_eq!(desc.name, "softmax");
408        assert_eq!(desc.version, 3);
409        assert_eq!(desc.precision, KernelPrecision::F16);
410        assert_eq!(desc.target, KernelTarget::Simd);
411        assert!((desc.flops_per_element - 4.0).abs() < f64::EPSILON);
412        assert_eq!(desc.tags, vec!["fused", "stable"]);
413    }
414
415    // -----------------------------------------------------------------------
416    // lookup — individual filters
417    // -----------------------------------------------------------------------
418
419    #[test]
420    fn test_lookup_by_name_substring() {
421        let mut r = TensorKernelRegistry::new();
422        reg_simple(&mut r, "matmul", 1, KernelPrecision::F32, KernelTarget::Cpu);
423        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
424        let q = KernelQuery {
425            name: Some("mat".to_string()),
426            ..Default::default()
427        };
428        let results = r.lookup(&q);
429        assert_eq!(results.len(), 1);
430        assert_eq!(results[0].name, "matmul");
431    }
432
433    #[test]
434    fn test_lookup_by_name_case_insensitive() {
435        let mut r = TensorKernelRegistry::new();
436        reg_simple(&mut r, "MatMul", 1, KernelPrecision::F32, KernelTarget::Cpu);
437        let q = KernelQuery {
438            name: Some("matmul".to_string()),
439            ..Default::default()
440        };
441        let results = r.lookup(&q);
442        assert_eq!(results.len(), 1);
443    }
444
445    #[test]
446    fn test_lookup_by_precision() {
447        let mut r = TensorKernelRegistry::new();
448        reg_simple(&mut r, "conv", 1, KernelPrecision::F16, KernelTarget::Cpu);
449        reg_simple(&mut r, "conv", 1, KernelPrecision::F32, KernelTarget::Cpu);
450        reg_simple(&mut r, "conv", 1, KernelPrecision::F64, KernelTarget::Cpu);
451        let q = KernelQuery {
452            precision: Some(KernelPrecision::F32),
453            ..Default::default()
454        };
455        let results = r.lookup(&q);
456        assert_eq!(results.len(), 1);
457        assert_eq!(results[0].precision, KernelPrecision::F32);
458    }
459
460    #[test]
461    fn test_lookup_by_target_simd() {
462        let mut r = TensorKernelRegistry::new();
463        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
464        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Simd);
465        let q = KernelQuery {
466            target: Some(KernelTarget::Simd),
467            ..Default::default()
468        };
469        let results = r.lookup(&q);
470        assert_eq!(results.len(), 1);
471        assert_eq!(results[0].target, KernelTarget::Simd);
472    }
473
474    #[test]
475    fn test_lookup_by_target_gpu() {
476        let mut r = TensorKernelRegistry::new();
477        reg_simple(&mut r, "gemm", 1, KernelPrecision::F32, KernelTarget::Gpu);
478        reg_simple(&mut r, "gemm", 1, KernelPrecision::F32, KernelTarget::Cpu);
479        let q = KernelQuery {
480            target: Some(KernelTarget::Gpu),
481            ..Default::default()
482        };
483        let results = r.lookup(&q);
484        assert_eq!(results.len(), 1);
485        assert_eq!(results[0].target, KernelTarget::Gpu);
486    }
487
488    #[test]
489    fn test_lookup_by_tag_exact_match() {
490        let mut r = TensorKernelRegistry::new();
491        reg_tagged(
492            &mut r,
493            "softmax",
494            1,
495            KernelPrecision::F32,
496            KernelTarget::Cpu,
497            vec!["stable", "blas"],
498        );
499        reg_tagged(
500            &mut r,
501            "relu",
502            1,
503            KernelPrecision::F32,
504            KernelTarget::Cpu,
505            vec!["experimental"],
506        );
507        let q = KernelQuery {
508            tag: Some("stable".to_string()),
509            ..Default::default()
510        };
511        let results = r.lookup(&q);
512        assert_eq!(results.len(), 1);
513        assert_eq!(results[0].name, "softmax");
514    }
515
516    #[test]
517    fn test_lookup_tag_no_partial_match() {
518        let mut r = TensorKernelRegistry::new();
519        reg_tagged(
520            &mut r,
521            "relu",
522            1,
523            KernelPrecision::F32,
524            KernelTarget::Cpu,
525            vec!["stable"],
526        );
527        let q = KernelQuery {
528            tag: Some("stab".to_string()), // not an exact match
529            ..Default::default()
530        };
531        let results = r.lookup(&q);
532        assert!(results.is_empty());
533    }
534
535    // -----------------------------------------------------------------------
536    // lookup — combined filters
537    // -----------------------------------------------------------------------
538
539    #[test]
540    fn test_lookup_combined_name_and_precision() {
541        let mut r = TensorKernelRegistry::new();
542        reg_simple(&mut r, "matmul", 1, KernelPrecision::F32, KernelTarget::Cpu);
543        reg_simple(&mut r, "matmul", 1, KernelPrecision::F64, KernelTarget::Cpu);
544        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
545        let q = KernelQuery {
546            name: Some("matmul".to_string()),
547            precision: Some(KernelPrecision::F32),
548            ..Default::default()
549        };
550        let results = r.lookup(&q);
551        assert_eq!(results.len(), 1);
552        assert_eq!(results[0].name, "matmul");
553        assert_eq!(results[0].precision, KernelPrecision::F32);
554    }
555
556    #[test]
557    fn test_lookup_combined_precision_and_target() {
558        let mut r = TensorKernelRegistry::new();
559        reg_simple(&mut r, "conv", 1, KernelPrecision::F32, KernelTarget::Simd);
560        reg_simple(&mut r, "conv", 1, KernelPrecision::F32, KernelTarget::Gpu);
561        reg_simple(&mut r, "conv", 1, KernelPrecision::F64, KernelTarget::Simd);
562        let q = KernelQuery {
563            precision: Some(KernelPrecision::F32),
564            target: Some(KernelTarget::Simd),
565            ..Default::default()
566        };
567        let results = r.lookup(&q);
568        assert_eq!(results.len(), 1);
569        assert_eq!(results[0].target, KernelTarget::Simd);
570        assert_eq!(results[0].precision, KernelPrecision::F32);
571    }
572
573    #[test]
574    fn test_lookup_all_filters() {
575        let mut r = TensorKernelRegistry::new();
576        reg_tagged(
577            &mut r,
578            "matmul",
579            2,
580            KernelPrecision::F32,
581            KernelTarget::Simd,
582            vec!["fast"],
583        );
584        reg_tagged(
585            &mut r,
586            "matmul",
587            2,
588            KernelPrecision::F32,
589            KernelTarget::Cpu,
590            vec!["fast"],
591        );
592        reg_tagged(
593            &mut r,
594            "gemm",
595            2,
596            KernelPrecision::F32,
597            KernelTarget::Simd,
598            vec!["fast"],
599        );
600        let q = KernelQuery {
601            name: Some("matmul".to_string()),
602            precision: Some(KernelPrecision::F32),
603            target: Some(KernelTarget::Simd),
604            tag: Some("fast".to_string()),
605        };
606        let results = r.lookup(&q);
607        assert_eq!(results.len(), 1);
608        assert_eq!(results[0].name, "matmul");
609        assert_eq!(results[0].target, KernelTarget::Simd);
610    }
611
612    // -----------------------------------------------------------------------
613    // lookup — sort order
614    // -----------------------------------------------------------------------
615
616    #[test]
617    fn test_lookup_sorted_name_asc_version_desc() {
618        let mut r = TensorKernelRegistry::new();
619        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
620        reg_simple(&mut r, "matmul", 2, KernelPrecision::F32, KernelTarget::Cpu);
621        reg_simple(&mut r, "matmul", 1, KernelPrecision::F32, KernelTarget::Cpu);
622        reg_simple(&mut r, "relu", 3, KernelPrecision::F32, KernelTarget::Cpu);
623        let q = KernelQuery::default();
624        let results = r.lookup(&q);
625        assert_eq!(results.len(), 4);
626        // name asc: matmul comes before relu
627        assert_eq!(results[0].name, "matmul");
628        assert_eq!(results[1].name, "matmul");
629        // within matmul, version desc: 2 before 1
630        assert_eq!(results[0].version, 2);
631        assert_eq!(results[1].version, 1);
632        // relu entries
633        assert_eq!(results[2].name, "relu");
634        assert_eq!(results[3].name, "relu");
635        // within relu, version desc: 3 before 1
636        assert_eq!(results[2].version, 3);
637        assert_eq!(results[3].version, 1);
638    }
639
640    #[test]
641    fn test_lookup_empty_registry_returns_nothing() {
642        let mut r = TensorKernelRegistry::new();
643        let results = r.lookup(&KernelQuery::default());
644        assert!(results.is_empty());
645    }
646
647    #[test]
648    fn test_lookup_increments_total_lookups() {
649        let mut r = TensorKernelRegistry::new();
650        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
651        assert_eq!(r.stats().total_lookups, 0);
652        r.lookup(&KernelQuery::default());
653        assert_eq!(r.stats().total_lookups, 1);
654        r.lookup(&KernelQuery::default());
655        assert_eq!(r.stats().total_lookups, 2);
656    }
657
658    // -----------------------------------------------------------------------
659    // best_for
660    // -----------------------------------------------------------------------
661
662    #[test]
663    fn test_best_for_prefers_simd_over_cpu() {
664        let mut r = TensorKernelRegistry::new();
665        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
666        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Simd);
667        let best = r.best_for("relu", KernelPrecision::F32).expect("must find");
668        assert_eq!(best.target, KernelTarget::Simd);
669    }
670
671    #[test]
672    fn test_best_for_prefers_cpu_over_gpu() {
673        let mut r = TensorKernelRegistry::new();
674        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Gpu);
675        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
676        let best = r.best_for("relu", KernelPrecision::F32).expect("must find");
677        assert_eq!(best.target, KernelTarget::Cpu);
678    }
679
680    #[test]
681    fn test_best_for_prefers_gpu_over_generic() {
682        let mut r = TensorKernelRegistry::new();
683        reg_simple(
684            &mut r,
685            "conv",
686            1,
687            KernelPrecision::F32,
688            KernelTarget::Generic,
689        );
690        reg_simple(&mut r, "conv", 1, KernelPrecision::F32, KernelTarget::Gpu);
691        let best = r.best_for("conv", KernelPrecision::F32).expect("must find");
692        assert_eq!(best.target, KernelTarget::Gpu);
693    }
694
695    #[test]
696    fn test_best_for_prefers_simd_highest_priority() {
697        let mut r = TensorKernelRegistry::new();
698        reg_simple(&mut r, "op", 1, KernelPrecision::F64, KernelTarget::Generic);
699        reg_simple(&mut r, "op", 1, KernelPrecision::F64, KernelTarget::Gpu);
700        reg_simple(&mut r, "op", 1, KernelPrecision::F64, KernelTarget::Cpu);
701        reg_simple(&mut r, "op", 1, KernelPrecision::F64, KernelTarget::Simd);
702        let best = r.best_for("op", KernelPrecision::F64).expect("must find");
703        assert_eq!(best.target, KernelTarget::Simd);
704    }
705
706    #[test]
707    fn test_best_for_prefers_higher_version_same_target() {
708        let mut r = TensorKernelRegistry::new();
709        reg_simple(&mut r, "matmul", 1, KernelPrecision::F32, KernelTarget::Cpu);
710        reg_simple(&mut r, "matmul", 5, KernelPrecision::F32, KernelTarget::Cpu);
711        reg_simple(&mut r, "matmul", 3, KernelPrecision::F32, KernelTarget::Cpu);
712        let best = r
713            .best_for("matmul", KernelPrecision::F32)
714            .expect("must find");
715        assert_eq!(best.version, 5);
716    }
717
718    #[test]
719    fn test_best_for_simd_higher_version_wins_over_simd_lower() {
720        let mut r = TensorKernelRegistry::new();
721        reg_simple(
722            &mut r,
723            "sigmoid",
724            2,
725            KernelPrecision::F32,
726            KernelTarget::Simd,
727        );
728        reg_simple(
729            &mut r,
730            "sigmoid",
731            7,
732            KernelPrecision::F32,
733            KernelTarget::Simd,
734        );
735        let best = r
736            .best_for("sigmoid", KernelPrecision::F32)
737            .expect("must find");
738        assert_eq!(best.version, 7);
739    }
740
741    #[test]
742    fn test_best_for_returns_none_when_no_name_match() {
743        let mut r = TensorKernelRegistry::new();
744        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
745        assert!(r.best_for("matmul", KernelPrecision::F32).is_none());
746    }
747
748    #[test]
749    fn test_best_for_returns_none_when_no_precision_match() {
750        let mut r = TensorKernelRegistry::new();
751        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
752        assert!(r.best_for("relu", KernelPrecision::F64).is_none());
753    }
754
755    #[test]
756    fn test_best_for_empty_registry_returns_none() {
757        let mut r = TensorKernelRegistry::new();
758        assert!(r.best_for("anything", KernelPrecision::F32).is_none());
759    }
760
761    #[test]
762    fn test_best_for_increments_total_lookups() {
763        let mut r = TensorKernelRegistry::new();
764        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
765        assert_eq!(r.stats().total_lookups, 0);
766        let _ = r.best_for("relu", KernelPrecision::F32);
767        assert_eq!(r.stats().total_lookups, 1);
768    }
769
770    // -----------------------------------------------------------------------
771    // remove
772    // -----------------------------------------------------------------------
773
774    #[test]
775    fn test_remove_returns_true_when_exists() {
776        let mut r = TensorKernelRegistry::new();
777        let id = reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
778        assert!(r.remove(id));
779    }
780
781    #[test]
782    fn test_remove_returns_false_when_not_exists() {
783        let mut r = TensorKernelRegistry::new();
784        assert!(!r.remove(9999));
785    }
786
787    #[test]
788    fn test_remove_decrements_total_kernels() {
789        let mut r = TensorKernelRegistry::new();
790        let id = reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
791        assert_eq!(r.stats().total_kernels, 1);
792        r.remove(id);
793        assert_eq!(r.stats().total_kernels, 0);
794    }
795
796    #[test]
797    fn test_remove_makes_kernel_unreachable() {
798        let mut r = TensorKernelRegistry::new();
799        let id = reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
800        r.remove(id);
801        assert!(r.get(id).is_none());
802    }
803
804    #[test]
805    fn test_remove_idempotent_second_call() {
806        let mut r = TensorKernelRegistry::new();
807        let id = reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
808        assert!(r.remove(id));
809        assert!(!r.remove(id));
810    }
811
812    // -----------------------------------------------------------------------
813    // get
814    // -----------------------------------------------------------------------
815
816    #[test]
817    fn test_get_returns_some_for_existing() {
818        let mut r = TensorKernelRegistry::new();
819        let id = reg_simple(
820            &mut r,
821            "conv",
822            1,
823            KernelPrecision::I8,
824            KernelTarget::Generic,
825        );
826        let desc = r.get(id);
827        assert!(desc.is_some());
828        assert_eq!(desc.expect("test: should succeed").name, "conv");
829    }
830
831    #[test]
832    fn test_get_returns_none_for_missing() {
833        let r = TensorKernelRegistry::new();
834        assert!(r.get(42).is_none());
835    }
836
837    // -----------------------------------------------------------------------
838    // stats — histograms
839    // -----------------------------------------------------------------------
840
841    #[test]
842    fn test_stats_by_target() {
843        let mut r = TensorKernelRegistry::new();
844        reg_simple(&mut r, "a", 1, KernelPrecision::F32, KernelTarget::Cpu);
845        reg_simple(&mut r, "b", 1, KernelPrecision::F32, KernelTarget::Cpu);
846        reg_simple(&mut r, "c", 1, KernelPrecision::F32, KernelTarget::Gpu);
847        let stats = r.stats();
848        assert_eq!(stats.by_target[&KernelTarget::Cpu], 2);
849        assert_eq!(stats.by_target[&KernelTarget::Gpu], 1);
850    }
851
852    #[test]
853    fn test_stats_by_precision() {
854        let mut r = TensorKernelRegistry::new();
855        reg_simple(&mut r, "a", 1, KernelPrecision::F32, KernelTarget::Cpu);
856        reg_simple(&mut r, "b", 1, KernelPrecision::F16, KernelTarget::Cpu);
857        reg_simple(&mut r, "c", 1, KernelPrecision::I32, KernelTarget::Cpu);
858        reg_simple(&mut r, "d", 1, KernelPrecision::I32, KernelTarget::Cpu);
859        let stats = r.stats();
860        assert_eq!(stats.by_precision[&KernelPrecision::F32], 1);
861        assert_eq!(stats.by_precision[&KernelPrecision::F16], 1);
862        assert_eq!(stats.by_precision[&KernelPrecision::I32], 2);
863    }
864
865    #[test]
866    fn test_stats_histograms_updated_on_remove() {
867        let mut r = TensorKernelRegistry::new();
868        let id = reg_simple(&mut r, "a", 1, KernelPrecision::F32, KernelTarget::Simd);
869        assert_eq!(r.stats().by_target[&KernelTarget::Simd], 1);
870        r.remove(id);
871        assert_eq!(
872            r.stats()
873                .by_target
874                .get(&KernelTarget::Simd)
875                .copied()
876                .unwrap_or(0),
877            0
878        );
879    }
880
881    #[test]
882    fn test_stats_total_lookups_combined() {
883        let mut r = TensorKernelRegistry::new();
884        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
885        r.lookup(&KernelQuery::default());
886        let _ = r.best_for("relu", KernelPrecision::F32);
887        r.lookup(&KernelQuery::default());
888        assert_eq!(r.stats().total_lookups, 3);
889    }
890
891    #[test]
892    fn test_lookup_i8_precision() {
893        let mut r = TensorKernelRegistry::new();
894        reg_simple(
895            &mut r,
896            "quantized_matmul",
897            1,
898            KernelPrecision::I8,
899            KernelTarget::Cpu,
900        );
901        reg_simple(&mut r, "relu", 1, KernelPrecision::F32, KernelTarget::Cpu);
902        let q = KernelQuery {
903            precision: Some(KernelPrecision::I8),
904            ..Default::default()
905        };
906        let results = r.lookup(&q);
907        assert_eq!(results.len(), 1);
908        assert_eq!(results[0].name, "quantized_matmul");
909    }
910}
ipfrs_tensorlogic/kernel_registry.rs

ipfrs_tensorlogic/
kernel_registry.rs
