mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-26 07:22:54 +08:00
33 lines
670 B
Python
33 lines
670 B
Python
import torch
|
|
import tilelang
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit(out_idx=-1)
|
|
def get_inf_kernel(dtype: str):
|
|
@T.prim_func
|
|
def main(A: T.Tensor((32,), dtype)):
|
|
with T.Kernel(1, threads=32):
|
|
T.fill(A, T.infinity(dtype))
|
|
|
|
return main
|
|
|
|
|
|
def _test_infinity(dtype: str):
|
|
kernel = get_inf_kernel(dtype)
|
|
output = kernel()
|
|
|
|
assert torch.all(output == torch.inf), f"check failed for {dtype=}"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_infinity():
|
|
_test_infinity("float16")
|
|
_test_infinity("bfloat16")
|
|
_test_infinity("float32")
|
|
_test_infinity("float64")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|