cubecl_tile/launch/spec/
base.rs

1//! The comptime half of an operand.
2
3use cubecl::zspace::SmallVec;
4
5use crate::Delivery;
6use crate::{Axis, Boundary, BoundaryPolicy, Field, Packing, Paging, Projection, Space, Storage};
7
8/// The comptime half of an operand: its axis mapping, edge checks, packing, storage and delivery.
9#[derive(Clone, PartialEq, Eq, Hash, Debug)]
10pub struct TileSpec {
11    /// How this operand's logical axes address its buffer's physical ones.
12    pub projection: Projection,
13    /// Out-of-bounds handling per coordinate axis; `None` or an empty list is unchecked.
14    pub boundaries: SmallVec<[Option<Boundary>; Space::MAX_RANK]>,
15    /// The launch's units per cube, `0` when unknown.
16    pub units: usize,
17    /// How this operand's values sit in its binding.
18    pub(crate) packing: Packing,
19    /// What this operand's storage tiles are to the windows it is read through.
20    pub storage: Storage,
21    /// How the buffer is stored: its stated storage tiles, finest first, as its
22    /// [`StoragePartitioning`](crate::StoragePartitioning) reads back. Empty for a plain buffer.
23    /// What a load of it covers is the kernel's to pick
24    /// ([`VectorTile::new`](crate::VectorTile::new)).
25    pub stored_tiles: Vec<(Axis, usize)>,
26    /// Who moves this operand into a stage ([`delivery`](Self::delivery)).
27    pub delivery: Delivery,
28    /// The axes a page table resolves where the buffer is a pool of pages read as sequences
29    /// ([`TileArgLaunch::indexed`](crate::TileArgLaunch::indexed)); `None` for a buffer read where
30    /// it lies.
31    pub paging: Option<Paging>,
32}
33
34impl TileSpec {
35    /// An operand's spec from its mapping, unchecked, plain, untiled, cube size unknown.
36    pub fn new(projection: Projection) -> Self {
37        TileSpec {
38            projection,
39            boundaries: SmallVec::new(),
40            units: 0,
41            packing: Packing::Plain,
42            storage: Storage::Strided,
43            stored_tiles: Vec::new(),
44            delivery: Delivery::SyncPerUnit,
45            paging: None,
46        }
47    }
48
49    /// [`new`](Self::new) over the [`direct`](Projection::direct) mapping.
50    pub fn direct(axes: &[Axis]) -> Self {
51        Self::new(Projection::direct(axes))
52    }
53
54    /// The axes this operand spans, in its logical order.
55    pub fn axes(&self) -> &[Axis] {
56        self.projection.logical_axes()
57    }
58
59    /// This spec with one boundary mode per coordinate axis; panics on any other count.
60    pub fn boundaries(mut self, boundaries: &[Option<Boundary>]) -> Self {
61        let coord_rank = self.projection.coordinate_rank();
62        assert!(
63            boundaries.len() == coord_rank,
64            "TileSpec::boundaries: {} modes stated but the operand has {coord_rank} coordinate \
65             axes",
66            boundaries.len()
67        );
68        self.boundaries = match boundaries.iter().any(Option::is_some) {
69            true => SmallVec::from_slice(boundaries),
70            false => SmallVec::new(),
71        };
72        self
73    }
74
75    /// This spec checked as `policy` says.
76    pub fn boundary(self, policy: BoundaryPolicy) -> Self {
77        let coord_rank = self.projection.coordinate_rank();
78        match policy {
79            BoundaryPolicy::Unchecked => self.boundaries(&vec![None; coord_rank]),
80            // A spec has no launch to say which axes some extent could overhang: every one.
81            BoundaryPolicy::Every(boundary) | BoundaryPolicy::AnyExtent(boundary) => {
82                self.boundaries(&vec![Some(boundary); coord_rank])
83            }
84        }
85    }
86
87    /// This spec with its binding holding `u32` words of `field`-wide values.
88    pub fn packed(mut self, field: impl Into<Field>) -> Self {
89        self.packing = Packing::Packed {
90            field: field.into(),
91        };
92        self
93    }
94
95    /// This spec with its stages moved by `delivery`: each unit its own lines, landed when the fill
96    /// returns or after it, or one unit the whole stage in bulk. The kernel does not change; the
97    /// fill does.
98    ///
99    /// # Panics
100    ///
101    /// [`Tma`](Delivery::Tma), which moves a tensor map rather than a tensor, and
102    /// [`Procedural`](Delivery::Procedural), which moves nothing.
103    pub fn delivery(mut self, delivery: Delivery) -> Self {
104        assert!(
105            matches!(
106                delivery,
107                Delivery::SyncPerUnit | Delivery::AsyncPerUnit | Delivery::AsyncBulk
108            ),
109            "TileSpec::delivery: a tensor is moved per unit or in bulk; {delivery:?} is a tensor \
110             map's or a coordinate tile's"
111        );
112        self.delivery = delivery;
113        self
114    }
115
116    /// Whether this operand has bounds checking enabled.
117    pub fn is_checked(&self) -> bool {
118        self.boundaries.iter().any(Option::is_some)
119    }
120}
cubecl_tile/launch/spec/base.rs

cubecl_tile/launch/spec/
base.rs
