Struct UnitTile 

pub struct UnitTile<E: Numeric> {
    pub data: Array<E>,
    pub layout: UnitTileLayout,
}

Fields

data: Array<E>

layout: UnitTileLayout

Implementations

impl<E: Numeric> UnitTile<E>

pub fn new(layout: UnitTileLayout) -> UnitTile<E>

pub fn zero(&mut self)

pub fn get(&self, row: u32, col: u32) -> E

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 tile is acting as a materialized mask fragment.

pub fn accumulate(&mut self, row: u32, col: u32, val: E)

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 copy_from<E2: Numeric>(&mut self, other: &UnitTile<E2>)

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

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

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

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

pub fn __expand_get( scope: &Scope, this: &<Self as CubeType>::ExpandType, row: <u32 as CubeType>::ExpandType, col: <u32 as CubeType>::ExpandType, ) -> <E 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 tile is acting as a materialized mask fragment.

pub fn __expand_accumulate( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, row: <u32 as CubeType>::ExpandType, col: <u32 as CubeType>::ExpandType, val: <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_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_copy_from<E2: Numeric>( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, other: &<UnitTile<E2> as CubeType>::ExpandType, )

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

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> UnitTile<E>

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

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

pub fn fill_zero(&mut self)

pub fn write_to<Lhs: Float>(&self, dest: &mut UnitTile<Lhs>)

Cast-copies this unit tile into dest. Used by per-variant softmax helpers when writing the post-softmax score into a same-storage destination.

pub fn exp_diff(&mut self, rowwise: &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, )

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

pub fn __expand_write_to<Lhs: Float>( scope: &Scope, this: &<Self as CubeType>::ExpandType, dest: &mut <UnitTile<Lhs> as CubeType>::ExpandType, )

Cast-copies this unit tile into dest. Used by per-variant softmax helpers when writing the post-softmax score into a same-storage destination.

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

impl<Acc: Float> UnitTile<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 the per-unit-full Unit variant. Each unit holds the whole tile in registers; reductions are register-local. Destination must be another UnitTile.

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 the per-unit-full Unit variant. Each unit holds the whole tile in registers; reductions are register-local. Destination must be another UnitTile.

Trait Implementations

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

type ExpandType = UnitTileExpand<E>