Mojo struct

TileScheduler

@register_passable(trivial) struct TileScheduler[offsets_layout: Layout, //, *, static_MN: Int, tile_shape: IndexList[3], cluster: IndexList[3] = Index(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, MutAnyOrigin]):
current_iter (Int32):
current_group_idx (UInt32):
current_dynamic_dim_cumsum (UInt32):
block_idx_start (UInt32):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, Movable, UnknownDestructibility

`comptime` members

`copyinitis_trivial`

comptime __copyinit__is_trivial = True

`delis_trivial`

comptime __del__is_trivial = True

`moveinitis_trivial`

comptime __moveinit__is_trivial = True

`cta_group_tile_shape`

comptime cta_group_tile_shape = Index((tile_shape.__getitem__[3, DType.int64, Int](0) * cta_group), (tile_shape.__getitem__[3, DType.int64, Int](1) * cta_group))

`div_dynamic_block`

comptime div_dynamic_block = FastDiv[DType.uint32](TileScheduler[static_MN=static_MN, tile_shape=tile_shape, cluster=cluster, cta_group=cta_group, swizzle=swizzle, swapAB=swapAB].cta_group_tile_shape.__getitem__[2, DType.int64, Int](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 = 16

`num_static_dim_blocks`

comptime num_static_dim_blocks = SIMD[DType.uint32, 1](ceildiv(static_MN, tile_shape.__getitem__[3, DType.int64, Int](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, MutAnyOrigin]) -> Self

`fetch_next_work`

fetch_next_work(mut self) -> WorkInfo

Returns:

WorkInfo

Fields​

Implemented traits​

comptime members​

__copyinit__is_trivial​

__del__is_trivial​

__moveinit__is_trivial​

cta_group_tile_shape​

div_dynamic_block​

dynamic_dim​

kNum1DBlocksPerGroup​

num_static_dim_blocks​

static_dim​

Methods​

__init__​

fetch_next_work​