pub struct PartitionScheduler {
pub m: AxisScheduler,
pub n: AxisScheduler,
pub k: AxisScheduler,
}Expand description
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: AxisSchedulerImplementations§
Source§impl PartitionScheduler
impl PartitionScheduler
Sourcepub fn new(
partition_index_m: u32,
partition_index_n: u32,
partition_size: PartitionSize,
partition_schedule_scheme: PartitionSchedulerScheme,
) -> 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.
Sourcepub fn __expand_new(
scope: &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_new( scope: &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
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.
Sourcepub fn __expand_map_m(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
i: <u32 as CubeType>::ExpandType,
) -> <u32 as CubeType>::ExpandType
pub fn __expand_map_m( scope: &Scope, this: &<Self as CubeType>::ExpandType, i: <u32 as CubeType>::ExpandType, ) -> <u32 as CubeType>::ExpandType
Maps a local M index to a global index.
Sourcepub fn __expand_map_n(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
i: <u32 as CubeType>::ExpandType,
) -> <u32 as CubeType>::ExpandType
pub fn __expand_map_n( scope: &Scope, this: &<Self as CubeType>::ExpandType, i: <u32 as CubeType>::ExpandType, ) -> <u32 as CubeType>::ExpandType
Maps a local N index to a global index.
Sourcepub fn __expand_map_k(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
i: <u32 as CubeType>::ExpandType,
) -> <u32 as CubeType>::ExpandType
pub fn __expand_map_k( scope: &Scope, this: &<Self as CubeType>::ExpandType, i: <u32 as CubeType>::ExpandType, ) -> <u32 as CubeType>::ExpandType
Maps a local K index to a global index.
Trait Implementations§
Auto Trait Implementations§
impl Freeze for PartitionScheduler
impl RefUnwindSafe for PartitionScheduler
impl Send for PartitionScheduler
impl Sync for PartitionScheduler
impl Unpin for PartitionScheduler
impl UnsafeUnpin for PartitionScheduler
impl UnwindSafe for PartitionScheduler
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more
impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more