For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo struct
TMAReduceExecutor
struct TMAReduceExecutor[c_type: DType, c_smem_dim0: Int, c_smem_dim1: Int, epc: EpilogueConfig, stage_contiguous_size: Int, c_swizzle: TensorMapSwizzle, batched: Bool = False]
Execute TMA reduce-add from SMEM to GMEM.
Mirrors TMAStoreExecutor but uses cp_async_bulk_tensor_reduce_global_shared_cta (add) instead of cp_async_bulk_tensor_global_shared_cta (store). Takes a typed TMATensorTile value (not a raw pointer) so the descriptor keeps its grid_constant provenance end-to-end -- otherwise the compiler drops the constant-memory optimization and each TMA issue refetches the descriptor. Only supports non-transpose path.
Implemented traits
AnyType,
Copyable,
Deinitable,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
c_smem_shape0
comptime c_smem_shape0 = c_smem_dim0
CG1_TMA_BM
comptime CG1_TMA_BM = TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].c_smem_shape0
CG2_TMA_BM
comptime CG2_TMA_BM = TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].c_smem_shape0 if (epc == Int(256)) else epc.BM
cta_group
comptime cta_group = epc.cta_group
stageN
comptime stageN = epc.stageN
TMA_BM
comptime TMA_BM = TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].c_smem_shape0 if (eq epc.MMA_M, 256) else epc.BM if (epc == Int(2)) else TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].CG1_TMA_BM
Methods
execute
static def execute[tma_rank: Int, tile_shape: IndexList[tma_rank], desc_shape: IndexList[tma_rank]](c_smem_tile: TileTensor[c_type, address_space=AddressSpace.SHARED, linear_idx_type=c_smem_tile.linear_idx_type], store_coords: TMAStoreCoords[epc, Self.c_smem_shape0, batched], c_tma_op: TMATensorTile[c_type, tma_rank, tile_shape, desc_shape], warp_id: UInt32, lane: UInt32)
Execute TMA reduce-add from SMEM to GMEM via typed descriptor.