Mojo struct
TMAStoreCoords
@register_passable(trivial)
struct TMAStoreCoords[BM: Int, BN: Int, MMA_M: Int, MMA_N: Int, stageN: Int, cta_group: Int, c_smem_shape0: Int, stage: Int, batched: Bool = False]
TMA store coordinates and warp election for SM100 epilogue.
When batched=True, includes a batch coordinate for 3D TMA stores.
Fields
- coord_m (
UInt): - coord_n (
UInt): - coord_b (
UInt): - elect_one_warp (
Bool): - c_smem_coord_m (
UInt):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
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
CG1_TMA_BM
comptime CG1_TMA_BM = c_smem_shape0
CG2_TMA_BM
comptime CG2_TMA_BM = c_smem_shape0 if (eq MMA_M._mlir_value, 256) else BM
stage_n_offset
comptime stage_n_offset = (stage * stageN)
TMA_BM
comptime TMA_BM = c_smem_shape0 if (eq MMA_M._mlir_value, 256) else BM if (eq cta_group._mlir_value, 2) else TMAStoreCoords[BM, BN, MMA_M, MMA_N, stageN, cta_group, c_smem_shape0, stage, batched].CG1_TMA_BM
Methods
__init__
__init__(c_coord: Tuple[UInt32, UInt32], warp_id: UInt32) -> Self
Compute TMA store coordinates from 2D tile coords and warp ID.
__init__(c_coord: Tuple[UInt32, UInt32, UInt32], warp_id: UInt32) -> Self
Compute TMA store coordinates from 3D tile coords and warp ID.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!