Struct PartitionScheduler 

pub struct PartitionScheduler {
    pub m: AxisScheduler,
    pub n: AxisScheduler,
    pub k: AxisScheduler,
}

Schedules global indices for M, N, and K axes in a partitioned matmul. Internally uses an AxisScheduler per axis.

Fields

m: AxisScheduler

n: AxisScheduler

k: AxisScheduler

Implementations

impl PartitionScheduler

pub fn new( partition_index_m: u32, partition_index_n: u32, partition_size: PartitionSize, partition_schedule_scheme: PartitionSchedulerScheme, ) -> PartitionScheduler

Creates a PartitionScheduler for a partition at (partition_index_m, partition_index_n).

• Offset: rotates indices per axis to spread memory access and avoid conflicts.
• Naive: uses simple row-major order for partitions.

pub fn map_m(&self, i: u32) -> u32

Maps a local M index to a global index.

pub fn map_n(&self, i: u32) -> u32

Maps a local N index to a global index.

pub fn map_k(&self, i: u32) -> u32

Maps a local K index to a global index.

pub fn __expand_new( scope: &mut Scope, partition_index_m: <u32 as CubeType>::ExpandType, partition_index_n: <u32 as CubeType>::ExpandType, partition_size: PartitionSize, partition_schedule_scheme: PartitionSchedulerScheme, ) -> <PartitionScheduler as CubeType>::ExpandType

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

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

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

Trait Implementations

impl CubeType for PartitionScheduler

type ExpandType = PartitionSchedulerExpand

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

Wrapper around the init method, necessary to type inference.