Mojo struct
WarpRole1D1D
@register_passable(trivial)
struct WarpRole1D1D
Warp role for 1D-1D kernel with 3-warp specialization.
Thread layout (192 threads total) - matches original kernel:
- Warps 0-3 (threads 0-127): Epilogue (4 warps)
- Warp 4 (threads 128-159): TMA Load
- Warp 5 (threads 160-191): MMA
This layout matches the original grouped_matmul_sm100_1d1d.mojo kernel
which uses WarpRole[has_scheduler=False]. The epilogue warps being at
0-3 is important because TMAStoreCoords uses warp_id == 0 for election.
No scheduler warp - work distribution uses linear grid traversal.
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
EPILOGUE_WARP_START
comptime EPILOGUE_WARP_START = 0
LOAD_WARP_START
comptime LOAD_WARP_START = 128
MMA_WARP_START
comptime MMA_WARP_START = 160
NUM_EPILOGUE_THREADS
comptime NUM_EPILOGUE_THREADS = 128
NUM_LOAD_THREADS
comptime NUM_LOAD_THREADS = 32
NUM_MMA_THREADS
comptime NUM_MMA_THREADS = 32
TOTAL_THREADS
comptime TOTAL_THREADS = 192
Methods
is_epilogue
static is_epilogue() -> Bool
Returns True if current thread is in an epilogue warp (warps 0-3).
Returns:
is_load
static is_load() -> Bool
Returns True if current thread is in the TMA load warp (warp 4).
Returns:
is_mma
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!