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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
//! Allocating the tile an operand is staged into, and filling a stage from a procedural source.
use cubecl::prelude::*;
use crate::*;
/// What a stage of an operand holds, and so how its fill lands it.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub(crate) enum StageElement {
/// The values the operand serves, which are what it stores: copied as they lie.
Served,
/// The values the operand serves, decoded by the fill: from scales, a codebook, or packed
/// fields as wide as their word, which take the same shared memory staged either way.
Decoded,
/// The stored form, words of several packed fields: copied as they lie.
Stored,
}
impl StageElement {
/// Whether the stage holds the values the operand serves, rather than its words.
pub(crate) fn holds_values(self) -> bool {
self != Self::Stored
}
}
#[cube]
impl Memory {
/// Cooperatively materialize a procedural source into this plain, direct scalar tile.
/// Every unit that shares its fill ([`FillUnits`]) must own this window.
pub(crate) fn fill_procedural<T: Numeric>(
&mut self,
src: &Procedural,
#[comptime] space: Space,
) {
comptime!(assert!(
self.store.packing == Packing::Plain
&& self.projection.is_direct()
&& self.store.vector_size == 1,
"Memory::fill_procedural: procedural sources require a plain, direct scalar destination"
));
// Runtime window, so a Dynamic extent witnessed by another operand still works.
let shape = self.window.extent.clone();
let fill = comptime!(self.access.fill);
let mut dst = self.flat_mut::<T, Const<1>>();
let total = dst.shape();
let stride = FillUnits::workers(fill);
let mut i = FillUnits::worker(fill);
while i < total {
let pos = shape.unravel(i.cast::<u32>());
// TODO: masked cells use Procedural's zero fallback, not the reduction's identity.
dst.write(
i,
Vector::cast_from(src.read_masked::<T>(&pos, comptime!(space.clone()))),
);
i += stride;
}
}
/// Allocate a tile staging one `divide()` sub-tile of `operand`, laid out as `storage` and
/// served `width` wide where stated, for `owner`: the cube, or each plane a copy of its own,
/// which only a plain stage of served values is.
pub(crate) fn stage<S: Numeric>(
operand: &Tile,
#[comptime] level: Level,
#[comptime] storage: StageStorage,
#[comptime] width: Option<usize>,
#[comptime] owner: StageOwner,
) -> Tile {
match comptime!(storage.clone()) {
// The plane holds the units at the depth of one region of `level`.
StageStorage::Lines { read } => {
comptime!(owner.refuse_for_a_plane("held in the plane's units"));
Tile {
kind: TileKind::new_Lines(Lines::new(
operand,
comptime!(level.clone()),
comptime!(read),
)),
factor: Factor::none(),
place: comptime!(Placement::new(
level.child(&operand.place.space),
operand.place.depth + 1,
operand.place.levels.clone()
)),
element: elem_type_of::<S>(),
}
}
_ => Memory::stage_memory::<S>(operand, level, storage, width, owner),
}
}
/// [`stage`](Memory::stage) in shared memory.
fn stage_memory<S: Numeric>(
operand: &Tile,
#[comptime] level: Level,
#[comptime] storage: StageStorage,
#[comptime] width: Option<usize>,
#[comptime] owner: StageOwner,
) -> Tile {
let element = operand.stage_element();
let decodes = comptime!(element == StageElement::Decoded);
match comptime!(element) {
StageElement::Served | StageElement::Decoded => {
let space = comptime!(level.child(&operand.place.space));
let projection = operand.projection();
let units = operand.units();
let source_width = operand.vector_size();
let vector_size = comptime!(match width {
Some(width) => {
assert!(
source_width == 1,
"Memory::stage: a padded stage assembles its lines from scalar \
source cells, so the operand it pads must be unvectorized (it is \
served {source_width} wide)"
);
assert!(
width > 1,
"Memory::stage: a padded stage width must widen the operand's own \
1-wide lines (got {width})"
);
width
}
// A decoding fill writes the stage value by value, so its lines need only
// divide the source's loads: as wide as every region below the stage cuts
// whole along its innermost axis, or a step would start mid-line.
None if decodes => {
let innermost = space.axis_at(space.rank() - 1);
space
.leaf(operand.place.below_level(&level))
.whole_run(innermost, source_width)
}
None => source_width,
});
// A TMA-filled stage's buffer must be TMA-aligned; TMA operands are always direct.
let delivery = operand.delivery();
let alignment = comptime!(match delivery.is_tma() {
true => Delivery::TMA_STAGE_ALIGNMENT,
false => 0usize,
});
// A decoded stage is dense over the values' axes whatever the source's map.
if comptime!(projection.is_direct() || decodes) {
Memory::smem_owned::<S>(space, vector_size, storage, units, alignment, owner)
} else {
comptime!(owner.refuse_for_a_plane("gathered"));
Memory::smem_gathered::<S>(
space,
vector_size,
storage,
units,
projection,
&operand.runtime_map(),
operand.window_signed(),
operand.window_boundaries(),
owner,
)
}
}
StageElement::Stored => {
comptime!(owner.refuse_for_a_plane("kept in its stored form"));
Memory::smem_stored::<S>(operand, level, storage, owner)
}
}
}
/// [`stage`](Memory::stage) in the element the operand is stored in, for the cube `owner`
/// names, whose units fill it.
fn smem_stored<S: Numeric>(
operand: &Tile,
#[comptime] level: Level,
#[comptime] storage: StageStorage,
#[comptime] owner: StageOwner,
) -> Tile {
let space = comptime!(level.child(&operand.place.space));
let vector_size = operand.vector_size();
let units = operand.units();
// A box lands where the descriptor's swizzle span starts.
let (packing, alignment) = match &operand.kind {
TileKind::Memory(g) => comptime!((g.store.packing, 0usize)),
TileKind::TmaGmem(t) => comptime!((t.packing, Delivery::TMA_STAGE_ALIGNMENT)),
TileKind::PlaneTile(_) | TileKind::PlanePartition(_) => {
panic!("Memory::smem_stored: a fragment is not a stage source")
}
TileKind::Procedural(_) | TileKind::Lines(_) => {
panic!(
"Memory::smem_stored: a procedural tile and the plane's units are not a stage source"
)
}
};
match comptime!(packing) {
Packing::Plain => {
// The values the bulk engine lands come packed; a plain stored stage keeps the
// copy's alignment.
Memory::smem_owned::<S>(space, vector_size, storage, units, 0usize, owner)
}
packing => Memory::smem_packed::<S>(
space,
vector_size,
storage,
units,
packing,
alignment,
owner,
),
}
}
}