Struct BounceTileExpand 

pub struct BounceTileExpand<N: Numeric> {
    pub cmma: <CmmaTile<N> as CubeType>::ExpandType,
    pub smem: <Shared<[N]> as CubeType>::ExpandType,
    pub fragment: <WhiteboxFragment<N> as CubeType>::ExpandType,
}

Fields

cmma: <CmmaTile<N> as CubeType>::ExpandType

smem: <Shared<[N]> as CubeType>::ExpandType

fragment: <WhiteboxFragment<N> as CubeType>::ExpandType

Implementations

impl<E: Float> BounceTileExpand<E>

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

Synchronizes the fragment view from the cmma fragment via smem. Call before any rowwise/elementwise op so the fragment reflects the current cmma state.

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

Synchronizes the cmma fragment from the fragment view via smem. Call after rowwise/elementwise edits to make the cmma side current for the next mma.

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_and_mask_method<M: Mask>( &mut self, scope: &Scope, scale: <E as CubeType>::ExpandType, mask: &<M as CubeType>::ExpandType, )

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

Zeros the cmma fragment. The fragment view is not the live storage at fill_zero call sites (always invoked before any cmma_to_fragment), so only cmma needs clearing.

pub fn __expand_write_fragment_to_method<Lhs: Float, Sc: TileScope>( &self, scope: &Scope, softmaxed: &mut <Tile<Lhs, Sc> as CubeType>::ExpandType, )

Writes the (already-softmaxed) fragment view of this bounce tile into softmaxed. The source fragment is plane-fragmented; for a Bounce destination this routes through the destination’s smem into its cmma fragment.

pub fn __expand_cmma_to_fragment(scope: &Scope, this: &mut Self)

Synchronizes the fragment view from the cmma fragment via smem. Call before any rowwise/elementwise op so the fragment reflects the current cmma state.

pub fn __expand_fragment_to_cmma(scope: &Scope, this: &mut Self)

Synchronizes the cmma fragment from the fragment view via smem. Call after rowwise/elementwise edits to make the cmma side current for the next mma.

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_and_mask<M: Mask>( scope: &Scope, this: &mut Self, scale: <E as CubeType>::ExpandType, mask: &<M as CubeType>::ExpandType, )

pub fn __expand_fill_zero(scope: &Scope, this: &mut Self)

Zeros the cmma fragment. The fragment view is not the live storage at fill_zero call sites (always invoked before any cmma_to_fragment), so only cmma needs clearing.

pub fn __expand_write_fragment_to<Lhs: Float, Sc: TileScope>( scope: &Scope, this: &Self, softmaxed: &mut <Tile<Lhs, Sc> as CubeType>::ExpandType, )

Writes the (already-softmaxed) fragment view of this bounce tile into softmaxed. The source fragment is plane-fragmented; for a Bounce destination this routes through the destination’s smem into its cmma fragment.

impl<N: Numeric> BounceTileExpand<N>

pub fn __expand_copy_from_method<SE: Numeric, SS: Size, Sc: TileScope>( &mut self, scope: &Scope, source: &<Tile<SE, Sc> as CubeType>::ExpandType, ident: StageIdent, )

Copies into the bounce tile’s cmma fragment from source. Bounce always loads through its CMMA representation (the WhiteboxFragment view is synced lazily on demand by softmax/scale ops); supported sources mirror CmmaTile::copy_from.

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

Zero-init the bounce tile (clears its cmma fragment).

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

Copies into the bounce tile’s cmma fragment from source. Bounce always loads through its CMMA representation (the WhiteboxFragment view is synced lazily on demand by softmax/scale ops); supported sources mirror CmmaTile::copy_from.

pub fn __expand_init_zero(scope: &Scope, this: &mut Self)

Zero-init the bounce tile (clears its cmma fragment).

impl<Acc: Float> BounceTileExpand<Acc>

pub fn __expand_softmax_method<Lhs: Float, M: Mask>( &mut self, scope: &Scope, mask: &<M as CubeType>::ExpandType, softmaxed: &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 for the Bounce variant. cmma → fragment once at entry so all subsequent rowwise ops read/write the fragment view; the post- exp values are still in the fragment at the end (we skip fragment_to_cmma on score because its cmma is cleared next iteration), and we stream straight into softmaxed via write_fragment_to.

pub fn __expand_softmax<Lhs: Float, M: Mask>( scope: &Scope, this: &mut Self, mask: &<M as CubeType>::ExpandType, softmaxed: &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 for the Bounce variant. cmma → fragment once at entry so all subsequent rowwise ops read/write the fragment view; the post- exp values are still in the fragment at the end (we skip fragment_to_cmma on score because its cmma is cleared next iteration), and we stream straight into softmaxed via write_fragment_to.

Trait Implementations

impl<N: Numeric> AsMutExpand for BounceTileExpand<N>

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

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

impl<N: Numeric> AsRefExpand for BounceTileExpand<N>

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

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

impl<N: Numeric> CubeDebug for BounceTileExpand<N>

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> ExpandTypeClone for BounceTileExpand<N>

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> IntoExpand for BounceTileExpand<N>

type Expand = BounceTileExpand<N>

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

impl<N: Numeric> IntoMut for BounceTileExpand<N>

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

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