Struct Tile 

pub struct Tile<N: Numeric, Sc: TileScope> { /* private fields */ }

Public tile type. Wraps a [TileKind] payload; the inner enum is crate-private and callers construct via Tile::new_*.

Implementations

impl<N: Numeric, Sc: TileScope> Tile<N, Sc>

pub fn new_SharedTile(t: SharedTile<N>) -> Tile<N, Sc>

pub fn new_Stage(t: StageTile<N>) -> Tile<N, Sc>

pub fn new_Partition(t: PartitionTile<N, Sc>) -> Tile<N, Sc>

pub fn new_Pipelined(t: PipelinedTile<N, Sc>) -> Tile<N, Sc>

pub fn new_None() -> Tile<N, Sc>

pub fn new_RowWise(t: RowWise<N>) -> Tile<N, Sc>

pub fn partition_tile_at_mut( &mut self, m: usize, n: usize, n_cols: usize, ) -> &mut Tile<N, Sc>

Mutable reference to the (m, n) element of a Partition tile.

pub fn __expand_new_SharedTile( scope: &Scope, t: <SharedTile<N> as CubeType>::ExpandType, ) -> <Tile<N, Sc> as CubeType>::ExpandType

pub fn __expand_new_Stage( scope: &Scope, t: <StageTile<N> as CubeType>::ExpandType, ) -> <Tile<N, Sc> as CubeType>::ExpandType

pub fn __expand_new_Partition( scope: &Scope, t: <PartitionTile<N, Sc> as CubeType>::ExpandType, ) -> <Tile<N, Sc> as CubeType>::ExpandType

pub fn __expand_new_Pipelined( scope: &Scope, t: <PipelinedTile<N, Sc> as CubeType>::ExpandType, ) -> <Tile<N, Sc> as CubeType>::ExpandType

pub fn __expand_new_None(scope: &Scope) -> <Tile<N, Sc> as CubeType>::ExpandType

pub fn __expand_new_RowWise( scope: &Scope, t: <RowWise<N> as CubeType>::ExpandType, ) -> <Tile<N, Sc> as CubeType>::ExpandType

pub fn __expand_partition_tile_at_mut<'infer, 'scope>( scope: &'scope Scope, this: &'infer mut <Self as CubeType>::ExpandType, m: usize, n: usize, n_cols: usize, ) -> &'infer mut <Tile<N, Sc> as CubeType>::ExpandType

Mutable reference to the (m, n) element of a Partition tile.

impl<N: Numeric, Sc: TileScope> Tile<N, Sc>

pub fn init_zero<L: Numeric, R: Numeric>(&mut self, ident: StageIdent)

Zero-initialize the tile in place. L/R are only consulted on MMA.

pub fn copy_from<SE: Numeric, SS: Size, L: Numeric, R: Numeric, A: Numeric>( &mut self, source: &Tile<SE, Sc>, ident: StageIdent, )

Copy source into self. SS is the smem vector size involved in the copy; L/R/A are only consulted on MMA paths.

pub fn __expand_init_zero<L: Numeric, R: Numeric>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, ident: StageIdent, )

Zero-initialize the tile in place. L/R are only consulted on MMA.

pub fn __expand_copy_from<SE: Numeric, SS: Size, L: Numeric, R: Numeric, A: Numeric>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, source: &<Tile<SE, Sc> as CubeType>::ExpandType, ident: StageIdent, )

Copy source into self. SS is the smem vector size involved in the copy; L/R/A are only consulted on MMA paths.

impl<E: Float> Tile<E, Plane>

pub fn scale_and_mask<M: Mask>(&mut self, scale: E, mask: &M)

Multiply by scale and add -inf at masked positions.

pub fn fill_zero(&mut self)

pub fn __expand_scale_and_mask<M: Mask>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, scale: <E as CubeType>::ExpandType, mask: &<M as CubeType>::ExpandType, )

Multiply by scale and add -inf at masked positions.

pub fn __expand_fill_zero( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, )

impl<N: Numeric, Sc: TileScope> Tile<N, Sc>

pub fn load_mask_from_strided_tile<E: Numeric, ES: Size>( &mut self, tile: &StridedTile<E, ES>, )

Materialize a mask fragment from a StridedTile into Unit or WhiteboxFragment.

pub fn __expand_load_mask_from_strided_tile<E: Numeric, ES: Size>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, tile: &<StridedTile<E, ES> as CubeType>::ExpandType, )

Materialize a mask fragment from a StridedTile into Unit or WhiteboxFragment.

impl<N: Numeric, Sc: TileScope> Tile<N, Sc>

pub fn mma<L: Numeric, R: Numeric>( &mut self, lhs: &Tile<L, Sc>, rhs: &Tile<R, Sc>, )

self += lhs · rhs. For (Stage, Stage, Partition) use Tile::mma_partition.

pub fn mma_partition<LhsS: Numeric, LhsSize: Size, LhsR: Numeric, RhsS: Numeric, RhsSize: Size, RhsR: Numeric, SEL: StageEventListener>( &mut self, lhs: &Tile<LhsS, Sc>, rhs: &Tile<RhsS, Sc>, a_fragment: &mut Sequence<Tile<LhsR, Sc>>, b_fragments: &mut Tile<RhsR, Sc>, partition_size_k: u32, listener: SEL, scheduler: &PartitionScheduler, )

