Struct TmaWeightLoader

pub struct TmaWeightLoader<MP: MatmulPrecision, S: StageConfig> {
    pub tensor_view: MappedTensorReader<MP::EI>,
    pub stage: Stage<MP::ES, TmaWeightTiling>,
    /* private fields */
}

Fields

tensor_view: MappedTensorReader<MP::EI>

stage: Stage<MP::ES, TmaWeightTiling>

Implementations

impl<MP: MatmulPrecision, S: StageConfig> TmaWeightLoader<MP, S>

pub fn new<G: GlobalConfig>( tensor: TensorMap<MP::EI>, x: u32, y: u32, quantization: CubeOption<Quantization<MP>>, runtime_args: &RuntimeArgs, config: G, ) -> Self

pub fn fill_stage(this: &mut Self, barrier: &Barrier<MP::ES>, config: S)

pub fn reader(this: &Self) -> TmaWeightReader<MP>

pub fn advance_view(this: &mut Self, k_offset: u32)

pub fn __expand_new<G: GlobalConfig>( context: &mut Scope, tensor: <TensorMap<MP::EI> as CubeType>::ExpandType, x: <u32 as CubeType>::ExpandType, y: <u32 as CubeType>::ExpandType, quantization: <CubeOption<Quantization<MP>> as CubeType>::ExpandType, runtime_args: <RuntimeArgs as CubeType>::ExpandType, config: G, ) -> <Self as CubeType>::ExpandType

pub fn __expand_fill_stage( context: &mut Scope, this: <Self as CubeType>::ExpandType, barrier: <Barrier<MP::ES> as CubeType>::ExpandType, config: S, ) -> <() as CubeType>::ExpandType

pub fn __expand_reader( context: &mut Scope, this: <Self as CubeType>::ExpandType, ) -> <TmaWeightReader<MP> as CubeType>::ExpandType

pub fn __expand_advance_view( context: &mut Scope, this: <Self as CubeType>::ExpandType, k_offset: <u32 as CubeType>::ExpandType, ) -> <() as CubeType>::ExpandType

Trait Implementations

impl<MP: MatmulPrecision, S: StageConfig> CubeType for TmaWeightLoader<MP, S>

type ExpandType = TmaWeightLoaderExpand<MP, S>

fn init(scope: &mut Scope, expand: Self::ExpandType) -> Self::ExpandType

Wrapper around the init method, necessary to type inference.