mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Elevator14BandLeiWang1999
49c857154e
* Fix various issues under int64_t static and dynamic shape. * Resolve reviewed issues. * Add unit test. * fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
67 lines
1.6 KiB
Python
67 lines
1.6 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def fill_symbolic(value: float, dtype="bfloat16"):
|
|
n = T.symbolic("n", "int64")
|
|
block_n = 512
|
|
|
|
@T.prim_func
|
|
def main(x: T.Tensor[n, dtype]):
|
|
# Initialize Kernel Context
|
|
with T.Kernel(T.ceildiv(n, block_n), threads=128) as bx:
|
|
# Doesn't yet work with int64-shaped global tensor
|
|
# T.fill(x[bx * block_n : (bx + 1) * block_n], value)
|
|
for i in T.Parallel(block_n):
|
|
x[bx * block_n + i] = value
|
|
|
|
return main
|
|
|
|
|
|
def run_fill_symbolic(n: int):
|
|
import torch
|
|
|
|
x = torch.zeros(n, dtype=torch.bfloat16, device="cuda")
|
|
fill_symbolic(1.0)(x)
|
|
assert x.min() == 1.0 and x.max() == 1.0
|
|
|
|
|
|
def test_fill_symbolic():
|
|
# Requires 8GB VRAM
|
|
run_fill_symbolic(2**32)
|
|
|
|
|
|
@tilelang.jit
|
|
def fill_static(n: int, value: float, dtype="bfloat16"):
|
|
block_n = 512
|
|
|
|
@T.prim_func
|
|
def main(x: T.Tensor[n, dtype]):
|
|
# Initialize Kernel Context
|
|
with T.Kernel(T.ceildiv(n, block_n), threads=128) as bx:
|
|
# Doesn't yet work with int64-shaped global tensor
|
|
# T.fill(x[bx * block_n : (bx + 1) * block_n], value)
|
|
for i in T.Parallel(block_n):
|
|
x[bx * block_n + i] = value
|
|
|
|
return main
|
|
|
|
|
|
def run_fill_static(n: int):
|
|
import torch
|
|
|
|
x = torch.zeros(n, dtype=torch.bfloat16, device="cuda")
|
|
fill_static(n, 1.0)(x)
|
|
assert x.min() == 1.0 and x.max() == 1.0
|
|
|
|
|
|
def test_fill_static():
|
|
# Requires 8GB VRAM
|
|
run_fill_static(2**32)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_fill_symbolic()
|
|
test_fill_static()
|