tilelang.tileop.gemm_sp.gemm_sp_wgmma

Classes

GemmSPWGMMA

Module Contents

class tilelang.tileop.gemm_sp.gemm_sp_wgmma.GemmSPWGMMA

Bases: tilelang.tileop.gemm_sp.gemm_sp_base.GemmSPBase

infer_shared_layout(continuity)

Parameters:

continuity (int)

Return type:

collections.abc.Callable[[tvm.tir.Buffer], tilelang.layout.Layout]

infer_layout(target, thread_nums)

Parameters:

• target (tvm.target.Target)

• thread_nums (int)

lower(layout_map, target, thread_nums, thread_var, mbar_phase_expr=None)

Parameters:

• layout_map (dict)

• target (tvm.target.Target)

• thread_nums (tvm.ir.Range)

• thread_var (tvm.tir.Var)

• mbar_phase_expr (tvm.tir.PrimExpr | None)

is_gemm_ss()

Return type:

bool

is_gemm_sr()

Return type:

bool

is_gemm_rs()

Return type:

bool

is_gemm_rr()

Return type:

bool