tilelang.tileop.gemm

Submodules

• tilelang.tileop.gemm.gemm_base
• tilelang.tileop.gemm.gemm_mfma
• tilelang.tileop.gemm.gemm_mma
• tilelang.tileop.gemm.gemm_mma_sm70
• tilelang.tileop.gemm.gemm_tcgen05
• tilelang.tileop.gemm.gemm_wgmma

Classes

GemmInst	Enum where members are also (and must be) ints
GemmPy

Functions

gemm_py_infer_layout(gemm_py, target, thread_bounds)
gemm_py_lower(gemm_py, layout_map, target, ...)

Package Contents

tilelang.tileop.gemm.gemm_py_infer_layout(gemm_py, target, thread_bounds)

tilelang.tileop.gemm.gemm_py_lower(gemm_py, layout_map, target, thread_bounds, thread_var)

class tilelang.tileop.gemm.GemmInst

Bases: enum.IntEnum

Enum where members are also (and must be) ints

MMA = 0

WGMMA = 1

TCGEN5MMA = 2

MFMA = 3

is_mma()

Return type:

bool

is_wgmma()

Return type:

bool

is_tcgen5mma()

Return type:

bool

is_mfma()

Return type:

bool

__repr__()

Return repr(self).

Return type:

str

class tilelang.tileop.gemm.GemmPy

Bases: tvm.ir.base.Node, tvm.runtime.Scriptable

property A

property B

property C

property APtr

property BPtr

property CPtr

property M

property N

property K

property trans_A

property trans_B

property stride_A

property stride_B

property offset_A

property offset_B

property clear_accum

property k_pack

property wg_wait

infer_layout(target, thread_nums)

Infer the layout for the GEMM operation based on target architecture.

Parameters:

• target (tvm.target.Target)

• thread_nums (int)

lower(layout_map, target, thread_nums, thread_var)

Lower the GEMM operation to TIR statements based on target architecture.

Parameters:

• layout_map (dict)

• target (tvm.target.Target)

• thread_nums (int)

• thread_var (tvm.tir.Var)