Struct WhiteboxFragment 

pub struct WhiteboxFragment<E: Numeric> {
    pub array: Array<E>,
    pub layout: WhiteboxFragmentLayout,
}

Assumes:

• unit_size * plane_dim = total_size (not dim wise but in total count)

Fields

array: Array<E>

layout: WhiteboxFragmentLayout

Implementations

impl<E: Numeric> WhiteboxFragment<E>

pub fn new(layout: WhiteboxFragmentLayout) -> WhiteboxFragment<E>

pub fn zero(&mut self)

pub fn load_from_slice(&mut self, smem_slice: &[E])

pub fn load_from_strided_tile<E2: Numeric, N: Size>( &mut self, strided_tile: &StridedTile<E2, N>, )

pub fn store_to<F: Float>(&self, smem_slice: &mut [F])

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

Reads the element at local_pos and casts to bool. Used by the Mask trait dispatcher when this fragment is acting as a materialized mask fragment.

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

pub fn rowwise_max(&self) -> RowWise<E>

pub fn rowwise_sum(&self) -> RowWise<E>

pub fn num_units_per_row(&self) -> u32

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

pub fn __expand_new( scope: &Scope, layout: WhiteboxFragmentLayout, ) -> <WhiteboxFragment<E> as CubeType>::ExpandType

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

pub fn __expand_load_from_slice( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, smem_slice: &<[E] as CubeType>::ExpandType, )

pub fn __expand_load_from_strided_tile<E2: Numeric, N: Size>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, strided_tile: &<StridedTile<E2, N> as CubeType>::ExpandType, )

pub fn __expand_store_to<F: Float>( scope: &Scope, this: &<Self as CubeType>::ExpandType, smem_slice: &mut <[F] as CubeType>::ExpandType, )

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

Reads the element at local_pos and casts to bool. Used by the Mask trait dispatcher when this fragment is acting as a materialized mask fragment.

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

pub fn __expand_rowwise_max( scope: &Scope, this: &<Self as CubeType>::ExpandType, ) -> <RowWise<E> as CubeType>::ExpandType

pub fn __expand_rowwise_sum( scope: &Scope, this: &<Self as CubeType>::ExpandType, ) -> <RowWise<E> as CubeType>::ExpandType

pub fn __expand_num_units_per_row( scope: &Scope, this: &<Self as CubeType>::ExpandType, ) -> u32

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, )

impl<E: Float> WhiteboxFragment<E>

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

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

impl<E: Float> WhiteboxFragment<E>

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

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

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, )

impl<Acc: Float> WhiteboxFragment<Acc>

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

Online softmax for a free-standing WhiteboxFragment score (the register-only variant of attention’s softmax). Writes the post-softmax values into softmaxed (which may be Bounce — routed through its smem into its cmma fragment — or another WhiteboxFragment).

pub fn __expand_softmax<Lhs: Float, M: Mask>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, 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 a free-standing WhiteboxFragment score (the register-only variant of attention’s softmax). Writes the post-softmax values into softmaxed (which may be Bounce — routed through its smem into its cmma fragment — or another WhiteboxFragment).

Trait Implementations

impl<E: Numeric> CubeType for WhiteboxFragment<E>

type ExpandType = WhiteboxFragmentExpand<E>