cubek_std/routine/
grid.rs

1use crate::Hardware;
2use crate::microkernel::{CubeFootprint, CubePlanes, KernelEstimate};
3use crate::routine::{ADMITTED_PLAN_RUNS, LaunchSettings};
4use cubecl::CubeCount;
5use cubecl::client::Client;
6use cubecl::tile::launch::Grid as LaunchGrid;
7use cubecl::tile::launch::Refusal;
8use cubecl::tile::{Launcher, Partitioning, Space};
9use cubecl_environment::sync::{Arc, SyncOnceCell};
10
11/// The cubes a launch spawns, the units each holds, and the shared memory each holds.
12///
13/// Beside a blueprint rather than inside it: it is read off the partitioning at the plane width
14/// this device runs, so a blueprint replayed on another device spawns that device's grid.
15#[derive(Debug, Clone)]
16pub struct Grid {
17    /// The cubes and their units, checked against the device as every family's launch is.
18    pub launch: LaunchSettings,
19    /// Bytes of shared memory one cube holds.
20    pub shared_memory: usize,
21    /// The launcher the first launch on this grid built, shared by its clones.
22    launcher: SettledLauncher,
23}
24
25/// A grid's launcher, built by its first launch and cloned by every later one: it reads nothing
26/// but the grid and the blueprint the grid was prepared beside, and building it settles the
27/// partitioning over the problem's extents, which costs more host time than a small product.
28#[derive(Clone, Default)]
29struct SettledLauncher(Arc<SyncOnceCell<Launcher>>);
30
31impl core::fmt::Debug for SettledLauncher {
32    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
33        f.write_str("SettledLauncher")
34    }
35}
36
37/// What one cube of a launch asks of the device.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct CubeResources {
40    /// Units in one plane: the cube's `x`.
41    pub plane_width: u32,
42    /// Bytes of shared memory one cube holds.
43    pub shared_memory: usize,
44}
45
46impl Grid {
47    /// The launcher a kernel compiled for the partitioning `settle` states runs on over the real
48    /// extents of the space it states beside it, spawning this grid: the cubes a launch's own
49    /// partitioning states, not ones cubek derives again. `settle` runs on the grid's first launch
50    /// alone: every later one clones the launcher it built.
51    ///
52    /// # Panics
53    ///
54    /// Where the device refuses the grid, which [`Grid::new`] already admitted for it.
55    pub fn launcher(
56        &self,
57        client: &Client,
58        settle: impl FnOnce() -> (Partitioning, Space),
59    ) -> Launcher {
60        self.launcher
61            .0
62            .get_or_init(|| {
63                let (kernel, concrete) = settle();
64                Launcher::new(
65                    client,
66                    kernel,
67                    &concrete,
68                    LaunchGrid::Stated {
69                        cube_count: self.launch.cube_count.clone(),
70                        cube_dim: self.launch.cube_dim,
71                    },
72                )
73                .and_then(|launcher| launcher.holding(self.shared_memory))
74                .expect(ADMITTED_PLAN_RUNS)
75            })
76            .clone()
77    }
78
79    /// The grid `partitioning` distributes to on this device, each cube holding `cube`.
80    ///
81    /// # Errors
82    ///
83    /// A cube or a grid this device does not spawn, one that spawns nothing, or a cube holding
84    /// more shared memory than the device gives one.
85    pub fn new(
86        hw: &Hardware,
87        partitioning: &Partitioning,
88        cube: CubeResources,
89    ) -> Result<Self, Refusal> {
90        let launch = LaunchSettings {
91            cube_dim: partitioning.cube_dim(cube.plane_width),
92            cube_count: partitioning.cube_count(),
93            vectors: (),
94        }
95        .validate(hw)?;
96        let available = hw.shared_memory_budget();
97        if cube.shared_memory > available {
98            return Err(Refusal::SharedMemoryPastDevice {
99                bytes: cube.shared_memory,
100                most: available,
101            });
102        }
103        Ok(Self {
104            launch,
105            shared_memory: cube.shared_memory,
106            launcher: SettledLauncher::default(),
107        })
108    }
109
110    /// Cubes the grid spawns; `None` where the count is only known at launch.
111    pub fn cubes(&self) -> Option<u64> {
112        match self.launch.cube_count {
113            CubeCount::Static(x, y, z) => Some(u64::from(x) * u64::from(y) * u64::from(z)),
114            CubeCount::Dynamic(_) => None,
115        }
116    }
117
118    /// The planes of one cube.
119    pub fn planes(&self) -> CubePlanes {
120        CubePlanes {
121            count: self.launch.cube_dim.y as usize,
122            width: self.launch.cube_dim.x as usize,
123        }
124    }
125}
126
127impl From<&Grid> for KernelEstimate {
128    /// The estimate of a kernel on `grid`, before any owner has set its part.
129    fn from(grid: &Grid) -> Self {
130        Self {
131            cubes: grid.cubes(),
132            cube: CubeFootprint {
133                units: grid.launch.units_per_cube(),
134                shared_memory: grid.shared_memory,
135            },
136            demand: None,
137            staging: None,
138            work: None,
139        }
140    }
141}
cubek_std/routine/grid.rs

cubek_std/routine/
grid.rs
