tilelang.primitives.gemm¶

Submodules¶

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 (Optional[int])

  • block_col_warps (Optional[int])

  • warp_row_tiles (Optional[int])

  • warp_col_tiles (Optional[int])

  • chunk (Optional[int])

  • policy (base.GemmWarpPolicy)

  • k_pack (int)