tilelang.language.gemm_op¶

GEMM (General Matrix Multiplication) operators exposed on the TileLang language surface.

Attributes¶

gemm

Functions¶

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

GEMM v1: use op tl.gemm.

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

GEMM v2: use op tl.gemm_py.

Module Contents¶

tilelang.language.gemm_op.gemm_v1(A, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, k_pack=1, wg_wait=0, mbar=None)¶

GEMM v1: use op tl.gemm.

Parameters:

  • A (tvm.tir.Buffer | tvm.tir.Var)

  • B (tvm.tir.Buffer | tvm.tir.Var)

  • C (tvm.tir.Buffer | tvm.tir.Var)

  • transpose_A (bool)

  • transpose_B (bool)

  • policy (tilelang.tileop.base.GemmWarpPolicy)

  • clear_accum (bool)

  • k_pack (int)

  • wg_wait (int)

  • mbar (tvm.tir.Buffer | None)

tilelang.language.gemm_op.gemm_v2(A, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, k_pack=1, wg_wait=0, mbar=None)¶

GEMM v2: use op tl.gemm_py.

Parameters:

  • A (tvm.tir.Buffer | tvm.tir.Var)

  • B (tvm.tir.Buffer | tvm.tir.Var)

  • C (tvm.tir.Buffer | tvm.tir.Var)

  • transpose_A (bool)

  • transpose_B (bool)

  • policy (tilelang.tileop.base.GemmWarpPolicy)

  • clear_accum (bool)

  • k_pack (int)

  • wg_wait (int)

  • mbar (tvm.tir.Buffer | None)

tilelang.language.gemm_op.gemm¶