cubecl_tile/launch/spec/
delivery.rs

1//! Who moves an operand's bytes into a stage, and whether they have landed when the fill
2//! returns: the [`Delivery`].
3
4use crate::{Refusal, Rendezvous, Space};
5
6/// Who moves an operand into a stage, and whether its bytes have landed when the fill returns.
7#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default)]
8pub enum Delivery {
9    /// Every unit loads its own lines and stores them, through its registers: landed when the fill
10    /// returns, and a line can be decoded on its way.
11    #[default]
12    SyncPerUnit,
13    /// Every unit hands its own lines to the copy engine (`cp.async`), memory to memory: they land
14    /// after the fill returns, as they lie, and something the slot waits on tracks them.
15    AsyncPerUnit,
16    /// One unit hands the whole stage to the copy engine as one contiguous run
17    /// (`cp.async.bulk`): it lands after the fill returns, counted in bytes on the slot's barrier.
18    AsyncBulk,
19    /// One unit hands the stage's box to the TMA engine through a tensor map
20    /// (`cp.async.bulk.tensor`): it lands after the fill returns, counted in bytes on the slot's
21    /// barrier.
22    Tma,
23    /// Every unit computes its values from their coordinates: nothing is read, and the stage
24    /// holds them when the fill returns.
25    Procedural,
26}
27
28/// CUDA's cap on each TMA box dimension.
29const TMA_MAX_BOX_DIM: usize = 256;
30
31impl Delivery {
32    /// The byte alignment a stage a [`Tma`](Delivery::Tma) box lands in starts on: the span a
33    /// 128-byte swizzle repeats over, which the engine keys off the address, so a stage's first
34    /// row takes the first key as the fragment reading it does. A launch sizing its shared memory
35    /// counts what aligning a stage to it can cost.
36    pub const TMA_STAGE_ALIGNMENT: usize = 1024;
37
38    /// Whether this delivery moves `stage` in one transfer, and why not where it cannot.
39    pub fn moves(self, stage: &Space) -> Result<(), Refusal> {
40        match self {
41            Delivery::SyncPerUnit
42            | Delivery::AsyncPerUnit
43            | Delivery::AsyncBulk
44            | Delivery::Procedural => Ok(()),
45            Delivery::Tma => match stage
46                .extents()
47                .into_iter()
48                .find(|&(_, edge)| edge > TMA_MAX_BOX_DIM)
49            {
50                Some((axis, edge)) => Err(Refusal::BoxPastDescriptor {
51                    axis,
52                    edge,
53                    most: TMA_MAX_BOX_DIM,
54                }),
55                None => Ok(()),
56            },
57        }
58    }
59
60    pub fn is_tma(&self) -> bool {
61        matches!(self, Delivery::Tma)
62    }
63
64    /// Whether the fill returns before its bytes land.
65    pub(crate) fn is_async(&self) -> bool {
66        match self {
67            Delivery::SyncPerUnit | Delivery::Procedural => false,
68            Delivery::AsyncPerUnit | Delivery::AsyncBulk | Delivery::Tma => true,
69        }
70    }
71
72    /// The synchronization required to materialize this source in a staging slot: an mbarrier for
73    /// every async delivery, counting a bulk copy's bytes or each unit's committed `cp.async`
74    /// copies.
75    pub(crate) fn rendezvous(&self) -> Rendezvous {
76        if self.is_async() {
77            Rendezvous::Barrier
78        } else {
79            Rendezvous::Cube
80        }
81    }
82
83    /// Whether the copy writes through the async proxy and completes on the slot's barrier as
84    /// transaction bytes (`cp.async.bulk`, TMA), so the barrier's initialization must be fenced
85    /// into that proxy before the first copy. A per-unit `cp.async` stays in the generic proxy.
86    pub(crate) fn through_async_proxy(&self) -> bool {
87        matches!(self, Delivery::AsyncBulk | Delivery::Tma)
88    }
89
90    /// Whether every unit of the cube takes part in the fill, rather than one elected issuer.
91    pub(crate) fn every_unit_fills(&self) -> bool {
92        match self {
93            Delivery::SyncPerUnit | Delivery::AsyncPerUnit | Delivery::Procedural => true,
94            Delivery::AsyncBulk | Delivery::Tma => false,
95        }
96    }
97}
cubecl_tile/launch/spec/delivery.rs

cubecl_tile/launch/spec/
delivery.rs
