Mojo struct
TileScheduler
@register_passable(trivial)
struct TileScheduler[problem_shape: IndexList[3], tile_shape: IndexList[3], grid_shape: IndexList[2], cluster: IndexList[3] = Index(1, 1, 1), raster_dim: UInt32 = 1, schedule: MatmulSchedule = MatmulSchedule(1)]
Fields
- idx (
UInt32): - prob_shape (
IndexList[3]): - num_waves_m (
UInt32): - num_waves_n (
UInt32): - log_num_waves_n (
FastDiv[DType.uint32]): - current_iter (
Int): - num_aligned_m_blocks (
UInt32): - num_blocks (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
alias __copyinit__is_trivial = True
__del__is_trivial
alias __del__is_trivial = True
__moveinit__is_trivial
alias __moveinit__is_trivial = True
kNum1DBlocksPerGroup
alias kNum1DBlocksPerGroup = 16
kNumNBlocks
alias kNumNBlocks = SIMD[DType.uint32, 1](ceildiv(problem_shape.__getitem__[3, DType.int64, Int](1), tile_shape.__getitem__[3, DType.int64, Int](1)))
num_grids
alias num_grids = SIMD[DType.uint32, 1]((grid_shape.__getitem__[2, DType.int64, Int](0) * grid_shape.__getitem__[2, DType.int64, Int](1)))
wave_shape
alias wave_shape = Index[dtype=DType.uint32]((grid_shape.__getitem__[2, DType.int64, Int](1) * tile_shape.__getitem__[3, DType.int64, Int](0)), (grid_shape.__getitem__[2, DType.int64, Int](0) * tile_shape.__getitem__[3, DType.int64, Int](1)))
Methods
__init__
__init__(prob_shape: IndexList[3]) -> Self
get_current_work_info
get_current_work_info(mut self) -> WorkInfo
Returns:
WorkInfo
advance
advance(mut self)
fetch_next_work
fetch_next_work(mut self) -> WorkInfo
Returns:
WorkInfo
num_output_tiles
fetch_next_work_ds
fetch_next_work_ds(mut self) -> WorkInfo
Returns:
WorkInfo
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!