cubecl_tile/launch/arg/last_arrival.rs
1//! A box the runs of a share hold parts of, merged by the last run to arrive: the [`LastArrival`].
2
3use cubecl::prelude::*;
4
5use crate::launch::arg::adding::adding_tile;
6use crate::launch::arg::tensor::TileArgExpand;
7use crate::space::Extent;
8use crate::*;
9
10/// One output box of a [`Levels::shared_by`] share, held whole by one run or in parts by several
11/// ([`Portion::holders`]). A run holding it whole hands its sum straight on to the output. Each run
12/// holding a part parks it in a slot of its own and counts itself in at the box; the last to count
13/// adds the box's slots in the order of their runs and hands the total on to the output.
14///
15/// No run waits on another, so the merge does not depend on how the device schedules cubes. The
16/// parts are added in the order of their runs, so the total is the same bits whichever run arrives
17/// last.
18///
19/// The slots are one operand over the output's axes, in its order, of two boxes a run
20/// ([`Partitioning::arrival_slots`]): the first holds the part of the box the run starts in, the
21/// second the part of the box it ends in. A run's other boxes are its own whole, so no part needs a
22/// third. The counters are one `u32` a box ([`Partitioning::arrival_counters`]), holding zero, and
23/// left holding zero. A count publishes the run's part at device scope first, so it runs only where
24/// the runtime hands one cube's writes to another within a dispatch (`device_memory_scope`).
25#[derive(CubeType)]
26pub struct LastArrival<'a> {
27 /// The slots, as the boxes of [`Partitioning::arrival_slots`]'s space.
28 slots: Tile,
29 /// The slots written adding into what they hold, as the last run sums into the first's.
30 adding: Tile,
31 /// The slot boxes, one flat index: a run's two are `2 ยท run` and the next.
32 boxes: Walk,
33 /// This run's slot of this box.
34 own: Tile,
35 arrivals: &'a [Atomic<u32>],
36 /// This run's box of the output.
37 region: Region,
38 /// This box's counter.
39 counter: usize,
40 /// The runs holding parts of this box, from the first.
41 first: usize,
42 holders: usize,
43 /// The first run's slot of this box: its second where it started in a box before.
44 first_slot: usize,
45}
46
47#[cube]
48impl<'a> TileArg<'a> {
49 /// This operand as the slots of a [`LastArrival`] at the `i`-th box `portion` touches, under
50 /// `partitioning`, whose cube level shares its steps ([`shared_by`](Levels::shared_by)). The
51 /// operand is laid out as [`Partitioning::arrival_slots`] says for its axes.
52 pub fn last_arrival<'t, V: Size>(
53 &self,
54 partitioning: &Partitioning,
55 portion: &Portion,
56 i: usize,
57 arrivals: &'t [Atomic<u32>],
58 ) -> LastArrival<'t> {
59 let boxes_space = comptime!(partitioning.clone().arrival_slots_space(self.spec.axes()));
60 let boxes = Partitioning {
61 space: Space::with_sizes(
62 comptime!(boxes_space.space().clone()),
63 partitioning.space.sizes.clone(),
64 ),
65 levels: comptime!(boxes_space.levels().to_vec()),
66 };
67 let (first, holders) = portion.holders(i);
68 let run = portion.run();
69 let first_slot = 2 * first + select(portion.starts_in(i, first), 0usize, 1usize);
70 let slots = self.tile(&boxes);
71 let define!(Served) = self.element;
72 let walk = boxes.walk();
73 let own_box = walk.region(2 * run + select(portion.starts_in(i, run), 0usize, 1usize));
74 let own = slots.at(&own_box);
75 // The merge adds every later run's part into the first run's slot, which already holds
76 // that run's part: no write replaces.
77 let first_turn = false;
78 LastArrival::<'t> {
79 slots,
80 adding: adding_tile::<Served, V>(self, &boxes, first_turn),
81 boxes: walk,
82 own,
83 arrivals,
84 region: portion.region(i),
85 counter: portion.index(i),
86 first,
87 holders,
88 first_slot,
89 }
90 }
91}
92
93#[cube]
94impl<'a> LastArrival<'a> {
95 /// This run's slot of its box: what a sum it hands on is opened on.
96 pub fn tile(&self) -> Tile {
97 self.own.clone()
98 }
99
100 /// Hand `sum`, opened on [`tile`](LastArrival::tile), on to `out` with a cast: drained straight
101 /// into it where this run holds the box whole, parked and merged by the last run to arrive
102 /// otherwise. Every unit of the cube reaches it.
103 pub fn hand_on(&self, sum: &Tile, out: &mut Tile) {
104 let slot = self.own.projection();
105 let output = out.projection();
106 comptime!(assert!(
107 slot.logical_axes() == output.logical_axes(),
108 "LastArrival::hand_on: the slots span {:?} and the output {:?}; a sum opened on a slot \
109 drains into the output axis for axis, so the slots span the output's axes in its order",
110 slot.logical_axes(),
111 output.logical_axes()
112 ));
113 if self.holders == 1 {
114 sum.drained_into(&out.at(&self.region));
115 } else {
116 sum.drained_into(&self.tile());
117 // The part lands at device scope before the run counts itself in.
118 sync_storage();
119 let mut arrived = Shared::<[u32]>::new_slice(1usize);
120 if UNIT_POS == 0 {
121 arrived[0] = self.arrivals[self.counter].fetch_add(1u32);
122 }
123 sync_cube();
124 // Uniform across the cube: only the last run merges.
125 if arrived[0] as usize + 1 == self.holders {
126 // What the other runs parked is acquired before it is read.
127 sync_storage();
128 let total = self.boxes.region(self.first_slot);
129 let mut sum_into = self.adding.at(&total);
130 // Every run after the first starts in this box: its part is in its first slot.
131 for run in self.first + 1..self.first + self.holders {
132 sum_into.copy_from(&self.slots.at(&self.boxes.region(2 * run)));
133 }
134 // Every line of the total is in before any is read back.
135 sync_storage();
136 let mut out_box = out.at(&self.region);
137 out_box.copy_cast_from(&self.slots.at(&total));
138 if UNIT_POS == 0 {
139 self.arrivals[self.counter].store(0u32);
140 }
141 }
142 }
143 }
144}
145
146impl Partitioning {
147 /// The arrival counters a [`LastArrival`] under this partitioning counts its runs in through:
148 /// one `u32` a box of its shared cube level, holding zero.
149 pub fn arrival_counters(&self) -> usize {
150 let level = &self.levels()[0];
151 self.space()
152 .axes()
153 .map(|axis| level.tiles(self.space(), axis))
154 .product()
155 }
156
157 /// The shape of a [`LastArrival`]'s slots for an output over `axes`, in that order: two cube
158 /// boxes a run of the shared cube level, stacked along the first of `axes`.
159 pub fn arrival_slots(&self, axes: &[Axis]) -> cubecl::zspace::Shape {
160 let boxes = self.arrival_slots_space(axes);
161 cubecl::zspace::Shape::from(axes.iter().map(|&axis| boxes.space().extent(axis)))
162 }
163
164 /// The space the slots of a [`LastArrival`] for an output over `axes` are boxes of, cut by
165 /// these levels: one cube box, with the first of `axes` two boxes a run long. Every axis the
166 /// cube level cuts, a batch one included, is one box long, so the slot boxes are counted along
167 /// that first axis alone.
168 fn arrival_slots_space(&self, axes: &[Axis]) -> Partitioning {
169 let level = &self.levels()[0];
170 let runs = level.shared_by().expect(
171 "LastArrival: the cube level distributes no grid as one index; say `shared_by`",
172 );
173 let stacked = axes[0];
174 let cube_box = level.child(self.space());
175 let Extent::Static(edge) = cube_box.extent_raw(stacked) else {
176 panic!(
177 "LastArrival: the cube box's {stacked:?} extent is known only at launch; the slots \
178 stack cube boxes along it, so the cube level cuts it by a stated tile"
179 )
180 };
181 let extents: Vec<(Axis, Extent)> = cube_box
182 .axes()
183 .map(|axis| match axis == stacked {
184 true => (axis, Extent::Static(2 * runs * edge)),
185 false => (axis, cube_box.extent_raw(axis)),
186 })
187 .collect();
188 Partitioning::new(Space::from_extents(&extents), self.levels().to_vec())
189 }
190}