cubecl_tile/launch/arg/
output.rs

1//! What a kernel leaves by: [`Output`].
2
3use cubecl::prelude::*;
4
5use crate::{AccumulateArg, AccumulateArgLaunch, Bound, Partitioning, Tile, TileArg};
6
7/// The operand a kernel leaves by: a tensor read and replaced ([`TileArg`]), or an atomic tensor
8/// added into and never read ([`AccumulateArg`]).
9#[derive(CubeType, CubeLaunch)]
10pub enum Output<'a, E: Numeric> {
11    Tensor(TileArg<'a>),
12    Atomic(AccumulateArg<'a, E>),
13}
14
15#[cube]
16impl<'a, E: Numeric> Output<'a, E> {
17    /// Serve this operand as a [`Tile`] under the kernel's one `partitioning`.
18    pub fn tile(&self, partitioning: &Partitioning) -> Tile {
19        match self {
20            Output::Tensor(arg) => arg.tile(partitioning),
21            Output::Atomic(arg) => arg.tile(partitioning),
22        }
23    }
24}
25
26impl Bound {
27    /// This operand as an [`Output::Tensor`]: read and replaced.
28    pub fn output<E: Numeric>(self) -> OutputArgs<'static, E> {
29        OutputArgs::Tensor(self.arg())
30    }
31
32    /// This operand as an [`Output::Atomic`]; the buffer must arrive holding zero, since it only
33    /// ever adds.
34    pub fn atomic<E: Numeric>(self) -> OutputArgs<'static, E> {
35        let (spec, width) = (self.spec.clone(), self.vector_size);
36        OutputArgs::Atomic(AccumulateArgLaunch::new(self.tensor(), width, spec))
37    }
38}
cubecl_tile/launch/arg/output.rs

cubecl_tile/launch/arg/
output.rs
