tilelang.cuda.intrinsics.macro.tcgen05_macro_generator¶

Attributes¶

lift

Classes¶

TCGEN05DescriptorParams

Pre-computed parameters for TCGEN05 descriptor initialization and atom offset computation.

TCGEN05TensorMemoryParams

Pre-computed TMEM addressing for one operand Region.

TensorCoreIntrinEmitter

Intrinsic emitter for Blackwell (SM100) TCGEN5MMA instructions.

Functions¶

compute_umma_descriptor(tl_layout, buffer, transposed)

Port of CuTe make_umma_desc (mma_traits_sm100.hpp).

Module Contents¶

tilelang.cuda.intrinsics.macro.tcgen05_macro_generator.lift¶
class tilelang.cuda.intrinsics.macro.tcgen05_macro_generator.TCGEN05DescriptorParams¶

Pre-computed parameters for TCGEN05 descriptor initialization and atom offset computation.

Returned by compute_tcgen05_*_desc_params() and consumed by init_tcgen05_*_desc() and tcgen05_*_atom() methods.

swizzle_mode: tilelang.layout.SwizzleMode¶

Canonical swizzle mode; project to the descriptor field via tcgen05_layout_type().

leading_byte_offset: int¶

LBO >> 4, ready to pass to T.initialize_tcgen05_descriptor.

stride_byte_offset: int¶

SBO >> 4, ready to pass to T.initialize_tcgen05_descriptor.

swizzle_atom_elems: int¶

Number of elements per swizzle atom along the non-K dimension.

k_atom_size: int¶

