mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import torch
|
|
import pytest
|
|
import tilelang.testing
|
|
|
|
|
|
@tilelang.jit
|
|
def tilelang_rand_1d(M=1024, seed=42, generator="curandStatePhilox4_32_10_t"):
|
|
num_per_thread = 128
|
|
threads = 1
|
|
blk_M = num_per_thread * threads
|
|
|
|
@T.prim_func
|
|
def rand_kernel(
|
|
A: T.Tensor((M,), "uint32"),
|
|
B: T.Tensor((M,), "float32"),
|
|
C: T.Tensor((M,), "float64"),
|
|
D: T.Tensor((M,), "float32"),
|
|
E: T.Tensor((M,), "float64"),
|
|
):
|
|
with T.Kernel(T.ceildiv(M, threads * num_per_thread), threads=threads) as bx:
|
|
tx = T.get_thread_binding()
|
|
T.rng_init(seed, 0, bx * blk_M + tx * num_per_thread, generator=generator)
|
|
for i, j in T.Parallel(threads, num_per_thread):
|
|
offsets = (bx * threads + i) * num_per_thread
|
|
idx = offsets + j
|
|
if idx < M:
|
|
A[idx] = T.rng_rand()
|
|
for i, j in T.Parallel(threads, num_per_thread):
|
|
offsets = (bx * threads + i) * num_per_thread
|
|
idx = offsets + j
|
|
if idx < M:
|
|
B[idx] = T.rng_rand_float()
|
|
for i, j in T.Parallel(threads, num_per_thread):
|
|
offsets = (bx * threads + i) * num_per_thread
|
|
idx = offsets + j
|
|
if idx < M:
|
|
C[idx] = T.rng_rand_float(bit=64)
|
|
for i, j in T.Parallel(threads, num_per_thread):
|
|
offsets = (bx * threads + i) * num_per_thread
|
|
idx = offsets + j
|
|
if idx < M:
|
|
D[idx] = T.rng_rand_float(dist="normal")
|
|
for i, j in T.Parallel(threads, num_per_thread):
|
|
offsets = (bx * threads + i) * num_per_thread
|
|
idx = offsets + j
|
|
if idx < M:
|
|
E[idx] = T.rng_rand_float(bit=64, dist="normal")
|
|
|
|
return rand_kernel
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
"M, seed, generator", [(1024, 42, "curandStateMRG32k3a_t"), (512, 123, "curandStatePhilox4_32_10_t"), (128, 0, "curandStateXORWOW_t")]
|
|
)
|
|
def test_rand_1d(M, seed, generator):
|
|
kernel = tilelang_rand_1d(M, seed, generator)
|
|
A = torch.empty(M, dtype=torch.uint32, device="cuda")
|
|
B = torch.empty(M, dtype=torch.float32, device="cuda")
|
|
C = torch.empty(M, dtype=torch.float64, device="cuda")
|
|
D = torch.empty(M, dtype=torch.float32, device="cuda")
|
|
E = torch.empty(M, dtype=torch.float64, device="cuda")
|
|
kernel(A, B, C, D, E)
|
|
|
|
|
|
@tilelang.jit
|
|
def tilelang_rand_blockwise(M=64, seed=42, generator="curandStatePhilox4_32_10_t"):
|
|
threads = 32
|
|
|
|
@T.prim_func
|
|
def rand_kernel(A: T.Tensor((M,), "uint32")):
|
|
with T.Kernel(M, threads=threads) as bx:
|
|
tx = T.get_thread_binding()
|
|
T.rng_init(seed, 0, bx, generator=generator)
|
|
if tx == 0:
|
|
A[bx] = T.rng_rand()
|
|
|
|
return rand_kernel
|
|
|
|
|
|
@tilelang.jit
|
|
def tilelang_rand_guarded_cumsum(M=64, seed=42, generator="curandStatePhilox4_32_10_t"):
|
|
threads = 32
|
|
|
|
@T.prim_func
|
|
def rand_kernel(
|
|
A: T.Tensor((M,), "uint32"),
|
|
n: T.int32,
|
|
):
|
|
with T.Kernel(M, threads=threads) as bx:
|
|
tx = T.get_thread_binding()
|
|
s = T.alloc_shared((threads,), "int32")
|
|
# rng_init inside a runtime guard, with a shared-memory cumsum
|
|
# between init and use: sync legalization hoists __syncthreads()
|
|
# out of the guard and splits it into sibling blocks, so the
|
|
# curand state must be declared at function scope to stay visible.
|
|
if bx < n:
|
|
T.rng_init(seed, 0, bx, generator=generator)
|
|
s[tx] = 1
|
|
T.cumsum(s, dim=0)
|
|
if tx == 0:
|
|
A[bx] = T.rng_rand() + T.cast(s[threads - 1], "uint32") * 0
|
|
|
|
return rand_kernel
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("generator", ["curandStateMRG32k3a_t", "curandStatePhilox4_32_10_t", "curandStateXORWOW_t"])
|
|
def test_rand_init_in_split_guard(generator):
|
|
M, seed, n = 64, 42, 37
|
|
guarded = tilelang_rand_guarded_cumsum(M, seed, generator)
|
|
baseline = tilelang_rand_blockwise(M, seed, generator)
|
|
|
|
sentinel = 0xDEADBEEF
|
|
A = torch.full((M,), sentinel, dtype=torch.uint32, device="cuda")
|
|
guarded(A, n)
|
|
A_ref = torch.empty(M, dtype=torch.uint32, device="cuda")
|
|
baseline(A_ref)
|
|
|
|
assert torch.equal(A[:n], A_ref[:n]), "guarded rng output differs from unguarded baseline"
|
|
assert (A[n:] == sentinel).all(), "rows outside the guard must stay untouched"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|