tilelang.language.experimental.gemm_sp¶

The language interface for tl programs.

函式¶

`gemm_sp`(A_sparse, E, B, C[, transpose_A, transpose_B, ...])	Perform a Sparse General Matrix Multiplication (GEMM-sp) operation.
`gemm_sp_v2`(A_sparse, E, B, C[, transpose_A, ...])	Perform a General Matrix Multiplication (GEMM) operation.

Module Contents¶

tilelang.language.experimental.gemm_sp.gemm_sp(A_sparse, E, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, k_pack=1, wg_wait=0)¶

Perform a Sparse General Matrix Multiplication (GEMM-sp) operation.

This function computes C = A @ B where A and B can optionally be transposed. The operation supports various warp policies and accumulation modes.

參數:

A_sparse (Union[BufferLikeType, tir.Var]) -- First input matrix dense values
E (Union[BufferLikeType, tir.Var]) -- First input matrix sparse metadata
B (Union[BufferLikeType, tir.Var]) -- Second input matrix
C (Union[BufferLikeType, tir.Var]) -- Output matrix for results
transpose_A (bool, optional) -- Whether to transpose matrix A. Defaults to False.
transpose_B (bool, optional) -- Whether to transpose matrix B. Defaults to False.
policy (GemmWarpPolicy, optional) -- Warp execution policy. Defaults to GemmWarpPolicy.Square.
clear_accum (bool, optional) -- Whether to clear accumulator before computation. Defaults to False.
k_pack (int, optional) -- Number of k dimensions packed into a single warp. Defaults to 1.
wg_wait (int, optional) -- Warp group wait count. Defaults to 0.

回傳:

A handle to the GEMM operation

回傳型別:

tir.Call

引發:

AssertionError -- If the K dimensions of matrices A and B don't match

tilelang.language.experimental.gemm_sp.gemm_sp_v2(A_sparse, E, B, C, transpose_A=False, transpose_B=False, transpose_E=False, policy=GemmWarpPolicy.Square, clear_accum=False, k_pack=1, wg_wait=0)¶

Perform a General Matrix Multiplication (GEMM) operation.