cubecl_tile/layout/paging/
base.rs

1//! [`Paging`]: which axes of an operand a page table resolves, and what an operand read
2//! through one owes.
3
4use crate::*;
5
6/// Which axes of an operand a page table resolves: `sequence` names the page a table row lists,
7/// `position` a run of `page_size` positions within one.
8#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
9pub struct Paging {
10    pub sequence: Axis,
11    pub position: Axis,
12    /// Positions one page holds: the pool's extent along `position`.
13    pub page_size: usize,
14}
15
16impl Paging {
17    /// The paging of a pool of `shape`, its pages along the dim `sequence` labels and their
18    /// positions along the one `position` labels, a page its extent there.
19    ///
20    /// # Errors
21    ///
22    /// What [`validate`](Self::validate) refuses, and a pool of no pages.
23    pub(crate) fn of_pool(
24        sequence: Axis,
25        position: Axis,
26        shape: &[usize],
27        projection: &Projection,
28        stored_tiles: &[(Axis, usize)],
29        delivery: Delivery,
30    ) -> Result<Paging, Refusal> {
31        let unsized_paging = Paging {
32            sequence,
33            position,
34            page_size: 0,
35        };
36        // Where the axes name no dims, `validate` refuses before it reads the page.
37        let page_size = unsized_paging
38            .dims(projection)
39            .map_or(0, |(_, position)| shape[position]);
40        let paging = Paging {
41            page_size,
42            ..unsized_paging
43        };
44        paging.validate(projection, stored_tiles, delivery)?;
45        let (sequence, _) = paging.dims(projection)?;
46        match shape[sequence] {
47            0 => Err(Refusal::PagingOfAnEmptyPool),
48            _ => Ok(paging),
49        }
50    }
51
52    /// The buffer dims `projection` carries the two axes on, `(sequence, position)`.
53    ///
54    /// # Errors
55    ///
56    /// Either axis not one dim's identity, or both on one dim: a page is resolved coordinate by
57    /// coordinate.
58    pub(crate) fn dims(&self, projection: &Projection) -> Result<(usize, usize), Refusal> {
59        let dim = |axis: Axis| {
60            (0..projection.physical_rank())
61                .find(|&pa| projection.physical_axis(pa).is_identity(axis))
62                .ok_or(Refusal::PagingOfAnUnlabelledAxis(axis))
63        };
64        let (sequence, position) = (dim(self.sequence)?, dim(self.position)?);
65        match sequence == position {
66            true => Err(Refusal::PagingOfOneDim(self.sequence)),
67            false => Ok((sequence, position)),
68        }
69    }
70
71    /// What an operand read through a table owes, checked once for the build and the kernel both:
72    /// a direct, untiled mapping over a pool of pages, moved a line at a time, its lines never
73    /// along the pages. A line is whole within a page because the binding serves it whole within
74    /// the dim a page spans. A packed pool is paged like a plain one: its words lie along the
75    /// innermost dim, as its lines do.
76    ///
77    /// # Errors
78    ///
79    /// The [`Refusal`] naming the first thing the table cannot resolve.
80    pub(crate) fn validate(
81        &self,
82        projection: &Projection,
83        stored_tiles: &[(Axis, usize)],
84        delivery: Delivery,
85    ) -> Result<(), Refusal> {
86        if !projection.is_direct() || projection.is_tiled() || !stored_tiles.is_empty() {
87            return Err(Refusal::PagingOfAStatedLayout);
88        }
89        if !matches!(delivery, Delivery::SyncPerUnit | Delivery::AsyncPerUnit) {
90            return Err(Refusal::PagingMovedByRuns);
91        }
92        let (sequence, _) = self.dims(projection)?;
93        if sequence == projection.physical_rank() - 1 {
94            return Err(Refusal::PagesAlongALine(self.sequence));
95        }
96        match self.page_size {
97            0 => Err(Refusal::PagingOfAnEmptyPool),
98            _ => Ok(()),
99        }
100    }
101}
cubecl_tile/layout/paging/base.rs

cubecl_tile/layout/paging/
base.rs
