mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
|
|
M, N, K = 128, 128, 33 # K = BLOCK_K + 1 -> 1-element residual K-tile
|
|
BM, BN, BK = 64, 64, 32
|
|
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((M, K), T.float16), B: T.Tensor((K, N), T.float16), C: T.Tensor((M, N), T.float32)):
|
|
with T.Kernel(T.ceildiv(N, BN), T.ceildiv(M, BM), threads=128) as (bx, by):
|
|
As = T.alloc_shared((BM, BK), "float16")
|
|
Bs = T.alloc_shared((BK, BN), "float16")
|
|
Cl = T.alloc_fragment((BM, BN), "float32")
|
|
T.clear(Cl)
|
|
for k in T.Pipelined(T.ceildiv(K, BK), num_stages=2):
|
|
T.copy(A[by * BM, k * BK], As)
|
|
T.copy(B[k * BK, bx * BN], Bs)
|
|
T.gemm(As, Bs, Cl)
|
|
T.copy(Cl, C[by * BM, bx * BN])
|
|
|
|
|
|
def test_issue_2365():
|
|
tilelang.compile(
|
|
main,
|
|
out_idx=[2],
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
},
|
|
) # crashes here
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_issue_2365()
|