cubek_std/microkernel/
base.rs

1use core::fmt::{Debug, Display};
2use cubecl::tile::launch::{Refusal, Vectorization};
3
4use crate::Hardware;
5use crate::microkernel::{CubePlanes, KernelEstimate};
6
7/// What one owner of a work item runs.
8///
9/// A value: the choice the hardware settled, sized by its settings, holding the microkernel it is
10/// built from where it is built from one. Its scope is in its name.
11pub trait Microkernel: Debug + Clone + PartialEq + Sized {
12    /// What one owner is handed, as far as the choice of microkernel depends on it.
13    type WorkItem: Debug + Clone;
14    /// The knobs a measurement settles.
15    type Settings: Debug + Clone + Display;
16    /// How the family names its operands, for [`Vectorization`].
17    type Operand: Debug + Copy + Eq;
18    /// The part of the problem one cube covers, as much of it as this microkernel reads: `()` for
19    /// one that reads only its own tiling. Where the cubes lie is its routine's placement.
20    type CubeSpace: Debug + Clone;
21
22    /// Why [`candidates`](Self::candidates) comes back empty on a device, when it does.
23    const DECLINE: &'static str;
24
25    /// The settings worth trying on this device, its default first: its own knobs, each setting
26    /// holding the one option it picked of the microkernel it is built from. A routine built on it
27    /// lists these, each at its own default, before it varies its own knobs; a microkernel built on
28    /// it picks from them by name.
29    fn settings(hw: &Hardware) -> Vec<Self::Settings>;
30
31    /// Every microkernel this device can run `item` as, best first; empty where it can run none.
32    ///
33    /// An iterator, and one that allocates nothing: [`best`](Self::best) takes its first item on
34    /// every `prepare`, so a list built to be read once would be built on every launch.
35    fn candidates(
36        hw: &Hardware,
37        item: &Self::WorkItem,
38        settings: &Self::Settings,
39    ) -> impl Iterator<Item = Self>;
40
41    /// Whether this microkernel serves `item` on any device, and why not where it does not: what
42    /// [`best`](Self::best) reports before [`DECLINE`](Self::DECLINE), which then names what a
43    /// device lacks. Every item, where it can say nothing finer.
44    ///
45    /// # Errors
46    ///
47    /// The reason this microkernel serves no such item.
48    fn admits(_item: &Self::WorkItem) -> Result<(), Refusal> {
49        Ok(())
50    }
51
52    /// Why [`candidates`](Self::candidates) comes back empty at `settings`: [`DECLINE`](Self::DECLINE),
53    /// where no setting asks the device for something finer.
54    fn decline(_settings: &Self::Settings) -> &'static str {
55        Self::DECLINE
56    }
57
58    /// The best candidate.
59    ///
60    /// # Errors
61    ///
62    /// [`admits`](Self::admits)'s reason, then [`decline`](Self::decline)'s where this device runs
63    /// none.
64    fn best(
65        hw: &Hardware,
66        item: &Self::WorkItem,
67        settings: &Self::Settings,
68    ) -> Result<Self, Refusal> {
69        Self::admits(item)?;
70        Self::candidates(hw, item, settings)
71            .next()
72            .ok_or(Refusal::Declined {
73                why: Self::decline(settings),
74            })
75    }
76
77    /// Enqueue what the accesses this microkernel's own code makes to the tiles it is handed
78    /// require of their widths, its cube covering `cube_space`, then hand
79    /// `vectorization` to the microkernels it is built from. Where those tiles live is its
80    /// walker's to know, not its own.
81    fn vectorize(
82        &self,
83        cube_space: &Self::CubeSpace,
84        vectorization: &mut Vectorization<Self::Operand>,
85    );
86
87    /// Set in `estimate` the parts this microkernel owns of a kernel running `item`, each of its
88    /// cubes of `planes` covering `cube_space`, then hand `estimate` to the microkernel it is built from:
89    /// the staging where it places the stage's reads, the work where its units hold it.
90    fn estimate(
91        &self,
92        item: &Self::WorkItem,
93        cube_space: &Self::CubeSpace,
94        planes: CubePlanes,
95        estimate: &mut KernelEstimate,
96    );
97}
cubek_std/microkernel/base.rs

cubek_std/microkernel/
base.rs
