tilelang.tileop.gemm.gemm_mfma

類別

GemmMFMA

Base class for GEMM tile operators.

Module Contents

class tilelang.tileop.gemm.gemm_mfma.GemmMFMA

Bases: 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.

infer_layout(target, thread_nums)

參數:

• target (tvm.target.Target)

• thread_nums (int)

lower(layout_map, target, thread_bounds, thread_var)

參數:

• layout_map (dict)

• target (tvm.target.Target)

• thread_bounds (tvm.ir.Range)

• thread_var (tvm.tir.Var)

is_gemm_ss()

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

回傳型別:

bool

is_gemm_sr()

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

回傳型別:

bool

is_gemm_rs()

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

回傳型別:

bool

is_gemm_rr()

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

回傳型別:

bool