cubecl_tile/launch/arg/relay.rs
1//! A destination the cubes of a split contraction write in turns: the [`Relay`].
2
3use cubecl::prelude::*;
4
5use crate::launch::arg::adding::adding_tile;
6use crate::launch::arg::tensor::TileArgExpand;
7use crate::space::partition::GridCount;
8use crate::*;
9
10/// One output box, relayed by the cubes a [`Levels::across`] split hands it: each contracts its
11/// own run, then takes its turn at the box's carry in the order of the runs. The first turn writes
12/// its partial, each after adds its own into what is there a line at a time, and the last hands the
13/// sum on to the output. The carry is the [`TileArg`] the relay was opened on; the turns are passed
14/// on through `turns`, one counter a cube ([`Partitioning::relay_counters`]).
15///
16/// A cube opens its accumulator on the relayed [`tile`](Relay::tile), takes its turn with
17/// [`take`](Relay::take), drains into the tile, and passes it on with [`pass`](Relay::pass), every
18/// unit of the cube reaching both. A sum whose partials must be rescaled against one another before
19/// they add, as an online softmax's are, reads and rescales what the turns before it left
20/// ([`carried`](Relay::carried)) in its turn, then drains as any other. The relay publishes a
21/// turn's lines at device scope before passing the box on and acquires them after taking it, so it
22/// runs only where the runtime hands one cube's writes to another within a dispatch
23/// (`device_memory_scope`), and reads its turn with a `u32` atomic add, which every GPU runtime
24/// offers. A waiting cube holds its place on the device: the relay rests on a box's earlier runs
25/// having started by the time a later one waits, which a grid whose runs are neighbours gives.
26#[derive(CubeType)]
27pub struct Relay<'a> {
28 /// The carry, read by the last turn and by a turn that rescales it.
29 carry: Tile,
30 /// The carry written in turns: replaced by the first, added into by the rest.
31 sink: Tile,
32 turns: &'a [Atomic<u32>],
33 /// The cube level's walk, whose first region is this cube's box.
34 walk: Walk,
35 /// This box's counter.
36 counter: usize,
37 /// This cube's turn: its run's position along the split axis.
38 turn: usize,
39 /// The turns whose run holds a tile: the leading ones.
40 holders: usize,
41}
42
43#[cube]
44impl<'a> TileArg<'a> {
45 /// This operand as the carry of a [`Relay`]: the box the cubes splitting the contraction
46 /// along its [`across`](Levels::across) axis take their turns at, under `partitioning`, whose
47 /// first level is the cube level. `turns` arrives holding zero and is left holding zero.
48 pub fn relay<'t, V: Size>(
49 &self,
50 partitioning: &Partitioning,
51 turns: &'t [Atomic<u32>],
52 ) -> Relay<'t> {
53 let space = comptime!(partitioning.clone());
54 let split = comptime!(Relay::split(&space, &self.spec));
55 let walk = partitioning.walk();
56 let counter = Relay::counter(partitioning, &walk, split);
57 let turn = walk.position_at(split);
58 let holders = Relay::holders(partitioning, split);
59
60 let carry = self.tile(partitioning);
61 let define!(Served) = self.element;
62 let sink = adding_tile::<Served, V>(self, partitioning, turn == 0usize);
63
64 Relay::<'t> {
65 carry,
66 sink,
67 turns,
68 walk,
69 counter,
70 turn,
71 holders,
72 }
73 }
74}
75
76#[cube]
77impl<'a> Relay<'a> {
78 /// Whether this cube's run holds a tile, and so takes a turn: a cube past the last such run
79 /// holds nothing and never meets the box. Uniform across the cube.
80 pub fn has_turn(&self) -> bool {
81 self.turn < self.holders
82 }
83
84 /// This cube's box of the carry, relayed: what its accumulator opens on, starting from the
85 /// monoid's identity, and what it drains into once it has [`take`](Relay::take)n its turn. The
86 /// first turn replaces what the box held; every later one adds into it.
87 pub fn tile(&self) -> Tile {
88 self.sink.at(&self.walk.region(0usize))
89 }
90
91 /// This cube's box of the carry as the turns before this one left it, to read and to rescale
92 /// in place once this cube has [`take`](Relay::take)n its turn and before it drains. It holds
93 /// nothing on the [first](Relay::is_first) turn, whose drain replaces it.
94 pub fn carried(&self) -> Tile {
95 self.carry.at(&self.walk.region(0usize))
96 }
97
98 /// Whether this cube takes its box's first turn, whose drain replaces what the carry held.
99 /// Uniform across the cube.
100 pub fn is_first(&self) -> bool {
101 self.turn == 0
102 }
103
104 /// Whether this cube takes its box's last turn, whose carry [`pass`](Relay::pass) hands on to
105 /// the output. Uniform across the cube.
106 pub fn is_last(&self) -> bool {
107 self.turn + 1 == self.holders
108 }
109
110 /// Wait for this cube's turn at its box, and acquire what the turns before it wrote. Every
111 /// unit of the cube reaches it.
112 ///
113 /// The counter is read with `fetch_add(0)`. On Metal, an atomic load spinning on the counter
114 /// can keep returning the value it first read long after the previous turn's store has
115 /// landed, and the cube then never leaves the loop. A plain `load` serves once the Metal
116 /// backend's atomic load reads the counter coherently on every spin.
117 pub fn take(&self) {
118 if UNIT_POS == 0 {
119 loop {
120 if self.turns[self.counter].fetch_add(0u32) == self.turn as u32 {
121 break;
122 }
123 }
124 }
125 sync_storage();
126 }
127
128 /// Hand the box on: to the next turn, or, on the last, into `out` with a cast, once every
129 /// plane's lines of the carry are in. Every unit of the cube reaches it.
130 pub fn pass(&self, out: &mut Tile) {
131 let last = self.is_last();
132 // Publishes this turn's lines before the next turn, or the last's copy, reads them.
133 sync_storage();
134 if last {
135 let region = self.walk.region(0usize);
136 let mut out_box = out.at(®ion);
137 out_box.copy_cast_from(&self.carry.at(®ion));
138 }
139 if UNIT_POS == 0 {
140 // The last turn leaves the counter at zero, ready for the next launch.
141 if last {
142 self.turns[self.counter].store(0u32);
143 } else {
144 self.turns[self.counter].store(self.turn as u32 + 1);
145 }
146 }
147 }
148
149 /// The counter of this cube's box: the cube's position over every axis the cube level
150 /// distributes but the `split` one, as one mixed-radix index, the split's coordinate left at
151 /// zero.
152 fn counter(partitioning: &Partitioning, walk: &Walk, #[comptime] split: usize) -> usize {
153 let space = comptime!(partitioning.clone());
154 let level = comptime!(space.levels()[0].clone());
155 let mut counter = 0usize;
156 #[unroll]
157 for p in 0..comptime!(space.space().rank()) {
158 let axis = comptime!(space.space().axis_at(p));
159 if comptime!(level.distributes(axis)) {
160 let radix = Self::radix(partitioning, comptime!(level.clone()), p);
161 let position = match comptime!(p == split) {
162 true => 0usize.runtime(),
163 false => walk.position_at(p),
164 };
165 counter = counter.times(radix).plus(position);
166 }
167 }
168 counter
169 }
170
171 /// The turns along `split` whose run holds a tile: the leading ones, a run being
172 /// `ceil(grid / workers)` tiles.
173 fn holders(partitioning: &Partitioning, #[comptime] split: usize) -> usize {
174 let space = comptime!(partitioning.clone());
175 let level = comptime!(space.levels()[0].clone());
176 let axis = comptime!(space.space().axis_at(split));
177 let workers = comptime!(match level.count(axis) {
178 Some(Count::AllAcross(workers)) => workers,
179 _ => unreachable!("Relay::split names an axis spread across cubes"),
180 });
181 let grid = Self::grid(partitioning, comptime!(level.clone()), split);
182 match comptime!(level.cut(axis).spread) {
183 Spread::Contiguous => {
184 let run = grid
185 .plus(comptime!(workers - 1).runtime())
186 .divided_by(workers.runtime())
187 .max_with(1usize.runtime());
188 grid.plus(run.minus(1usize.runtime())).divided_by(run)
189 }
190 Spread::Interleaved => grid.min_with(workers.runtime()),
191 }
192 }
193
194 /// The instances the cube level distributes axis `p` to: its workers where spread across them,
195 /// its tile count otherwise.
196 fn radix(partitioning: &Partitioning, #[comptime] level: Level, #[comptime] p: usize) -> usize {
197 let axis = comptime!(partitioning.space().axis_at(p));
198 match comptime!(level.count(axis)) {
199 Some(Count::AllAcross(workers)) => workers.runtime(),
200 _ => Self::grid(partitioning, level, p),
201 }
202 }
203
204 /// The tiles the cube level cuts axis `p` into.
205 fn grid(partitioning: &Partitioning, #[comptime] level: Level, #[comptime] p: usize) -> usize {
206 let axis = comptime!(partitioning.space().axis_at(p));
207 match comptime!(level.grid(axis)) {
208 GridCount::Const(n) => n.runtime(),
209 GridCount::Extent(tile) => partitioning.space.count(p, tile),
210 }
211 }
212}
213
214impl Relay<'_> {
215 /// The position in `space` of the one axis a relay splits: spread across cubes by the first
216 /// level, the cube level, and not spanned by the carry. Refuses any other partitioning.
217 fn split(space: &Partitioning, spec: &TileSpec) -> usize {
218 let level = &space.levels()[0];
219 assert!(
220 level.coverage() == Coverage::Distribute(ComputeScope::Cube)
221 && level.shared_by().is_none(),
222 "TileArg::relay: the first level must hand each cube its own boxes; a relay's turns are \
223 taken by the cubes of one box"
224 );
225 let split: Vec<usize> = (0..space.space().rank())
226 .filter(|&p| {
227 let axis = space.space().axis_at(p);
228 matches!(level.count(axis), Some(Count::AllAcross(_)))
229 && !spec.axes().contains(&axis)
230 })
231 .collect();
232 assert!(
233 split.len() == 1,
234 "TileArg::relay: the carry must leave out exactly one axis the cube level spreads across \
235 cubes (`Levels::across`), the one whose runs take the turns; it leaves out {}",
236 split.len()
237 );
238 split[0]
239 }
240}
241
242impl Partitioning {
243 /// The turn counters a [`Relay`] under this partitioning passes its boxes on through: one
244 /// `u32` a cube, holding zero.
245 pub fn relay_counters(&self) -> usize {
246 self.instances(Coverage::Distribute(ComputeScope::Cube)) as usize
247 }
248}