cubecl_tile/launch/arg/
destination.rs

1//! The [`Destination`] seam, through which a kernel stores into [`Output`] or a caller's own
2//! argument.
3
4use cubecl::prelude::*;
5
6use crate::{
7    Bound, Output, OutputArgs, Partitioning, Refusal, Tile, TileArg, TileArgLaunch, TmaOperand,
8    kind::Write, launch::Unbound,
9};
10
11/// Which launch argument carries a kernel's output and how the kernel serves it as a [`Tile`].
12#[cube]
13pub trait Destination: Send + Sync + 'static {
14    /// The launch argument carrying the output and its spec, at element `E`.
15    type Arg<E: Numeric>: LaunchArg + CubeType;
16
17    /// Serve the argument as a [`Tile`] under the kernel's one `partitioning`, stored in lines of
18    /// `V` where the argument does not state its own width.
19    fn tile<E: Numeric, V: Size>(arg: &Self::Arg<E>, partitioning: &Partitioning) -> Tile;
20}
21
22/// [`Destination`]'s host half: how what the launch built becomes the argument.
23pub trait DestinationLaunch: Destination {
24    /// What the argument is made from, before the kernel's element types it.
25    type Operand;
26
27    fn arg<E: Numeric>(operand: Self::Operand) -> <Self::Arg<E> as LaunchArg>::RuntimeArg;
28}
29
30/// A kernel prepared to store through a [`Destination`] the caller builds afterwards, at
31/// [`width`](Self::width): the caller binds the output only once it knows the kernel's spec.
32pub trait SinkLaunch {
33    /// The width the kernel hands the destination its lines in.
34    fn width(&self) -> usize;
35
36    /// The dense output's element count, which a destination standing for it must cover exactly.
37    fn elements(&self) -> usize;
38
39    /// Launch through the destination `O`, `sink` building the output's argument from its derived
40    /// spec and geometry.
41    fn launch<O: DestinationLaunch>(self, sink: impl FnOnce(Unbound) -> O::Operand);
42}
43
44/// What a kernel named by a value rather than a type is handed to, once prepared for a sink: each
45/// value prepares a launch of its own type, so what receives it takes any [`SinkLaunch`].
46pub trait SinkLaunchVisitor {
47    type Output;
48
49    /// Receive the launch the value named, or why it could not be prepared.
50    fn visit<L: SinkLaunch>(self, prepared: Result<L, Refusal>) -> Self::Output;
51}
52
53/// The default [`Destination`]: an [`Output`] buffer.
54pub struct Buffered;
55
56#[cube]
57impl Destination for Buffered {
58    type Arg<E: Numeric> = Output<'static, E>;
59
60    fn tile<E: Numeric, V: Size>(arg: &Self::Arg<E>, partitioning: &Partitioning) -> Tile {
61        arg.tile(partitioning)
62    }
63}
64
65/// [`Write::Replace`] binds an [`Output::Tensor`], [`Write::Accumulate`] an [`Output::Atomic`].
66impl DestinationLaunch for Buffered {
67    type Operand = (Bound, Write);
68
69    fn arg<E: Numeric>((bound, write): (Bound, Write)) -> OutputArgs<'static, E> {
70        match write {
71            Write::Replace => bound.output(),
72            Write::Accumulate => bound.atomic(),
73            Write::Exclusive(_) => panic!(
74                "Buffered: a relayed output is opened in the kernel with `TileArg::relay`, which \
75                 takes its turn counters beside it; bind the carry as a plain `TileArg`"
76            ),
77        }
78    }
79}
80
81/// A [`Destination`] the TMA engine stores: the output is a tensor map, served as a tile a
82/// shared-memory box is bulk-stored into ([`Tile::smem_box`], [`Tile::copy_from`]), a box at a
83/// time, which the descriptor clips at the tensor's edge. Its width is the descriptor's cell, so
84/// a launch binds it one element wide.
85pub struct TmaStored;
86
87#[cube]
88impl Destination for TmaStored {
89    type Arg<E: Numeric> = TileArg<'static>;
90
91    fn tile<E: Numeric, V: Size>(arg: &Self::Arg<E>, partitioning: &Partitioning) -> Tile {
92        arg.tile(partitioning)
93    }
94}
95
96impl DestinationLaunch for TmaStored {
97    type Operand = TmaOperand;
98
99    fn arg<E: Numeric>(operand: TmaOperand) -> TileArgLaunch {
100        operand.arg()
101    }
102}
cubecl_tile/launch/arg/destination.rs

cubecl_tile/launch/arg/
destination.rs
