mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Gongen-Ali
4e53b50750
* Fix some bugs on ci and open rocm ci test * fix pre-commit * fix build issue on hip * refactor on hip * pre-commit fix * fix test with pdl * fix pdl test
34 lines
694 B
Python
34 lines
694 B
Python
import torch
|
|
import tilelang
|
|
import tilelang.testing
|
|
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(T.float16)
|
|
_test_infinity(T.bfloat16)
|
|
_test_infinity(T.float32)
|
|
_test_infinity(T.float64)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|