For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

TMAStoreCoords

struct TMAStoreCoords[epc: EpilogueConfig, 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 (Int):
coord_n (Int):
coord_b (Int):
elect_one_warp (Bool):
c_smem_coord_m (Int):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

`comptime` members

`BM`

comptime BM = epc.BM

`BN`

comptime BN = epc.BN

`CG1_TMA_BM`

comptime CG1_TMA_BM = c_smem_shape0

`CG2_TMA_BM`

comptime CG2_TMA_BM = c_smem_shape0 if (epc == Int(256)) else epc.BM

`cta_group`

comptime cta_group = epc.cta_group

`MMA_M`

comptime MMA_M = epc.MMA_M

`MMA_N`

comptime MMA_N = epc.MMA_N

`stage_n_offset`

comptime stage_n_offset = (stage * epc)

`stageN`

comptime stageN = epc.stageN

`TMA_BM`

comptime TMA_BM = c_smem_shape0 if (eq epc.MMA_M, 256) else epc.BM if (epc == Int(2)) else TMAStoreCoords[epc, c_smem_shape0, stage, batched].CG1_TMA_BM

Methods

`init`

def __init__(c_coord: Tuple[UInt32, UInt32], warp_id: UInt32) -> Self

Compute TMA store coordinates from 2D tile coords and warp ID.

def __init__(c_coord: Tuple[UInt32, UInt32, UInt32], warp_id: UInt32) -> Self

Compute TMA store coordinates from 3D tile coords and warp ID.

Fields​

Implemented traits​

comptime members​

BM​

BN​

CG1_TMA_BM​

CG2_TMA_BM​

cta_group​

MMA_M​

MMA_N​

stage_n_offset​

stageN​

TMA_BM​

Methods​

__init__​