Struct RegisterTile 

pub struct RegisterTile<N: Numeric> {
    pub tile: UnitTile<N>,
    pub matrix_layout: MatrixLayout,
    pub tile_size: TileSize,
    pub product_type: ProductType,
}

Register-resident matmul tile. Built on top of UnitTile (per-unit full register copy of the data). Holds only the minimal comptime data the tile body uses (tile_size + product_type); rowwise / elementwise ops are delegated to the inner UnitTile. The matmul-level configuration lives in cubek-matmul as RegisterMatmul.

Fields

tile: UnitTile<N>

matrix_layout: MatrixLayout

tile_size: TileSize

product_type: ProductType

Implementations

impl<E: Float> RegisterTile<E>

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_and_mask<M: Mask>(&mut self, scale: E, mask: &M)

pub fn fill_zero(&mut self)

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

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

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

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

impl<A: Numeric> RegisterTile<A>

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

Executes lhs · rhs, accumulating into self via the configured inner/outer software product.

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

Executes lhs · rhs, accumulating into self via the configured inner/outer software product.

impl<N: Numeric> RegisterTile<N>

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

Copies into the register tile from source. Supported sources: Shared (product-type-aware load) and None (zero-init).

pub fn init_zero(&mut self, ident: StageIdent)

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

Copies into the register tile from source. Supported sources: Shared (product-type-aware load) and None (zero-init).

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

impl<Acc: Float> RegisterTile<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 Register variant (legacy direct-register attention path). Destination must be another RegisterTile.

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 Register variant (legacy direct-register attention path). Destination must be another RegisterTile.

Trait Implementations

impl<N: Numeric> CubeType for RegisterTile<N>

type ExpandType = RegisterTileExpand<N>