Skip to main content

tract_linalg/
routines.rs

1//! Cross-arch registry of the single-winner kernels: one function, one datum type, one best
2//! implementation per machine.
3//!
4//! Every element-wise activation, scalar-parameter kernel and reduction declares itself here as
5//! data, so the whole function x target matrix is enumerable on any host -- with the
6//! `foreign-inventory` feature, including the trees this build cannot run. [`best_for`] answers
7//! which one a machine would use, and it takes the machine as an argument rather than reading
8//! the host, so the same query serves dispatch and introspection.
9//!
10//! This is deliberately not the model [`crate::mmm_routines`] uses. A matmul has many co-valid
11//! kernels per machine and the winner depends on the shape, so mmm keeps a pool and a tier
12//! ladder. These have no shape to weigh: one kernel wins outright, by what it is written
13//! against.
14//!
15//! A tree's kernels are declared whether or not this build compiled their bodies, so a
16//! descriptor being here means "such a kernel exists", not "it runs here". Only [`best_for`]'s
17//! answer for the *native* machine may be executed; anything else is metadata and would bail.
18//! What a build could not assemble at all is not declared, so an absent descriptor means "no
19//! such kernel", never "this toolchain skipped it".
20use crate::element_wise::{ElementWise, ElementWiseKer};
21use crate::isa::{Arch, Isa, IsaReq, IsaSet, LEVEL_BOOST};
22use crate::lut::Lut;
23use crate::reduce::{MapReduce, MapReduceKer, Reduce, ReduceKer};
24use tract_data::internal::*;
25
26/// A function a routine computes. One variant per function, whatever the datum types or the
27/// number of implementations behind it.
28#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
29pub enum Func {
30    Sigmoid,
31    Tanh,
32    Silu,
33    Gelu,
34    Erf,
35    Ln,
36    Exp,
37    Hardswish,
38    LeakyRelu,
39    MulByScalar,
40    ReduceMax,
41    ReduceMin,
42    ReduceSum,
43    Softmax2,
44    RmsNorm,
45    Lut,
46    /// A binary operation with the scalar broadcast over the whole slice.
47    BinByScalar(crate::BinOp),
48    /// A binary operation between two slices of the same length.
49    BinUnicast(crate::BinOp),
50    /// Depthwise convolution over an axis the output is contiguous on.
51    DepthwiseW,
52}
53
54impl Func {
55    /// Every binary operation, in both layouts: what [`Self::ALL`] ends with.
56    const BIN: [Func; 12] = {
57        use crate::BinOp::*;
58        [
59            Func::BinByScalar(Min),
60            Func::BinByScalar(Max),
61            Func::BinByScalar(Add),
62            Func::BinByScalar(Mul),
63            Func::BinByScalar(Sub),
64            Func::BinByScalar(SubF),
65            Func::BinUnicast(Min),
66            Func::BinUnicast(Max),
67            Func::BinUnicast(Add),
68            Func::BinUnicast(Mul),
69            Func::BinUnicast(Sub),
70            Func::BinUnicast(SubF),
71        ]
72    };
73
74    pub const ALL: [Func; 29] = [
75        Func::Sigmoid,
76        Func::Tanh,
77        Func::Silu,
78        Func::Gelu,
79        Func::Erf,
80        Func::Ln,
81        Func::Exp,
82        Func::Hardswish,
83        Func::ReduceMax,
84        Func::ReduceMin,
85        Func::ReduceSum,
86        Func::Softmax2,
87        Func::RmsNorm,
88        Func::Lut,
89        Func::LeakyRelu,
90        Func::MulByScalar,
91        Func::BIN[0],
92        Func::BIN[1],
93        Func::BIN[2],
94        Func::BIN[3],
95        Func::BIN[4],
96        Func::BIN[5],
97        Func::BIN[6],
98        Func::BIN[7],
99        Func::BIN[8],
100        Func::BIN[9],
101        Func::BIN[10],
102        Func::BIN[11],
103        Func::DepthwiseW,
104    ];
105
106    /// Where this function sits in [`Self::ALL`], which is what indexes the dispatch table.
107    fn slot(self) -> usize {
108        match self {
109            Func::Sigmoid => 0,
110            Func::Tanh => 1,
111            Func::Silu => 2,
112            Func::Gelu => 3,
113            Func::Erf => 4,
114            Func::Ln => 5,
115            Func::Exp => 6,
116            Func::Hardswish => 7,
117            Func::ReduceMax => 8,
118            Func::ReduceMin => 9,
119            Func::ReduceSum => 10,
120            Func::Softmax2 => 11,
121            Func::RmsNorm => 12,
122            Func::Lut => 13,
123            Func::LeakyRelu => 14,
124            Func::MulByScalar => 15,
125            Func::BinByScalar(op) => 16 + op as usize,
126            Func::BinUnicast(op) => 22 + op as usize,
127            Func::DepthwiseW => 28,
128        }
129    }
130
131    /// The name the matrix and the logs use.
132    pub fn name(&self) -> &'static str {
133        match self {
134            Func::Sigmoid => "sigmoid",
135            Func::Tanh => "tanh",
136            Func::Silu => "silu",
137            Func::Gelu => "gelu",
138            Func::Erf => "erf",
139            Func::Ln => "ln",
140            Func::Exp => "exp",
141            Func::Hardswish => "hardswish",
142            Func::LeakyRelu => "leaky_relu",
143            Func::MulByScalar => "mul_by_scalar",
144            Func::ReduceMax => "reduce_max",
145            Func::ReduceMin => "reduce_min",
146            Func::ReduceSum => "reduce_sum",
147            Func::Softmax2 => "softmax2",
148            Func::RmsNorm => "rms_norm",
149            Func::Lut => "lut",
150            Func::BinByScalar(op) => match op {
151                crate::BinOp::Min => "by_scalar_min",
152                crate::BinOp::Max => "by_scalar_max",
153                crate::BinOp::Add => "by_scalar_add",
154                crate::BinOp::Mul => "by_scalar_mul",
155                crate::BinOp::Sub => "by_scalar_sub",
156                crate::BinOp::SubF => "by_scalar_subf",
157            },
158            Func::BinUnicast(op) => match op {
159                crate::BinOp::Min => "unicast_min",
160                crate::BinOp::Max => "unicast_max",
161                crate::BinOp::Add => "unicast_add",
162                crate::BinOp::Mul => "unicast_mul",
163                crate::BinOp::Sub => "unicast_sub",
164                crate::BinOp::SubF => "unicast_subf",
165            },
166            Func::DepthwiseW => "depthwise_w",
167        }
168    }
169
170    /// The kernel this host runs for a function and datum type. An unfilled pair is an error rather
171    /// than a substitution: what a machine has no kernel for is what the matrix is there to show, and
172    /// a caller that quietly computed something else would hide it.
173    fn best_here(self, dt: DatumType) -> TractResult<&'static Routine> {
174        native_best(self, dt).with_context(|| {
175            format!("No {} kernel for {dt:?} on {:?}", self.name(), crate::isa::native())
176        })
177    }
178
179    /// The f32 kernel this host runs for `func`.
180    pub fn ew_f32(self) -> TractResult<Box<dyn ElementWise<f32>>> {
181        match self.best_here(DatumType::F32)?.factory {
182            RoutineFactory::F32(f) => Ok(f()),
183            // `Routine::dt` reads the arm, so `best_for` already filtered the datum type. The
184            // shape it cannot filter: asking a scalar-parameter routine for a plain one is a
185            // caller's mistake, not a missing kernel.
186            _ => bail!("{} is not a plain element-wise kernel", self.name()),
187        }
188    }
189
190    /// The f16 kernel this host runs for `func`.
191    pub fn ew_f16(self) -> TractResult<Box<dyn ElementWise<f16>>> {
192        match self.best_here(DatumType::F16)?.factory {
193            RoutineFactory::F16(f) => Ok(f()),
194            _ => bail!("{} is not a plain element-wise kernel", self.name()),
195        }
196    }
197
198    /// The f32 kernel this host runs for `func`, which takes a scalar parameter.
199    pub fn ew_f32_param(self) -> TractResult<Box<dyn ElementWise<f32, f32>>> {
200        match self.best_here(DatumType::F32)?.factory {
201            RoutineFactory::F32Param(f) => Ok(f()),
202            _ => bail!("{} is not a scalar-parameter kernel", self.name()),
203        }
204    }
205
206    /// The f16 kernel this host runs for `func`, which takes a scalar parameter.
207    pub fn ew_f16_param(self) -> TractResult<Box<dyn ElementWise<f16, f16>>> {
208        match self.best_here(DatumType::F16)?.factory {
209            RoutineFactory::F16Param(f) => Ok(f()),
210            _ => bail!("{} is not a scalar-parameter kernel", self.name()),
211        }
212    }
213
214    /// The f32 reduction this host runs for `func`.
215    pub fn reduce_f32(self) -> TractResult<Box<dyn Reduce<f32>>> {
216        match self.best_here(DatumType::F32)?.factory {
217            RoutineFactory::F32Reduce(f) => Ok(f()),
218            _ => bail!("{} is not a reduction", self.name()),
219        }
220    }
221
222    /// The f16 reduction this host runs for `func`.
223    pub fn reduce_f16(self) -> TractResult<Box<dyn Reduce<f16>>> {
224        match self.best_here(DatumType::F16)?.factory {
225            RoutineFactory::F16Reduce(f) => Ok(f()),
226            _ => bail!("{} is not a reduction", self.name()),
227        }
228    }
229
230    /// The f32 map-reduction this host runs for `func`.
231    pub fn map_reduce_f32(self) -> TractResult<Box<dyn MapReduce<f32, f32>>> {
232        match self.best_here(DatumType::F32)?.factory {
233            RoutineFactory::F32MapReduce(f) => Ok(f()),
234            _ => bail!("{} is not a map-reduction", self.name()),
235        }
236    }
237
238    /// The binary kernel this host runs for `func` and datum type, `None` when it has none. Unlike
239    /// the other accessors this one is optional rather than fallible: its callers rewrite a model
240    /// only when a kernel exists, and having none is an ordinary answer rather than a failure.
241    pub fn bin(self, dt: DatumType) -> Option<Box<crate::BinFn>> {
242        match native_best(self, dt)?.factory {
243            RoutineFactory::BinF32 { make, .. } | RoutineFactory::BinF16 { make, .. } => {
244                Some(make())
245            }
246            _ => None,
247        }
248    }
249}
250
251/// One depthwise output run: `len` contiguous output points, each the bias plus every
252/// `taps[t] * input[offsets[t] + i * in_stride]`. `taps` and `offsets` are one kernel tap each
253/// and are the same length; `in_stride` is the input step one output point costs, in elements.
254///
255/// # Safety
256/// `input.offset(offsets[t] + i * in_stride)` must be readable for every tap and every `i` below
257/// `len`, and `len` output points writable from `output`.
258pub type DepthwiseWF32 = unsafe fn(*const f32, *mut f32, &[f32], &[isize], f32, usize, isize);
259
260/// Builds the kernel behind a descriptor. The arm is what says which datum type the descriptor
261/// is for, so nothing repeats it as a field.
262#[allow(clippy::type_complexity)]
263pub enum RoutineFactory {
264    F32(fn() -> Box<dyn ElementWise<f32>>),
265    F16(fn() -> Box<dyn ElementWise<f16>>),
266    /// A kernel taking one scalar of its own datum type, applied to every element.
267    F32Param(fn() -> Box<dyn ElementWise<f32, f32>>),
268    F16Param(fn() -> Box<dyn ElementWise<f16, f16>>),
269    /// A kernel folding a slice to one value.
270    F32Reduce(fn() -> Box<dyn Reduce<f32>>),
271    F16Reduce(fn() -> Box<dyn Reduce<f16>>),
272    /// A kernel rewriting a slice and folding it in the same pass.
273    F32MapReduce(fn() -> Box<dyn MapReduce<f32, f32>>),
274    /// A kernel that is a plain function rather than a boxed object, so its name is a field:
275    /// there is no object to ask for one.
276    RmsNormF32 {
277        name: &'static str,
278        run: fn(&mut [f32], f32),
279    },
280    /// A kernel built around a table, which the caller owns and hands over per op.
281    LutU8 {
282        name: fn() -> &'static str,
283        make: fn(&[u8]) -> Box<dyn Lut>,
284    },
285    /// A binary kernel, over two tensor views. Type-erased like the views themselves, so the
286    /// datum type is the arm and the name a field.
287    BinF32 {
288        name: fn() -> &'static str,
289        make: fn() -> Box<crate::BinFn>,
290    },
291    BinF16 {
292        name: fn() -> &'static str,
293        make: fn() -> Box<crate::BinFn>,
294    },
295    /// A depthwise inner-loop kernel, a plain function rather than a boxed object, so its name
296    /// is a field.
297    DepthwiseWF32 {
298        name: &'static str,
299        run: DepthwiseWF32,
300    },
301}
302
303/// One kernel, enumerable uniformly on every host.
304pub struct Routine {
305    pub func: Func,
306    /// Architecture the kernel is written for, `None` for generic Rust every target builds.
307    pub arch: Option<Arch>,
308    /// What the instruction set must offer for this kernel to run at all. Runnability only:
309    /// a preference spelled here would also move the kernel in the matrix.
310    pub isa: IsaReq,
311    /// Where this kernel sits against its siblings, when the instruction set it needs does not
312    /// say it. Zero for the kernels whose ladder step already ranks them correctly; a measured
313    /// exception spells the steps it disagrees with, via [`crate::isa::peer_of`] or
314    /// [`crate::isa::NEVER_PREFERRED`].
315    pub boost: isize,
316    /// Whether it reaches an f16 answer by converting a chunk to f32, running an f32 kernel over
317    /// it and converting it back. Set by the declaration that writes the round trip, never by
318    /// hand.
319    pub round_trip: bool,
320    pub factory: RoutineFactory,
321}
322
323inventory::collect!(Routine);
324
325impl Routine {
326    pub fn dt(&self) -> DatumType {
327        match self.factory {
328            RoutineFactory::F32(_)
329            | RoutineFactory::F32Param(_)
330            | RoutineFactory::F32Reduce(_)
331            | RoutineFactory::F32MapReduce(_)
332            | RoutineFactory::RmsNormF32 { .. }
333            | RoutineFactory::DepthwiseWF32 { .. } => DatumType::F32,
334            RoutineFactory::F16(_) | RoutineFactory::F16Param(_) | RoutineFactory::F16Reduce(_) => {
335                DatumType::F16
336            }
337            RoutineFactory::LutU8 { .. } => DatumType::U8,
338            RoutineFactory::BinF32 { .. } => DatumType::F32,
339            RoutineFactory::BinF16 { .. } => DatumType::F16,
340        }
341    }
342
343    /// The kernel's own name. Read from the built object rather than declared, so it cannot
344    /// disagree with the kernel it names. Building is metadata-only work and safe anywhere;
345    /// running what it builds is not.
346    pub fn name(&self) -> &'static str {
347        match self.factory {
348            RoutineFactory::F32(f) => f().name(),
349            RoutineFactory::F16(f) => f().name(),
350            RoutineFactory::F32Param(f) => f().name(),
351            RoutineFactory::F16Param(f) => f().name(),
352            RoutineFactory::F32Reduce(f) => f().name(),
353            RoutineFactory::F16Reduce(f) => f().name(),
354            RoutineFactory::F32MapReduce(f) => f().name(),
355            RoutineFactory::RmsNormF32 { name, .. } => name,
356            RoutineFactory::DepthwiseWF32 { name, .. } => name,
357            RoutineFactory::LutU8 { name, .. } => name(),
358            RoutineFactory::BinF32 { name, .. } | RoutineFactory::BinF16 { name, .. } => name(),
359        }
360    }
361
362    /// Whether `isa` describes a machine this kernel runs on: its architecture, and every
363    /// feature it needs.
364    pub fn runnable_on(&self, isa: &IsaSet) -> bool {
365        self.arch.is_none_or(|a| Some(a) == isa.arch()) && self.isa.satisfied_by(*isa)
366    }
367
368    /// What this kernel is worth on a machine that can run it: its ladder step, plus whatever
369    /// a measurement said the step gets wrong. An arch kernel always outranks a generic one,
370    /// which is a different question and is compared before this.
371    fn preference(&self) -> isize {
372        self.isa.level() as isize * LEVEL_BOOST + self.boost
373    }
374}
375
376/// Every routine this build compiled, whichever architecture it speaks for.
377pub fn declared() -> impl Iterator<Item = &'static Routine> {
378    inventory::iter::<Routine>()
379}
380
381/// The kernel `isa` would run for this function and datum type: an architecture kernel over a
382/// generic one, then the most capable instruction set, then the name, which only settles ties
383/// `inventory`'s link order would otherwise settle differently between builds. `None` when
384/// nothing is declared for the pair at all.
385pub fn best_for(func: Func, dt: DatumType, isa: &IsaSet) -> Option<&'static Routine> {
386    declared()
387        .filter(|r| r.func == func && r.dt() == dt && r.runnable_on(isa))
388        .max_by_key(|r| (r.arch.is_some(), r.preference(), r.name()))
389}
390
391/// A cell of the matrix closed on purpose: on machines offering `isa`, this pair runs the kernel
392/// named here, and a kernel written for that instruction set would not be an improvement.
393///
394/// Declared beside the kernels of the tree it speaks for, so a build carries it exactly when it
395/// carries them, and it pins its winner by name: a machine that resolves the pair to something
396/// else is a different verdict, and needs its own settlement or none.
397pub struct Settled {
398    /// The rung this speaks for, which is also the column the matrix draws it under.
399    pub isa: Isa,
400    pub func: Func,
401    pub dt: DatumType,
402    /// The winner this keeps, by the name the kernel answers with.
403    pub kernel: &'static str,
404    /// One line, in the terms someone would need to disagree with it.
405    pub why: &'static str,
406}
407
408inventory::collect!(Settled);
409
410impl Settled {
411    /// Whether this speaks for `isa`: a machine of its architecture, offering what it names, and
412    /// resolving its pair to the kernel it pins.
413    pub fn covers(&self, isa: &IsaSet) -> bool {
414        Some(self.isa.arch()) == isa.arch()
415            && isa.has(self.isa)
416            && best_for(self.func, self.dt, isa).is_some_and(|r| r.name() == self.kernel)
417    }
418}
419
420/// Every settlement this build compiled, whichever architecture it speaks for.
421pub fn settlements() -> impl Iterator<Item = &'static Settled> {
422    inventory::iter::<Settled>()
423}
424
425/// Why what `isa` runs for this pair is the answer we mean, when a settlement says so. `None`
426/// when nothing settles it, which is what leaves such a cell a gap.
427pub fn settled_for(func: Func, dt: DatumType, isa: &IsaSet) -> Option<&'static Settled> {
428    settlements().find(|s| s.func == func && s.dt == dt && s.covers(isa))
429}
430
431/// What a machine's answer for one pair amounts to, which is what the matrix colours and what
432/// says whether a kernel is left to write.
433#[derive(Copy, Clone, PartialEq, Eq, Debug)]
434pub enum Standing {
435    /// Nothing this build declares serves the pair here.
436    Missing,
437    /// A kernel written for this machine's own rung of the ladder.
438    Dedicated,
439    /// Correct code that was not written for this machine: portable Rust, or an architecture
440    /// kernel from a rung below it.
441    Unspecialized,
442    /// Portable f16 code, whose every operation converts to f32 and back: the right answer at a
443    /// per-element cost no machine has to pay.
444    Emulated,
445    /// Whatever it runs, declared as the answer we mean. [`settled_for`] says why.
446    Settled,
447}
448
449/// What `isa` amounts to for this pair before any settlement speaks for it, which is the question
450/// a settlement answers.
451fn unsettled(func: Func, dt: DatumType, isa: &IsaSet) -> Standing {
452    let Some(routine) = best_for(func, dt, isa) else { return Standing::Missing };
453    // Portable f16 arithmetic goes through `f16`'s operators, which convert to f32 and back
454    // around every single one, whatever the machine underneath. A conversion-free portable f16
455    // kernel -- a table, or bit work -- would need saying so here.
456    if routine.arch.is_none() && dt == DatumType::F16 {
457        Standing::Emulated
458    } else if routine.round_trip && !isa.fp16_arithmetic() {
459        Standing::Settled
460    } else if routine.arch.is_some() && routine.isa.level() == isa.level() {
461        Standing::Dedicated
462    } else {
463        Standing::Unspecialized
464    }
465}
466
467/// What this machine's answer for the pair amounts to.
468pub fn standing(func: Func, dt: DatumType, isa: &IsaSet) -> Standing {
469    let answer = unsettled(func, dt, isa);
470    let settleable = matches!(answer, Standing::Unspecialized | Standing::Emulated);
471    if settleable && settled_for(func, dt, isa).is_some() { Standing::Settled } else { answer }
472}
473
474/// What every chunked round trip answers for on a machine with no f16 arithmetic: converting a
475/// chunk, computing it in f32 and converting it back is the technique, not a compromise.
476const NO_FP16_ARITHMETIC: &str = "no f16 arithmetic here: a chunk through an f32 kernel is it";
477
478/// Why this cell is closed, when it is: what a settlement declared, or the standing answer every
479/// f32 round trip carries on a machine that cannot compute in f16.
480pub fn settled_why(func: Func, dt: DatumType, isa: &IsaSet) -> Option<&'static str> {
481    if standing(func, dt, isa) != Standing::Settled {
482        return None;
483    }
484    Some(settled_for(func, dt, isa).map_or(NO_FP16_ARITHMETIC, |settled| settled.why))
485}
486
487/// The kernel this host runs for a function and datum type, resolved once. Dispatch happens per
488/// eval, so the scan over every declared routine runs at first use and the answers are kept in a
489/// flat table nothing has to hash.
490fn native_best(func: Func, dt: DatumType) -> Option<&'static Routine> {
491    const SLOTS: usize = Func::ALL.len() * 3;
492    static NATIVE: std::sync::OnceLock<[Option<&'static Routine>; SLOTS]> =
493        std::sync::OnceLock::new();
494    let dt_slot = match dt {
495        DatumType::F32 => 0,
496        DatumType::F16 => 1,
497        DatumType::U8 => 2,
498        _ => return None,
499    };
500    NATIVE.get_or_init(|| {
501        let isa = crate::isa::native();
502        let mut table = [None; SLOTS];
503        for func in Func::ALL {
504            for (dt_slot, dt) in [DatumType::F32, DatumType::F16, DatumType::U8].iter().enumerate()
505            {
506                table[func.slot() * 3 + dt_slot] = best_for(func, *dt, &isa);
507            }
508        }
509        table
510    })[func.slot() * 3 + dt_slot]
511}
512
513/// The fused row-wise RmsNorm this host runs: the row and the epsilon, in place.
514pub fn rms_norm_f32() -> TractResult<fn(&mut [f32], f32)> {
515    match Func::RmsNorm.best_here(DatumType::F32)?.factory {
516        RoutineFactory::RmsNormF32 { run, .. } => Ok(run),
517        _ => bail!("rms_norm is not a plain function"),
518    }
519}
520
521/// The depthwise inner-loop kernel this host runs, `None` where none is written. Optional rather
522/// than fallible, like [`Func::bin`]: its caller keeps the scalar loop it needs anyway for the
523/// zones no kernel covers, so a machine without one is an ordinary answer.
524pub fn depthwise_w_f32() -> Option<DepthwiseWF32> {
525    match native_best(Func::DepthwiseW, DatumType::F32)?.factory {
526        RoutineFactory::DepthwiseWF32 { run, .. } => Some(run),
527        _ => None,
528    }
529}
530
531/// The look-up table kernel this host runs, over the table the caller owns.
532pub fn lut_u8(table: &[u8]) -> TractResult<Box<dyn Lut>> {
533    match Func::Lut.best_here(DatumType::U8)?.factory {
534        RoutineFactory::LutU8 { make, .. } => Ok(make(table)),
535        _ => bail!("lut is not a table kernel"),
536    }
537}
538
539/// File the descriptor of a kernel declared elsewhere, under the leading architecture ident the
540/// `routine_*` declaration macros take, or with no ident at all for generic Rust every target
541/// builds. Those macros end here; write it directly for a kernel no shape macro emits, and it
542/// carries no test module of its own. The first argument names the factory arm, which is what
543/// says the kernel's shape and datum type; `isa` is omitted for a kernel whose architecture
544/// needs nothing extra to run it, `boost` for one its ladder step already ranks right.
545macro_rules! submit_routine {
546    (arm; $($rest:tt)*) => { submit_routine!(@ Some($crate::isa::Arch::Arm); $($rest)*); };
547    (aarch64; $($rest:tt)*) => { submit_routine!(@ Some($crate::isa::Arch::Aarch64); $($rest)*); };
548    (x86_64; $($rest:tt)*) => { submit_routine!(@ Some($crate::isa::Arch::X86_64); $($rest)*); };
549    (riscv64; $($rest:tt)*) => { submit_routine!(@ Some($crate::isa::Arch::RiscV64); $($rest)*); };
550    (wasm32; $($rest:tt)*) => { submit_routine!(@ Some($crate::isa::Arch::Wasm32Simd128); $($rest)*); };
551    (generic; $($rest:tt)*) => { submit_routine!(@ None; $($rest)*); };
552
553    ($factory:ident, $($rest:tt)*) => { submit_routine!(@ None; $factory, $($rest)*); };
554
555    // One arm per kernel shape: the trait a kernel implements decides how it is built, and
556    // nothing else here varies. The clauses are spelled out rather than forwarded as tokens
557    // because a `path` fragment may only be followed by a comma.
558    (@ $arch:expr; RmsNormF32, $func:ident, $name:literal, $run:path
559     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
560        submit_routine!(@@ $arch, $func,
561            $crate::routines::RoutineFactory::RmsNormF32 { name: $name, run: $run }
562            $(, isa($($isa),+))? $(, boost($boost))?);
563    };
564    (@ $arch:expr; DepthwiseWF32, $func:ident, $name:literal, $run:path
565     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
566        submit_routine!(@@ $arch, $func,
567            $crate::routines::RoutineFactory::DepthwiseWF32 { name: $name, run: $run }
568            $(, isa($($isa),+))? $(, boost($boost))?);
569    };
570    (@ $arch:expr; BinF32, BinByScalar($op:ident), $ker:path
571     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
572        submit_routine!(@@ $arch, BinByScalar($crate::BinOp::$op),
573            $crate::routines::RoutineFactory::BinF32 {
574                name: <$ker as $crate::element_wise::ElementWiseKer<f32, f32>>::name,
575                make: <$ker as $crate::by_scalar::ByScalarKer<f32>>::bin,
576            }
577            $(, isa($($isa),+))? $(, boost($boost))?);
578    };
579    (@ $arch:expr; BinF16, BinByScalar($op:ident), $ker:path
580     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
581        submit_routine!(@@ $arch, BinByScalar($crate::BinOp::$op),
582            $crate::routines::RoutineFactory::BinF16 {
583                name: <$ker as $crate::element_wise::ElementWiseKer<$crate::f16, $crate::f16>>::name,
584                make: <$ker as $crate::by_scalar::ByScalarKer<$crate::f16>>::bin,
585            }
586            $(, isa($($isa),+))? $(, boost($boost))?);
587    };
588    (@ $arch:expr; BinF32, BinUnicast($op:ident), $ker:path
589     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
590        submit_routine!(@@ $arch, BinUnicast($crate::BinOp::$op),
591            $crate::routines::RoutineFactory::BinF32 {
592                name: <$ker as $crate::unicast::UnicastKer<f32>>::name,
593                make: <$ker as $crate::unicast::UnicastKer<f32>>::bin,
594            }
595            $(, isa($($isa),+))? $(, boost($boost))?);
596    };
597    (@ $arch:expr; BinF16, BinUnicast($op:ident), $ker:path
598     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
599        submit_routine!(@@ $arch, BinUnicast($crate::BinOp::$op),
600            $crate::routines::RoutineFactory::BinF16 {
601                name: <$ker as $crate::unicast::UnicastKer<$crate::f16>>::name,
602                make: <$ker as $crate::unicast::UnicastKer<$crate::f16>>::bin,
603            }
604            $(, isa($($isa),+))? $(, boost($boost))?);
605    };
606    (@ $arch:expr; LutU8, $func:ident, $ker:path
607     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
608        submit_routine!(@@ $arch, $func,
609            $crate::routines::RoutineFactory::LutU8 {
610                name: <$ker as $crate::lut::LutKer>::name,
611                make: |table| $crate::routines::lut_of::<$ker>(table),
612            }
613            $(, isa($($isa),+))? $(, boost($boost))?);
614    };
615    (@ $arch:expr; F32Reduce, $func:ident, $ker:path
616     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
617        submit_routine!(@@ $arch, $func,
618            $crate::routines::RoutineFactory::F32Reduce(|| $crate::routines::reduce_of::<$ker, _>())
619            $(, isa($($isa),+))? $(, boost($boost))?);
620    };
621    (@ $arch:expr; F16Reduce, $func:ident, $ker:path
622     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
623        submit_routine!(@@ $arch, $func,
624            $crate::routines::RoutineFactory::F16Reduce(|| $crate::routines::reduce_of::<$ker, _>())
625            $(, isa($($isa),+))? $(, boost($boost))?);
626    };
627    (@ $arch:expr; F32MapReduce, $func:ident, $ker:path
628     $(, isa($($isa:ident),+))? $(, boost($boost:expr))?) => {
629        submit_routine!(@@ $arch, $func,
630            $crate::routines::RoutineFactory::F32MapReduce(
631                || $crate::routines::map_reduce_of::<$ker, _>()
632            )
633            $(, isa($($isa),+))? $(, boost($boost))?);
634    };
635    (@ $arch:expr; $factory:ident, $func:ident, $ker:path
636     $(, isa($($isa:ident),+))? $(, boost($boost:expr))? $(, round_trip($round_trip:expr))?) => {
637        submit_routine!(@@ $arch, $func,
638            $crate::routines::RoutineFactory::$factory(
639                || $crate::routines::factory_of::<$ker, _, _>()
640            )
641            $(, isa($($isa),+))? $(, boost($boost))? $(, round_trip($round_trip))?);
642    };
643
644    (@@ $arch:expr, $func:ident $(($($payload:tt)*))?, $factory:expr
645     $(, isa($($isa:ident),+))? $(, boost($boost:expr))? $(, round_trip($round_trip:expr))?) => {
646        inventory::submit! {
647            $crate::routines::Routine {
648                func: $crate::routines::Func::$func $(($($payload)*))?,
649                arch: $arch,
650                isa: $crate::isa::IsaReq::ANY $(.needing(&[$($crate::isa::Isa::$isa),+]))?,
651                boost: {
652                    #[allow(unused_mut, unused_assignments)]
653                    let mut boost = 0;
654                    $(boost = $boost;)?
655                    boost
656                },
657                round_trip: {
658                    #[allow(unused_mut, unused_assignments)]
659                    let mut round_trip = false;
660                    $(round_trip = $round_trip;)?
661                    round_trip
662                },
663                factory: $factory,
664            }
665        }
666    };
667}
668
669/// Close one cell of the matrix on purpose: the machines offering `$isa` run `$kernel` for this
670/// pair, and that is the answer we mean. Write it beside the kernels of the tree it speaks for,
671/// and say in `$why` what a kernel written for the instruction set would have to beat.
672macro_rules! settled {
673    ($isa:ident, $func:ident $(($op:ident))?, $dt:ident, $kernel:ident, $why:literal) => {
674        inventory::submit! {
675            $crate::routines::Settled {
676                isa: $crate::isa::Isa::$isa,
677                func: $crate::routines::Func::$func $(($crate::BinOp::$op))?,
678                dt: tract_data::prelude::DatumType::$dt,
679                kernel: stringify!($kernel),
680                why: $why,
681            }
682        }
683    };
684}
685
686/// A look-up table kernel over `table`, as its factory arm wants it.
687pub fn lut_of<K: crate::lut::LutKer + 'static>(table: &[u8]) -> Box<dyn Lut> {
688    Box::new(crate::lut::LutImpl::<K>::new(table))
689}
690
691/// The `red()` of a reduction kernel, as its factory arm wants it.
692pub fn reduce_of<K, T>() -> Box<dyn Reduce<T>>
693where
694    T: crate::LADatum,
695    K: ReduceKer<T> + Clone,
696{
697    K::red()
698}
699
700/// The `red()` of a map-reduction kernel, as its factory arm wants it.
701pub fn map_reduce_of<K, T>() -> Box<dyn MapReduce<T, T>>
702where
703    T: crate::LADatum,
704    K: MapReduceKer<T, T> + Clone,
705{
706    K::red()
707}
708
709pub fn factory_of<K, T, P>() -> Box<dyn ElementWise<T, P>>
710where
711    T: crate::LADatum,
712    P: Copy + Send + Sync + std::fmt::Debug + 'static + Default,
713    K: ElementWiseKer<T, P> + Clone,
714{
715    K::ew()
716}
717
718#[cfg(test)]
719mod tests {
720    use super::*;
721
722    /// The dispatch table is indexed by [`Func::slot`], so every function must own one slot
723    /// inside it, and the cache must answer what a fresh scan would.
724    #[test]
725    fn every_func_owns_a_slot() {
726        let mut slots = std::collections::HashSet::new();
727        for func in Func::ALL {
728            assert!(func.slot() < Func::ALL.len(), "{} is out of the table", func.name());
729            assert!(slots.insert(func.slot()), "{} shares a slot", func.name());
730        }
731        assert_eq!(slots.len(), Func::ALL.len());
732        let isa = crate::isa::native();
733        for func in Func::ALL {
734            for dt in [DatumType::F32, DatumType::F16, DatumType::U8] {
735                assert_eq!(
736                    native_best(func, dt).map(|r| r.name()),
737                    best_for(func, dt, &isa).map(|r| r.name()),
738                    "{} {dt:?}",
739                    func.name()
740                );
741            }
742        }
743    }
744
745    /// A declared kernel no machine would ever choose is either a mistake -- the wrong instruction
746    /// set, or a sibling that dominates it -- or one kept for its tests, which says so by
747    /// declining. Nothing else is reachable, and nothing would notice.
748    #[test]
749    fn a_kernel_nothing_can_choose_says_so() {
750        let mut chosen = std::collections::HashSet::new();
751        for isa in IsaSet::every_ladder() {
752            for func in Func::ALL {
753                for dt in [DatumType::F32, DatumType::F16, DatumType::U8] {
754                    if let Some(r) = best_for(func, dt, &isa) {
755                        chosen.insert((func, dt, r.name()));
756                    }
757                }
758            }
759        }
760        for r in declared() {
761            assert!(
762                chosen.contains(&(r.func, r.dt(), r.name())) || r.boost < 0,
763                "{} {:?} {} can never be chosen, and does not decline",
764                r.func.name(),
765                r.dt(),
766                r.name()
767            );
768        }
769    }
770
771    /// A pair with no kernel fails rather than falling back on something that computes a
772    /// different thing. f16 erf is the standing example: no tree has one, and core builds a
773    /// look-up table from the f32 kernel instead of asking for it.
774    #[test]
775    fn an_unfilled_pair_fails() {
776        let err = Func::Erf.ew_f16().unwrap_err().to_string();
777        assert!(err.starts_with("No erf kernel for F16 on "), "{err}");
778        // No tree has an f16 minimum either, and no consumer asks for one.
779        let err = Func::ReduceMin.reduce_f16().unwrap_err().to_string();
780        assert!(err.starts_with("No reduce_min kernel for F16 on "), "{err}");
781        // Asking for the wrong shape is a caller's mistake, and says so.
782        let err = Func::ReduceMax.ew_f32().unwrap_err().to_string();
783        assert_eq!(err, "reduce_max is not a plain element-wise kernel");
784    }
785    /// Whatever this machine declares, it can build and run: the registry is dispatch now, so a
786    /// pair whose accessor fails is a cell nothing would notice was dead. Also holds `best_for`
787    /// and the accessors to the same answer, they being two ways to ask one question.
788    #[test]
789    fn what_this_machine_declares_it_can_build() {
790        let isa = crate::isa::native();
791        for func in Func::ALL {
792            for dt in [DatumType::F32, DatumType::F16, DatumType::U8] {
793                let Some(routine) = best_for(func, dt, &isa) else { continue };
794                let built = match routine.factory {
795                    RoutineFactory::F32(_) => func.ew_f32().map(|k| k.name()),
796                    RoutineFactory::F16(_) => func.ew_f16().map(|k| k.name()),
797                    RoutineFactory::F32Param(_) => func.ew_f32_param().map(|k| k.name()),
798                    RoutineFactory::F16Param(_) => func.ew_f16_param().map(|k| k.name()),
799                    RoutineFactory::F32Reduce(_) => func.reduce_f32().map(|k| k.name()),
800                    RoutineFactory::F16Reduce(_) => func.reduce_f16().map(|k| k.name()),
801                    RoutineFactory::F32MapReduce(_) => func.map_reduce_f32().map(|k| k.name()),
802                    RoutineFactory::RmsNormF32 { name, .. }
803                    | RoutineFactory::DepthwiseWF32 { name, .. } => Ok(name),
804                    RoutineFactory::LutU8 { name, .. } => lut_u8(&[0u8; 256]).map(|_| name()),
805                    RoutineFactory::BinF32 { name, .. } | RoutineFactory::BinF16 { name, .. } => {
806                        func.bin(dt).map(|_| name()).ok_or_else(|| format_err!("no bin kernel"))
807                    }
808                };
809                assert_eq!(
810                    built.map_err(|e| e.to_string()),
811                    Ok(routine.name()),
812                    "{} {dt:?}",
813                    func.name()
814                );
815            }
816        }
817    }
818
819    /// A settlement answers for a cell the matrix would otherwise read as an invitation, so one
820    /// closing no such cell is either wrong about its machine or has outlived the kernel it
821    /// pinned -- someone wrote the kernel it says nobody should.
822    #[test]
823    fn every_settlement_closes_a_cell() {
824        for s in settlements() {
825            assert!(
826                IsaSet::every_ladder().any(|m| s.covers(&m)
827                    && matches!(
828                        unsettled(s.func, s.dt, &m),
829                        Standing::Unspecialized | Standing::Emulated
830                    )),
831                "{} {:?} on {} settles nothing, {} being what it keeps",
832                s.func.name(),
833                s.dt,
834                s.isa,
835                s.kernel
836            );
837        }
838    }
839
840    /// Two settlements over one cell would each answer for it, and the matrix would print
841    /// whichever `inventory` happened to link first.
842    #[test]
843    fn a_cell_is_settled_once() {
844        for isa in IsaSet::every_ladder() {
845            for func in Func::ALL {
846                for dt in [DatumType::F32, DatumType::F16, DatumType::U8] {
847                    let count = settlements()
848                        .filter(|s| s.func == func && s.dt == dt && s.covers(&isa))
849                        .count();
850                    assert!(
851                        count <= 1,
852                        "{} {dt:?} on {isa:?} is settled {count} times",
853                        func.name()
854                    );
855                }
856            }
857        }
858    }
859}