Mojo struct
TileScheduler
struct TileScheduler[problem_shape: IndexList[3], tile_shape: IndexList[3], grid_shape: IndexList[2], cluster: IndexList[3] = Index[Int, Int, Int](1, 1, 1), raster_dim: UInt32 = UInt32(1), schedule: MatmulSchedule = MatmulSchedule.TILE2D]
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,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
kNum1DBlocksPerGroup
comptime kNum1DBlocksPerGroup = UInt32(16)
kNumNBlocks
comptime kNumNBlocks = SIMD(ceildiv(problem_shape[1], tile_shape[1]))
num_grids
comptime num_grids = SIMD((grid_shape[0] * grid_shape[1]))
wave_shape
comptime wave_shape = Index[Int, Int, dtype=DType.uint32]((tile_shape[0] * grid_shape[1]), (tile_shape[1] * grid_shape[0]))
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!