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