Mojo struct
TileScheduler
@register_passable(trivial)
struct TileScheduler[tile_shape: Index[3], grid_shape: Index[2], cluster: Index[3] = Index(1, 1, 1), raster_dim: SIMD[uint32, 1] = __init__[__mlir_type.!pop.int_literal](1), schedule: MatmulSchedule = MatmulSchedule(__init__[__mlir_type.!pop.int_literal](1))]
Aliases
wave_shape = Index((grid_shape.__getitem__[::Indexer](1) * tile_shape.__getitem__[::Indexer](0)), (grid_shape.__getitem__[::Indexer](0) * tile_shape.__getitem__[::Indexer](1))):num_grids = SIMD((grid_shape.__getitem__[::Indexer](0) * grid_shape.__getitem__[::Indexer](1))):
Fields
- idx (
SIMD[uint32, 1]): - prob_shape (
Index[3]): - num_waves_m (
SIMD[uint32, 1]): - num_waves_n (
SIMD[uint32, 1]): - log_num_waves_n (
FastDiv[uint32]):
Implemented traits
AnyType,
UnknownDestructibility
Methods
__init__
__init__(prob_shape: Index[3]) -> Self
get_current_work_info
get_current_work_info(self) -> WorkInfo
advance
advance(mut self)
fetch_next_work
fetch_next_work(mut self) -> WorkInfo
num_output_tiles
num_output_tiles(self) -> UInt
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!