tilelang.intrinsics.tcgen05_macro_generator¶

Attributes¶

lift

Classes¶

SwizzleMode

Enum where members are also (and must be) ints

TensorCoreIntrinEmitter

To eliminate Python syntax within TIR Macro.

Module Contents¶

tilelang.intrinsics.tcgen05_macro_generator.lift¶
class tilelang.intrinsics.tcgen05_macro_generator.SwizzleMode¶

Bases: enum.IntEnum

Enum where members are also (and must be) ints

NONE = 0¶
SWIZZLE_128B = 2¶
SWIZZLE_64B = 4¶
SWIZZLE_32B = 6¶
is_none()¶
Return type:

bool

is_swizzle_32b()¶
Return type:

bool

is_swizzle_64b()¶
Return type:

bool

is_swizzle_128b()¶
Return type:

bool

swizzle_byte_size()¶
Return type:

int

swizzle_atom_size()¶
Return type:

int

class tilelang.intrinsics.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.intrinsics.mma_macro_generator.TensorCoreIntrinEmitter

To eliminate Python syntax within TIR Macro.

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.tir.Var | None)

tcgen05_prefix: str¶
a_shared_layout: tilelang.layout.Layout = None¶
b_shared_layout: tilelang.layout.Layout = None¶
tcgen05mma(A_buf, B_buf, C_local_buf, mbar, clear_accum=False)¶
Parameters:

  • A_buf (tvm.tir.Buffer)

  • B_buf (tvm.tir.Buffer)

  • C_local_buf (tvm.tir.Buffer)

  • clear_accum (tvm.tir.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 (tir.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)¶

Create the TCGEN5 tensor-memory layout used to store MMA accumulators.

Parameters:

tmem_buf (tir.Buffer) – The local buffer representing tensormemory of a mma’s output

Returns:

Layout object describing how logical (i, j) coordinates map to the swizzled tensor-memory offsets required by TCGEN5MMA.

Return type:

Layout

Raises:

AssertionError – If tmem_buf is not detected to be a tensor-memory buffer.

get_tcgen5_mma_meta(m, n, k)¶
Parameters:

  • m (int)

  • n (int)

  • k (int)

get_tcgen5_instr_desc(atom_m, atom_n, atom_k, a_is_k_major, b_is_k_major, scale_in_a, scale_in_b)¶
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.tir.PrimExpr