tilelang.language.gemm

The language interface for tl programs.

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.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.primitives.gemm.base.GemmWarpPolicy)

• clear_accum (bool)

• k_pack (int)

• wg_wait (int)

• mbar (tvm.tir.Buffer | None)

tilelang.language.gemm.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.primitives.gemm.base.GemmWarpPolicy)

• clear_accum (bool)

• k_pack (int)

• wg_wait (int)

• mbar (tvm.tir.Buffer | None)

tilelang.language.gemm.gemm