Struct RegisterTileExpand 

pub struct RegisterTileExpand<N: Numeric> {
    pub tile: <UnitTile<N> as CubeType>::ExpandType,
    pub matrix_layout: MatrixLayout,
    pub tile_size: TileSize,
    pub product_type: ProductType,
}

Fields

tile: <UnitTile<N> as CubeType>::ExpandType

matrix_layout: MatrixLayout

tile_size: TileSize

product_type: ProductType

Implementations

impl<E: Float> RegisterTileExpand<E>

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)

pub fn __expand_write_to_method<Lhs: Float>( &self, scope: &Scope, 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.

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)

pub fn __expand_write_to<Lhs: Float>( scope: &Scope, this: &Self, 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> RegisterTileExpand<A>

pub fn __expand_mma_method<L: Numeric, R: Numeric>( &mut self, scope: &Scope, 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.

pub fn __expand_mma<L: Numeric, R: Numeric>( scope: &Scope, this: &mut Self, 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> RegisterTileExpand<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 register tile from source. Supported sources: Shared (product-type-aware load) and None (zero-init).

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

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 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, ident: StageIdent)

impl<Acc: Float> RegisterTileExpand<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 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, 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> AsMutExpand for RegisterTileExpand<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 RegisterTileExpand<N>

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

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

impl<N: Numeric> CubeDebug for RegisterTileExpand<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 RegisterTileExpand<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 RegisterTileExpand<N>

type Expand = RegisterTileExpand<N>

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

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

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

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