tilelang.cuda.language.gemm_op

CUDA dialect of T.gemm: the common GEMM plus CUDA-specific knobs.

Functions

gemm(A, B, C[, transpose_A, transpose_B, policy, ...])

TileLang GEMM operator for CUDA.

gemm_sp(A_sparse, E, B, C[, transpose_A, transpose_E, ...])

Sparse GEMM (2:4 structured sparsity) for CUDA.

Module Contents

tilelang.cuda.language.gemm_op.gemm(A, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, mbar=None, annotations=None)

TileLang GEMM operator for CUDA.

Same semantics as the common tilelang.language.gemm_op.gemm(): the default synchronous GEMM. On Hopper, if the compiler selects WGMMA lowering, TileLang inserts the corresponding wait implicitly. On Blackwell TCGEN5MMA, TileLang inserts the corresponding mbarrier_wait_parity(...) implicitly after issue.

For manual asynchronous scheduling, use T.wgmma_gemm(...) with T.wait_wgmma(...) on Hopper, or T.tcgen05_gemm(...) with T.mbarrier_wait_parity(...) on Blackwell.

Parameters:

  • A (BufferLikeType, i.e. Buffer | BufferLoad | BufferRegion, or Var) – Input buffer A.

  • B (BufferLikeType) – Input buffer B.

  • C (BufferLikeType) – Output buffer C.

  • transpose_A (bool) – Whether to transpose A. Defaults to False.

  • transpose_B (bool) – Whether to transpose B. Defaults to False.

  • policy (GemmWarpPolicy) – GEMM warp partition policy.

  • clear_accum (bool) – Whether to clear the accumulator.

  • mbar (BarrierType, i.e. Buffer | BufferLoad, or Var, optional) – Mbarrier in Blackwell. Required when this GEMM lowers to TCGEN5MMA. Defaults to None.

  • annotations (Optional[dict]) – Additional annotations.

Returns:

A handle to the GEMM operation.

Return type:

tirx.Call

tilelang.cuda.language.gemm_op.gemm_sp(A_sparse, E, B, C, transpose_A=False, transpose_E=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, wg_wait=0, annotations=None)

Sparse GEMM (2:4 structured sparsity) for CUDA.

Same semantics as the common tilelang.language.experimental.gemm_sp_op.gemm_sp(). wg_wait is the Hopper warpgroup wait count consumed when the WGMMA SP lowering is selected (-1 defers the wait to an explicit T.wait_wgmma); it rides in the tile-op annotations.

Parameters:

  • A_sparse (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Compressed sparse matrix containing only non-zero elements.

  • E (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Metadata tensor encoding the sparsity pattern of A.

  • B (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Dense input matrix.

  • C (tilelang._typing.BufferLikeType | tvm.tirx.Var) – Output accumulator matrix.

  • transpose_A (bool) – Whether to transpose A. Defaults to False.

  • transpose_E (bool) – Whether to transpose E. Defaults to False.

  • transpose_B (bool) – Whether to transpose B. Defaults to False.

  • policy (tilelang.language.gemm_op.GemmWarpPolicy) – Warp partition policy. Defaults to GemmWarpPolicy.Square.

  • clear_accum (bool) – Whether to zero the accumulator before computation. Defaults to False.

  • wg_wait (int) – Warp group wait count. Defaults to 0.

  • annotations (dict | None) – Additional annotations; values in it take precedence.

Returns:

A handle to the sparse GEMM operation.

Return type:

tirx.Call