max(swizzle_atom_elems // micro_size_k, 1).

elem_bits: int¶

Bit width of a single logical element.

is_k_major: bool¶

Whether the matrix is stored in K-major order (affects offset formula branching).

slice_byte_offset: object = 0¶

Physical byte offset (raw bytes) of the operand slice origin within its buffer; passed to T.increase_descriptor_offset after building the descriptor from the buffer base. 0 for a whole-buffer / base-origin operand. Computed by compute_umma_descriptor() from the CuTe layout.

class tilelang.cuda.intrinsics.macro.tcgen05_macro_generator.TCGEN05TensorMemoryParams¶

Pre-computed TMEM addressing for one operand Region.

The TMEM analog of TCGEN05DescriptorParams: built by compute_tcgen05_{a,c}_tmem_params() from the operand’s CuTe fragment and Region, and consumed by tcgen05_*_atom() to form each MMA atom’s raw TMEM address.

data: tvm.tirx.Var¶

The TMEM buffer’s data Var, carrying the base address.

origin: tuple[tvm.tirx.PrimExpr, tvm.tirx.PrimExpr]¶

Physical (datapath, value-column) of the Region origin.

tile: tilelang.layout.cute.Layout¶

Region-local matrix coordinates to physical coordinate steps.

value_bits: int¶

Bit width of one value; converts value columns to raw b32 columns.

classmethod from_region(fragment, region)¶

Restrict a TMEM fragment to one operand Region.

cute.restrict applies the Region exactly like the shared-memory descriptor path: it yields the Region origin’s physical (datapath, column) plus the sliced sublayout, and each atom origin is then one evaluation of that sublayout.

Parameters:
Return type:

TCGEN05TensorMemoryParams

atom_offset(row_offset, col_offset)¶

Raw TMEM address of the atom at a Region-local matrix origin.

The column coordinate counts values of the buffer’s dtype; scaling it to raw b32 columns here is the same conversion TMEM allocation uses.

Parameters:

  • row_offset (tvm.tirx.PrimExpr)

  • col_offset (tvm.tirx.PrimExpr)

Return type:

tvm.tirx.PrimExpr

tilelang.cuda.intrinsics.macro.tcgen05_macro_generator.compute_umma_descriptor(tl_layout, buffer, transposed, micro_size_k=16, region=None)¶

Port of CuTe make_umma_desc (mma_traits_sm100.hpp).

The Blackwell analog of compute_gmma_descriptor(): decode an arbitrary shared-memory tl_layout for buffer and compute the UMMA descriptor parameters, accepting any UMMA-canonical layout. A non-canonical layout is a programming error, so the canonicity checks assert (mirroring CuTe’s ``static_assert``s).

SBO/LBO are read in uint128 units from the canonical layout (dtype-agnostic, exactly like make_umma_desc); elem_bits is carried for the atom-offset math. The SWIZZLE_128B_BASE32B UMMA mode is not produced by tilelang’s shared-layout makers, so only the standard {NONE,32B,64B,128B} atoms appear. Unlike make_umma_desc (which sees one atom), the decoded layout here is the whole operand tile; 2SM/block-scaled atom splitting is handled by the caller.

transposed selects which logical axis is MN vs K (shape only): default [MN, K], transposed [K, MN]. The operand is required to be row-major, i.e. K-major iff not transposed; the contiguity detected from the layout is asserted to agree.

region is the operand’s per-axis ranges (one tvm.ir.Range per logical buffer mode), used to restrict the decoded layout to a sliced operand. It may be None for a full-buffer operand.

Parameters:

  • transposed (bool)

  • micro_size_k (int)

Return type:

TCGEN05DescriptorParams

class tilelang.cuda.intrinsics.macro.tcgen05_macro_generator.TensorCoreIntrinEmitter(a_dtype='float16', b_dtype='float16', accum_dtype='float16', a_transposed=False, b_transposed=False, block_row_warps=2, block_col_warps=2, warp_row_tiles=8, warp_col_tiles=8, chunk=16, reduce_k=1, num_elems_per_byte=1, is_m_first=False, thread_var=None)¶

Bases: tilelang.cuda.intrinsics.macro.mma_macro_generator.TensorCoreIntrinEmitter

Intrinsic emitter for Blackwell (SM100) TCGEN5MMA instructions.

Generates TIR macros that lower to tcgen05.mma PTX instructions for both the SS (Shared-Shared) and TS (TensorMemory-Shared) GEMM variants. Also provides layout helpers for tensor-memory (TMEM) buffers.

Parameters:

  • a_dtype (str)

  • b_dtype (str)

  • accum_dtype (str)

  • a_transposed (bool)

  • b_transposed (bool)

  • block_row_warps (int)

  • block_col_warps (int)

  • warp_row_tiles (int)

  • warp_col_tiles (int)

  • chunk (int)

  • reduce_k (int)

  • num_elems_per_byte (int)

  • is_m_first (bool)

  • thread_var (tvm.tirx.Var | None)

tcgen05_prefix: str¶
meta: tuple = ()¶
a_shared_layout: tilelang.layout.Layout = None¶
b_shared_layout: tilelang.layout.Layout = None¶
a_tmem_layout: tilelang.layout.cute.Layout = None¶
c_tmem_layout: tilelang.layout.cute.Layout = None¶
validate_tcgen05_operand_regions(A_region, B_region, C_region, *, is_ts)¶

Enforce the TCGEN5MMA operand contract at the emitter boundary.

Every operand is a complete BufferRegion whose leading modes are pinned to one element and whose two trailing matrix modes are an exact union of instruction atoms, atom-aligned in origin. CUTLASS forms sliced views only after partitioning into atoms; enforcing the same contract here — once, before any address math — prevents descriptor-aligned but atom-misaligned windows (for example BF16 K16 at K=8) from silently producing incorrect results, and lets the rest of the lowering assume well-formed Regions without re-checking.

Parameters:

  • A_region (tvm.tirx.BufferRegion)

  • B_region (tvm.tirx.BufferRegion)

  • C_region (tvm.tirx.BufferRegion)

  • is_ts (bool)

Return type:

None

compute_tcgen05_c_tmem_params(C_region)¶

Compute accumulator TMEM addressing for one operand Region.

The TMEM analog of compute_tcgen05_*_desc_params: shared by the SS, TS, and block-scaled emitters so the accumulator addressing convention lives in exactly one place.

Parameters:

C_region (tvm.tirx.BufferRegion)

Return type:

TCGEN05TensorMemoryParams

compute_tcgen05_a_tmem_params(A_region)¶

Compute TS A TMEM addressing for one operand Region.

Parameters:

A_region (tvm.tirx.BufferRegion)

Return type:

TCGEN05TensorMemoryParams

tcgen05mma(A_region, B_region, C_region, mbar, clear_accum=False)¶

Emit a TCGEN5MMA operation, dispatching to SS or TS variant based on A’s memory scope.

If A_region resides in tensor memory (shared.tmem), the TS variant is emitted; otherwise the SS variant is used (both A and B from shared memory).

Parameters:

  • A_region (BufferRegion) – Operand A — either in shared memory (SS) or tensor memory (TS).

  • B_region (BufferRegion) – Operand B in shared memory.

  • C_region (BufferRegion) – Accumulator Region in tensor memory.

  • mbar (PrimExpr) – Memory barrier used for MMA completion signalling.

  • clear_accum (PrimExpr) – Whether to zero the accumulator before the first MMA.

tcgen05mma_ss(A_region, B_region, C_region, mbar, clear_accum=False)¶

Emit the SS (Shared-Shared) variant of TCGEN5MMA.

Reads operand A and B from shared memory via a descriptor.

Parameters:

  • A_region (BufferRegion) – Operand A in shared memory.

  • B_region (BufferRegion) – Operand B in shared memory.

  • C_region (BufferRegion) – Accumulator Region in tensor memory.

  • mbar (PrimExpr) – Memory barrier for MMA completion signalling. None keeps the MMA ordered in the issue stream without publishing an event.

  • clear_accum (PrimExpr) – Whether to zero the accumulator before the first MMA.

tcgen05mma_ts(A_region, B_region, C_region, mbar, clear_accum=False)¶

Emit the TS (TensorMemory-Shared) variant of TCGEN5MMA.

Reads operand A directly from tensor memory (TMEM) and operand B from shared memory via a descriptor. Every A and C atom origin is mapped through its CuTe TMEM fragment.

Parameters:

  • A_region (BufferRegion) – Operand A residing in tensor memory (shared.tmem).

  • B_region (BufferRegion) – Operand B in shared memory.

  • C_region (BufferRegion) – Accumulator Region in tensor memory.

  • mbar (PrimExpr) – Memory barrier for MMA completion signalling. None keeps the MMA ordered in the issue stream without publishing an event.

  • clear_accum (PrimExpr) – Whether to zero the accumulator before the first MMA.

tcgen05mma_blockscaled(A_region, B_region, C_region, SFA_tmem, SFB_tmem, mbar, sf_k_start, sf_a_granularity_k, sf_b_granularity_k, clear_accum=False)¶

Emit a block-scaled TCGEN5MMA (SS variant with TMEM scale factors).

Uses tcgen05.mma.cta_group::1|2.kind::mxf8f6f4.block_scale PTX instruction. Scale factors must already reside in tensor memory.

Parameters:

  • A_region (tvm.tirx.BufferRegion)

  • B_region (tvm.tirx.BufferRegion)

  • C_region (tvm.tirx.BufferRegion)

  • sf_k_start (tvm.tirx.PrimExpr)

  • sf_a_granularity_k (int)

  • sf_b_granularity_k (int)

  • clear_accum (tvm.tirx.PrimExpr)

get_tcgen5_blockscaled_instr_desc(atom_m, atom_n, a_is_k_major, b_is_k_major, scale_in_a, scale_in_b, a_sf_id, b_sf_id)¶

Build the block-scaled instruction descriptor via FFI.

Parameters:

  • atom_m (int)

  • atom_n (int)

  • a_is_k_major (bool)

  • b_is_k_major (bool)

  • scale_in_a (int)

  • scale_in_b (int)

  • a_sf_id (int)

  • b_sf_id (int)

Return type:

tvm.tirx.PrimExpr

abstract make_mma_load_layout(local_buf, matrix='A')¶

Create a layout function for storing MMA results into a fragment buffer. This layout is used in conjunction with inverse_mma_store_layout to map fragment indices to threads and local indices.

Parameters:

  • local_buf (tirx.Buffer) – The local buffer representing a fragment of a matrix.

  • matrix (str)

Returns:

A fragment object that describes how threads and indices in local_buf are laid out.

Return type:

T.Fragment

Raises:

AssertionError – If local_buf is not detected to be a fragment buffer.

make_mma_store_layout(tmem_buf, *, operand='C', is_ts=False)¶

Create the TMEM layout of one TCGEN5MMA operand for the layout map.

Requires self.meta initialized via get_tcgen5_mma_meta(); the selected instruction atom decides the fragment.

operand="A" builds the dense TS FrgTypeA (CuTe tmem_frg_{1,2}sm<A, A, ...>): 16-bit A remains expressed in value columns here; two such columns are packed into one raw b32 TMEM column only when an MMA atom address is formed. operand="C" builds the accumulator FrgTypeC; is_ts selects the TS traits’ allocation mode when the GEMM’s A operand also lives in TMEM.

Parameters:

  • tmem_buf (tvm.tirx.Buffer)

  • operand (str)

  • is_ts (bool)

Return type:

tilelang.layout.Layout

get_tcgen5_mma_meta(m, n, k, disable_2cta, disable_ws=False)¶

Query the FFI for TCGEN5MMA atom metadata (atom_m, atom_n, atom_k, enable_ws, enable_2cta), and record them in self.meta.

Parameters:

  • m (int)

  • n (int)

  • k (int)

  • disable_2cta (bool)

  • disable_ws (bool)

get_tcgen5_instr_desc(atom_m, atom_n, atom_k, a_is_k_major, b_is_k_major, scale_in_a, scale_in_b)¶

Build the 64-bit instruction descriptor for a tcgen05.mma PTX call.

Parameters:

  • atom_m (int)

  • atom_n (int)

  • atom_k (int)

  • a_is_k_major (bool)

  • b_is_k_major (bool)

  • scale_in_a (int)

  • scale_in_b (int)

Return type:

tvm.tirx.PrimExpr

property tcgen05_meta: tilelang.cuda.intrinsics.layout.mma_sm100_layout.TCGEN05Meta¶

The selected instruction atom, as named fields.

Requires self.meta to have been set via get_tcgen5_mma_meta().

Return type:

tilelang.cuda.intrinsics.layout.mma_sm100_layout.TCGEN05Meta

property tcgen05_num_inst_m: int¶

Number of TCGEN05MMA instruction atoms along M (SS variant).

Return type:

int

property tcgen05_num_inst_n: int¶

Number of TCGEN05MMA instruction atoms along N.

Return type:

int

property tcgen05_num_k_atoms: int¶

Number of K-dimension micro-steps (chunk // micro_size_k).

Return type:

int

compute_tcgen05_b_desc_params(B_buf)¶

Compute B descriptor parameters from the B shared buffer via the CuTe make_umma_desc port. The returned TCGEN05DescriptorParams is passed to init_tcgen05_b_desc() and tcgen05_*_atom().

Parameters:

B_buf (Buffer or BufferRegion) – The B operand in shared memory.

Return type:

TCGEN05DescriptorParams

compute_tcgen05_a_desc_params(A_buf)¶

Compute A descriptor parameters from the A shared buffer (SS variant) via the CuTe make_umma_desc port.

Parameters:

A_buf (Buffer or BufferRegion) – The A operand in shared memory.

Return type:

TCGEN05DescriptorParams

init_tcgen05_b_desc(desc_b, B_buf, b_params)¶

Emit TIR to initialize a pre-allocated TCGEN05 B descriptor.

Parameters:

  • desc_b (Buffer) – A descriptor buffer allocated via T.alloc_tcgen05_smem_desc().

  • B_buf (Buffer or BufferRegion) – The B operand in shared memory.

  • b_params (TCGEN05DescriptorParams) – Pre-computed parameters from compute_tcgen05_b_desc_params().

init_tcgen05_a_desc(desc_a, A_buf, a_params)¶

Emit TIR to initialize a pre-allocated TCGEN05 A descriptor (SS variant).

Parameters:

  • desc_a (Buffer) – A descriptor buffer allocated via T.alloc_tcgen05_smem_desc().

  • A_buf (Buffer or BufferRegion) – The A operand in shared memory.

  • a_params (TCGEN05DescriptorParams) – Pre-computed parameters from compute_tcgen05_a_desc_params().

compute_tcgen05_instr_desc()¶

Compute the 64-bit instruction descriptor using current meta.

Requires self.meta to have been set via get_tcgen5_mma_meta().

Return type:

tvm.tirx.PrimExpr

tcgen05_atom_arrive(mbar)¶

Emit tcgen05_mma_arrive(mbar).

tcgen05_ss_atom(desc_a, desc_b, inst_m_idx, inst_n_idx, ki, a_params, b_params, c_params, instr_desc, clear_accum=False)¶

Emit a single TCGEN05MMA SS instruction for atom (inst_m_idx, inst_n_idx, ki).

Must be called after descriptor initialization and before tcgen05_atom_arrive().

Parameters:

  • desc_a (Buffer) – Initialized A and B descriptors.

  • desc_b (Buffer) – Initialized A and B descriptors.

  • inst_m_idx (int) – M-dimension atom index (0 .. tcgen05_num_inst_m - 1).

  • inst_n_idx (int) – N-dimension atom index (0 .. tcgen05_num_inst_n - 1).

  • ki (int) – K-dimension atom index (0 .. tcgen05_num_k_atoms - 1).

  • a_params (TCGEN05DescriptorParams) – Pre-computed A descriptor parameters.

  • b_params (TCGEN05DescriptorParams) – Pre-computed B descriptor parameters.

  • c_params (TCGEN05TensorMemoryParams) – Pre-computed accumulator TMEM parameters from compute_tcgen05_c_tmem_params().

  • instr_desc (PrimExpr) – Instruction descriptor from compute_tcgen05_instr_desc().

  • clear_accum (PrimExpr) – Whether to zero the accumulator on the first K atom.

tcgen05_ts_atom(desc_b, inst_m_idx, inst_n_idx, ki, a_params, b_params, c_params, instr_desc, clear_accum=False)¶

Emit a single TCGEN05MMA TS instruction for atom (inst_m_idx, inst_n_idx, ki).

A resides in tensor memory; B in shared memory.

Parameters:

  • desc_b (Buffer) – Initialized B descriptor.

  • inst_m_idx (int) – M-dimension atom index.

  • inst_n_idx (int) – N-dimension atom index.

  • ki (int) – K-dimension atom index.

  • a_params (TCGEN05TensorMemoryParams) – Pre-computed A TMEM parameters from compute_tcgen05_a_tmem_params().

  • b_params (TCGEN05DescriptorParams) – Pre-computed B descriptor parameters.

  • c_params (TCGEN05TensorMemoryParams) – Pre-computed accumulator TMEM parameters from compute_tcgen05_c_tmem_params().

  • instr_desc (PrimExpr) – Instruction descriptor from compute_tcgen05_instr_desc().

  • clear_accum (PrimExpr) – Whether to zero the accumulator on the first K atom.

tcgen05_blockscaled_atom(desc_a, desc_b, sfa_data, sfb_data, inst_m_idx, inst_n_idx, ki, a_params, b_params, c_params, instr_desc, clear_accum=False)¶

Emit a single TCGEN05MMA block-scaled SS instruction.

Parameters:

  • desc_a (Buffer) – Initialized A and B descriptors.

  • desc_b (Buffer) – Initialized A and B descriptors.

  • sfa_data (Var) – Scale factor data pointers in tensor memory.

  • sfb_data (Var) – Scale factor data pointers in tensor memory.

  • inst_m_idx (int) – Atom indices.

  • inst_n_idx (int) – Atom indices.

  • ki (int) – Atom indices.

  • a_params (TCGEN05DescriptorParams) – Pre-computed descriptor parameters.

  • b_params (TCGEN05DescriptorParams) – Pre-computed descriptor parameters.

  • c_params (TCGEN05TensorMemoryParams) – Pre-computed accumulator TMEM parameters from compute_tcgen05_c_tmem_params().

  • instr_desc (PrimExpr) – Block-scaled instruction descriptor (with SF IDs already encoded).

  • clear_accum (PrimExpr) – Whether to zero the accumulator on the first K atom.