tilelang.primitives.gemm

Submodules

• tilelang.primitives.gemm.base
• tilelang.primitives.gemm.gemm_mma

Functions

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

Package Contents

tilelang.primitives.gemm.gemm(A, B, C, transpose_A=False, transpose_B=False, block_row_warps=None, block_col_warps=None, warp_row_tiles=None, warp_col_tiles=None, chunk=None, policy=GemmWarpPolicy.Square, k_pack=1)

Parameters:

• A (tvm.tir.Buffer)

• B (tvm.tir.Buffer)

• C (tvm.tir.Buffer)

• transpose_A (bool)

• transpose_B (bool)

• block_row_warps (int | None)

• block_col_warps (int | None)

• warp_row_tiles (int | None)

• warp_col_tiles (int | None)

• chunk (int | None)

• policy (base.GemmWarpPolicy)

• k_pack (int)