Mojo struct

TileScheduler

struct TileScheduler[offsets_layout: Layout, //, *, static_MN: Int, tile_shape: IndexList[3], cluster: IndexList[3] = Index[Int, Int, Int](1, 1, 1), cta_group: Int = 1, swizzle: Bool = False, swapAB: Bool = True]

Fields

num_active_experts (Int):
group_offsets (LayoutTensor[DType.uint32, offsets_layout, ImmutAnyOrigin]):
current_iter (Int32):
current_group_idx (UInt32):
current_dynamic_dim_cumsum (UInt32):
block_idx_start (UInt32):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable

`comptime` members

`cta_group_tile_shape`

comptime cta_group_tile_shape = Index[Int, Int]((tile_shape[0] * cta_group), (tile_shape[1] * cta_group))

`div_dynamic_block`

comptime div_dynamic_block = FastDiv(TileScheduler[static_MN=static_MN, tile_shape=tile_shape, cluster=cluster, cta_group=cta_group, swizzle=swizzle, swapAB=swapAB].cta_group_tile_shape[TileScheduler[static_MN=static_MN, tile_shape=tile_shape, cluster=cluster, cta_group=cta_group, swizzle=swizzle, swapAB=swapAB].dynamic_dim])

`dynamic_dim`

comptime dynamic_dim = 1 if swapAB else 0

`kNum1DBlocksPerGroup`

comptime kNum1DBlocksPerGroup = SIMD(16)

`num_static_dim_blocks`

comptime num_static_dim_blocks = SIMD(ceildiv(static_MN, tile_shape[TileScheduler[static_MN=static_MN, tile_shape=tile_shape, cluster=cluster, cta_group=cta_group, swizzle=swizzle, swapAB=swapAB].static_dim]))

`static_dim`

comptime static_dim = 0 if swapAB else 1

Methods

`init`

__init__(num_active_experts: Int, group_offsets: LayoutTensor[DType.uint32, offsets_layout, ImmutAnyOrigin]) -> Self

`fetch_next_work`

fetch_next_work(mut self) -> WorkInfo

Returns:

WorkInfo

Fields​

Implemented traits​

comptime members​

cta_group_tile_shape​

div_dynamic_block​

dynamic_dim​

kNum1DBlocksPerGroup​

num_static_dim_blocks​

static_dim​

Methods​

__init__​

fetch_next_work​