mma for (Stage, Stage, Partition) operands with rhs fragments held under TileKind::Pipelined.

pub fn __expand_mma<L: Numeric, R: Numeric>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, lhs: &<Tile<L, Sc> as CubeType>::ExpandType, rhs: &<Tile<R, Sc> as CubeType>::ExpandType, )

self += lhs · rhs. For (Stage, Stage, Partition) use Tile::mma_partition.

pub fn __expand_mma_partition<LhsS: Numeric, LhsSize: Size, LhsR: Numeric, RhsS: Numeric, RhsSize: Size, RhsR: Numeric, SEL: StageEventListener>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, lhs: &<Tile<LhsS, Sc> as CubeType>::ExpandType, rhs: &<Tile<RhsS, Sc> as CubeType>::ExpandType, a_fragment: &mut <Sequence<Tile<LhsR, Sc>> as CubeType>::ExpandType, b_fragments: &mut <Tile<RhsR, Sc> as CubeType>::ExpandType, partition_size_k: u32, listener: <SEL as CubeType>::ExpandType, scheduler: &<PartitionScheduler as CubeType>::ExpandType, )

mma for (Stage, Stage, Partition) operands with rhs fragments held under TileKind::Pipelined.

impl<N: Numeric, Sc: TileScope> Tile<N, Sc>

pub fn partition<P: Partitioner>( &self, compute_index: u32, plane_dim: u32, num_partitions_col: u32, ) -> Tile<N, P::OutputScope>

View of self at the partitioner’s output scope. Source must be a TileKind::Stage.

pub fn __expand_partition<P: Partitioner>( scope: &Scope, this: &<Self as CubeType>::ExpandType, compute_index: <u32 as CubeType>::ExpandType, plane_dim: u32, num_partitions_col: u32, ) -> <Tile<N, P::OutputScope> as CubeType>::ExpandType

View of self at the partitioner’s output scope. Source must be a TileKind::Stage.

impl<E: Float> Tile<E, Plane>

pub fn row_max(&self, acc: &mut RowWise<E>, base: &RowWise<E>)

pub fn row_sum(&self, acc: &mut RowWise<E>)

pub fn exp_diff(&mut self, rowwise: &RowWise<E>)

pub fn rowwise_scale(&mut self, scale: &RowWise<E>)

pub fn scale_mul<SM: Float>(&mut self, scale: &RowWise<SM>)

Multiply each row of self by scale[r]. The Bounce arm round-trips through smem to keep the cmma fragment current.

pub fn scale_div<SM: Float>(&mut self, running_state_l: &RowWise<SM>)

Divide each row by running_state_l[r]; fully-masked rows stay zero.

pub fn __expand_row_max( scope: &Scope, this: &<Self as CubeType>::ExpandType, acc: &mut <RowWise<E> as CubeType>::ExpandType, base: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_row_sum( scope: &Scope, this: &<Self as CubeType>::ExpandType, acc: &mut <RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_exp_diff( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, rowwise: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_rowwise_scale( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, scale: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_scale_mul<SM: Float>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, scale: &<RowWise<SM> as CubeType>::ExpandType, )

Multiply each row of self by scale[r]. The Bounce arm round-trips through smem to keep the cmma fragment current.

pub fn __expand_scale_div<SM: Float>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, running_state_l: &<RowWise<SM> as CubeType>::ExpandType, )

Divide each row by running_state_l[r]; fully-masked rows stay zero.

impl<Acc: Float> Tile<Acc, Plane>

pub fn softmax<Lhs: Float, M: Mask>( &mut self, mask: &M, softmaxed_tile: &mut Tile<Lhs, Plane>, state: &mut (RowWise<Acc>, RowWise<Acc>), head_dim_factor: Acc, ) -> RowWise<Acc>

Online softmax update fused with the precision-cast write into the value-matmul lhs tile.

pub fn write_results<DE: Float, DS: Size>(&self, dest: &mut Tile<DE, Plane>)

Copy self into dest.

pub fn __expand_softmax<Lhs: Float, M: Mask>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, mask: &<M as CubeType>::ExpandType, softmaxed_tile: &mut <Tile<Lhs, Plane> as CubeType>::ExpandType, state: &mut (<RowWise<Acc> as CubeType>::ExpandType, <RowWise<Acc> as CubeType>::ExpandType), head_dim_factor: <Acc as CubeType>::ExpandType, ) -> <RowWise<Acc> as CubeType>::ExpandType

Online softmax update fused with the precision-cast write into the value-matmul lhs tile.

pub fn __expand_write_results<DE: Float, DS: Size>( scope: &Scope, this: &<Self as CubeType>::ExpandType, dest: &mut <Tile<DE, Plane> as CubeType>::ExpandType, )

Copy self into dest.

Trait Implementations

impl<N: Numeric, Sc: TileScope> CubeType for Tile<N, Sc>

type ExpandType = TileExpand<N, Sc>

impl<E: Numeric, Sc: TileScope> Mask for Tile<E, Sc>

fn should_mask(&self, local_pos: Coords2d) -> bool

fn __expand_should_mask( scope: &Scope, this: &<Self as CubeType>::ExpandType, local_pos: <Coords2d as CubeType>::ExpandType, ) -> <bool as CubeType>::ExpandType