cubek_std/tile/ops/
copy.rs

1use cubecl::prelude::*;
2
3use crate::{
4    StageIdent,
5    tile::{Tile, TileExpand, TileKind, TileKindExpand, TileScope},
6};
7
8#[cube]
9impl<N: Numeric, Sc: TileScope> Tile<N, Sc> {
10    /// Zero-initialize the tile in place. `L`/`R` are only consulted on MMA.
11    pub fn init_zero<L: Numeric, R: Numeric>(&mut self, #[comptime] ident: StageIdent) {
12        match &mut self.kind {
13            TileKind::Cmma(t) => t.init_zero(),
14            TileKind::Bounce(b) => b.init_zero(),
15            TileKind::Mma(t) => t.init_zero::<L, R>(),
16            TileKind::Register(t) => t.init_zero(ident),
17            TileKind::PlaneVec(t) => t.init_zero(),
18            TileKind::Interleaved(t) => t.init_zero(),
19            TileKind::RowWise(t) => t.init_zero(),
20            _ => panic!("init_zero: unsupported tile variant"),
21        }
22    }
23
24    /// Copy `source` into `self`. `SS` is the smem vector size involved in
25    /// the copy; `L`/`R`/`A` are only consulted on MMA paths.
26    pub fn copy_from<SE: Numeric, SS: Size, L: Numeric, R: Numeric, A: Numeric>(
27        &mut self,
28        source: &Tile<SE, Sc>,
29        #[comptime] ident: StageIdent,
30    ) {
31        match &mut self.kind {
32            TileKind::Cmma(t) => t.copy_from::<SE, SS, Sc>(source, ident),
33            TileKind::Bounce(b) => b.copy_from::<SE, SS, Sc>(source, ident),
34            TileKind::Mma(t) => t.copy_from::<SE, SS, L, R, A, Sc>(source, ident),
35            TileKind::Register(t) => t.copy_from::<SE, SS, Sc>(source, ident),
36            TileKind::PlaneVec(t) => t.copy_from::<SE, SS, Sc>(source, ident),
37            TileKind::Interleaved(t) => t.copy_from::<SE, SS, Sc>(source, ident),
38            TileKind::SharedTile(shared) => {
39                shared.copy_from::<SE, SS, L, R, Sc>(source);
40            }
41            _ => panic!("copy_from: unsupported destination variant"),
42        }
43    }
44}
cubek_std/tile/ops/copy.rs

cubek_std/tile/ops/
copy.rs
