tilelang.tileop.gemm.gemm_tcgen05¶

Classes¶

GemmTCGEN5

GEMM operator for Blackwell (SM100) TCGEN5MMA instructions.

Module Contents¶

class tilelang.tileop.gemm.gemm_tcgen05.GemmTCGEN5¶

Bases: tilelang.tileop.gemm.gemm_base.GemmBase

GEMM operator for Blackwell (SM100) TCGEN5MMA instructions.

Supports the SS (Shared-Shared) and TS (TensorMemory-Shared) variants. Layout inference and lowering are dispatched based on the memory scopes of operands A and B.

infer_shared_layout(continuity)¶

Infer a standard shared-memory swizzle layout for TCGEN05 operands.

Parameters:

continuity (int)

Return type:

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

infer_layout(target, thread_nums)¶

Infer swizzled layouts for operands and accumulator.

For SS: both A and B get swizzled shared-memory layouts. For TS: A and C get TMEM store layouts, B gets a swizzled shared-memory layout.

Parameters:

  • target (tvm.target.Target)

  • thread_nums (int)

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

Lower the GEMM tile-op into a TIR prim_func containing TCGEN5MMA calls.

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)