tilelang.tileop.gemm.gemm_base¶

Classes¶

GemmBase

Base class for GEMM tile operators.

Module Contents¶

class tilelang.tileop.gemm.gemm_base.GemmBase¶

Base class for GEMM tile operators.

Classifies the GEMM variant by the memory scopes of operands A and B (SS, SR, RS, TS, RR) and provides common property accessors for the underlying gemm_node IR node.

gemm_node: tvm.ir.base.Node¶
abstract infer_layout(target, thread_nums)¶
Parameters:

  • target (tvm.target.Target)

  • thread_nums (int)

abstract lower(layout_map, target, thread_bounds, thread_var, mbar_phase_expr=None)¶
Parameters:

  • layout_map (dict)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

  • thread_var (tvm.tir.Var)

  • mbar_phase_expr (tvm.tir.PrimExpr | None)

is_gemm_ss()¶

Return True if both A and B are in shared memory (SS variant).

Return type:

bool

is_gemm_sr()¶

Return True if A is in shared memory and B is in registers (SR variant).

Return type:

bool

is_gemm_rs()¶

Return True if A is in registers and B is in shared memory (RS variant).

Return type:

bool

is_gemm_ts()¶

Return True if A is in tensor memory and B is in shared memory (TS variant).

Return type:

bool

is_gemm_rr()¶

Return True if both A and B are in registers (RR variant).

Return type:

bool

property M: int¶
Return type:

int

property N: int¶
Return type:

int

property K: int¶
Return type:

int

property trans_A: bool¶
Return type:

bool

property trans_B: bool¶
Return type:

bool

property in_dtype: str¶

Input data type for the multiplication.

For the TS variant, A resides in TMEM with the accumulator dtype, so the actual input dtype is derived from B.

Return type:

str

property accum_dtype: str¶
Return type:

str

property chunk: int¶
Return type:

int

property A: tvm.tir.Buffer¶
Return type:

tvm.tir.Buffer

property B: tvm.tir.Buffer¶
Return type:

tvm.tir.Buffer

property C: tvm.tir.Buffer¶
Return type:

tvm.tir.Buffer

property ARegion¶
property BRegion¶
property CRegion¶
property stride_A: int¶
Return type:

int

property stride_B: int¶
Return type:

int

property offset_A: int¶
Return type:

int

property offset_B: int¶
Return type:

int

property clear_accum: tvm.ir.PrimExpr¶
Return type:

tvm.ir.PrimExpr

property k_pack: int¶
Return type:

int

property wg_wait: int¶
Return type:

int

property is_tcgen05: bool¶
Return type:

bool

property policy: tilelang.tileop.base.GemmWarpPolicy¶
Return type:

tilelang.tileop.base.GemmWarpPolicy

property mbarptr: tvm.ir.PrimExpr¶
Return type:

tvm.ir.PrimExpr

property mbar: tvm.tir.BufferLoad | None¶
Return type:

tvm.tir.BufferLoad | None

property C_coords¶
get_region_base_offsets(region)¶

Get the base offset (start index) for each dimension from a BufferRegion.

For example, if region is A_shared[ko % 2, 0:128, 0:64], this returns [ko % 2, 0, 0]

Parameters:

region – BufferRegion object

Returns:

List of PrimExpr representing the base offset for each dimension

property A_base_offsets¶

Get base offsets for each dimension of A region

property B_base_offsets¶

Get base offsets for each dimension of B region

property C_base_offsets¶

Get base offsets for each dimension of C region