tilelang.tileop.gemm.gemm_tcgen05

類別

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_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.

參數:

  • target (tvm.target.Target)

  • thread_nums (int)

lower(layout_map, target, thread_bounds, thread_var)

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

參數:

  • layout_map (dict)

  • target (tvm.target.Target)

  • thread_bounds (tvm.ir.Range)

  • thread_var (tvm.tir.Var)