Struct TileExpand 

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

Implementations

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

pub fn __expand_partition_tile_at_mut_method<'infer, 'scope>( &'infer mut self, scope: &'scope Scope, 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.

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, 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> TileExpand<N, Sc>

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

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

pub fn __expand_copy_from_method<SE: Numeric, SS: Size, L: Numeric, R: Numeric, A: Numeric>( &mut self, scope: &Scope, 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.

pub fn __expand_init_zero<L: Numeric, R: Numeric>( scope: &Scope, this: &mut Self, 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, 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> TileExpand<E, Plane>

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

Multiply by scale and add -inf at masked positions.

pub fn __expand_fill_zero_method(&mut self, scope: &Scope)

pub fn __expand_scale_and_mask<M: Mask>( scope: &Scope, this: &mut Self, 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)

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

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

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, tile: &<StridedTile<E, ES> as CubeType>::ExpandType, )

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

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

pub fn __expand_mma_method<L: Numeric, R: Numeric>( &mut self, scope: &Scope, 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_method<LhsS: Numeric, LhsSize: Size, LhsR: Numeric, RhsS: Numeric, RhsSize: Size, RhsR: Numeric, SEL: StageEventListener>( &mut self, scope: &Scope, 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.

pub fn __expand_mma<L: Numeric, R: Numeric>( scope: &Scope, this: &mut Self, 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, 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> TileExpand<N, Sc>

pub fn __expand_partition_method<P: Partitioner>( &self, scope: &Scope, 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.

pub fn __expand_partition<P: Partitioner>( scope: &Scope, this: &Self, 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> TileExpand<E, Plane>

pub fn __expand_row_max_method( &self, scope: &Scope, acc: &mut <RowWise<E> as CubeType>::ExpandType, base: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_row_sum_method( &self, scope: &Scope, acc: &mut <RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_exp_diff_method( &mut self, scope: &Scope, rowwise: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_rowwise_scale_method( &mut self, scope: &Scope, scale: &<RowWise<E> as CubeType>::ExpandType, )

pub fn __expand_scale_mul_method<SM: Float>( &mut self, scope: &Scope, 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_method<SM: Float>( &mut self, scope: &Scope, running_state_l: &<RowWise<SM> as CubeType>::ExpandType, )

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

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

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

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

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

pub fn __expand_scale_mul<SM: Float>( scope: &Scope, this: &mut Self, 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, running_state_l: &<RowWise<SM> as CubeType>::ExpandType, )

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

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

pub fn __expand_softmax_method<Lhs: Float, M: Mask>( &mut self, scope: &Scope, 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_method<DE: Float, DS: Size>( &self, scope: &Scope, dest: &mut <Tile<DE, Plane> as CubeType>::ExpandType, )

Copy self into dest.

pub fn __expand_softmax<Lhs: Float, M: Mask>( scope: &Scope, this: &mut Self, 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, dest: &mut <Tile<DE, Plane> as CubeType>::ExpandType, )

Copy self into dest.

Trait Implementations

impl<N: Numeric, Sc: TileScope> AsMutExpand for TileExpand<N, Sc>

fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self

fn __expand_as_mut_method(&mut self, scope: &Scope) -> &mut T

impl<N: Numeric, Sc: TileScope> AsRefExpand for TileExpand<N, Sc>

fn __expand_ref_method(&self, _: &Scope) -> &Self

fn __expand_as_ref_method(&self, scope: &Scope) -> &T

impl<N: Numeric, Sc: TileScope> CubeDebug for TileExpand<N, Sc>

fn set_debug_name(&self, scope: &Scope, name: &'static str)

Set the debug name of this type’s expansion. Should do nothing for types that don’t appear at runtime

impl<N: Numeric, Sc: TileScope> ExpandTypeClone for TileExpand<N, Sc>

fn clone_unchecked(&self) -> Self

Unchecked clone that only clones the conceptual runtime value. Should only be used in cases where each copy is used in a mutually exclusive branch (i.e. match, runtime enums). This is intentionally separated from Rust’s Clone semantics and should only be used for the conceptual expand values, never real data. Using two values in the same branch is undefined behaviour.

impl<N: Numeric, Sc: TileScope> IntoExpand for TileExpand<N, Sc>

type Expand = TileExpand<N, Sc>

fn into_expand(self, _: &Scope) -> Self

impl<N: Numeric, Sc: TileScope> IntoMut for TileExpand<N, Sc>

fn into_mut(self, scope: &Scope) -> Self

Convert the variable into a potentially new mutable variable in scope, copying if needed.

impl<E: Numeric, Sc: TileScope> MaskExpand for TileExpand<E, Sc>

fn __expand_should_mask_method( &self, scope: &Scope, local_pos: <Coords2d as CubeType>::ExpandType, ) -> <bool as CubeType>::ExpandType