tilelang.primitives.gemm package

Submodules

• tilelang.primitives.gemm.base module
  • GemmBaseParams
    • GemmBaseParams.A
    • GemmBaseParams.B
    • GemmBaseParams.C
    • GemmBaseParams.block_col_warps
    • GemmBaseParams.block_row_warps
    • GemmBaseParams.chunk
    • GemmBaseParams.class_attributes
    • GemmBaseParams.get_warp_size()
    • GemmBaseParams.infer_block_partition()
    • GemmBaseParams.k_pack
    • GemmBaseParams.params_as_dict()
    • GemmBaseParams.policy
    • GemmBaseParams.transpose_A
    • GemmBaseParams.transpose_B
    • GemmBaseParams.warp_col_tiles
    • GemmBaseParams.warp_row_tiles
  • GemmWarpPolicy
    • GemmWarpPolicy.FullCol
    • GemmWarpPolicy.FullRow
    • GemmWarpPolicy.Square
    • GemmWarpPolicy.compute_warp_partition()
    • GemmWarpPolicy.from_warp_partition()
    • GemmWarpPolicy.is_full_col()
    • GemmWarpPolicy.is_full_row()
    • GemmWarpPolicy.is_square()
    • GemmWarpPolicy.to_prime_factors()
• tilelang.primitives.gemm.gemm_mma module
  • GemmPrimitiveMMA
    • GemmPrimitiveMMA.accum_dtype
    • GemmPrimitiveMMA.gemm_rrr()
    • GemmPrimitiveMMA.gemm_rsr()
    • GemmPrimitiveMMA.gemm_srr()
    • GemmPrimitiveMMA.gemm_ssr()
    • GemmPrimitiveMMA.in_dtype
    • GemmPrimitiveMMA.invoke()

Module contents

tilelang.primitives.gemm.gemm(A: Buffer, B: Buffer, C: Buffer, transpose_A: bool = False, transpose_B: bool = False, block_row_warps: Optional[int] = None, block_col_warps: Optional[int] = None, warp_row_tiles: Optional[int] = None, warp_col_tiles: Optional[int] = None, chunk: Optional[int] = None, policy: GemmWarpPolicy = GemmWarpPolicy.Square, k_pack: int = 1)