tilelang.metal.op.gemm.gemm_metal¶

Attributes¶

GEMM_INST_METAL

GEMM_INST_METAL_COOPERATIVE_TENSOR

Classes¶

GemmMetalSimdGroup

Base class for GEMM tile operators.

GemmMetal

Base class for GEMM tile operators.

Module Contents¶

tilelang.metal.op.gemm.gemm_metal.GEMM_INST_METAL = 'metal.simdgroup'¶
tilelang.metal.op.gemm.gemm_metal.GEMM_INST_METAL_COOPERATIVE_TENSOR = 'metal.cooperative_tensor'¶
class tilelang.metal.op.gemm.gemm_metal.GemmMetalSimdGroup¶

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.

is_gemm_ss()¶

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

Return type:

bool

infer_layout(target, thread_nums)¶
Parameters:

  • target (tvm.target.Target)

  • thread_nums (int)

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

  • layout_map (dict)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

  • thread_index (tvm.tirx.PrimExpr)

  • mbar_phase_expr (tvm.tirx.PrimExpr | None)

class tilelang.metal.op.gemm.gemm_metal.GemmMetal¶

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.

is_gemm_ss()¶

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

Return type:

bool

is_gemm_gg()¶
Return type:

bool

infer_layout(target, thread_nums)¶
Parameters:

  • target (tvm.target.Target)

  • thread_nums (int)

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

  • layout_map (dict)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

  • thread_index (tvm.tirx.PrimExpr)

  • mbar_phase_expr (tvm.tirx.PrimExpr | None)