mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Xuyang Li
d36ec37c59
`T.infinity` aborted with an internal `LOG(FATAL)` ("Cannot decide
infinity for type float8_e5m2") because `infinity_op` had no fp8 branch.
e5m2 is the only fp8 format with a representable inf, so add a branch
returning +inf, and emit it in codegen via its bit pattern
(`::bitcast(0x7c)`, mirroring the tfloat32 path) since inf/NaN cannot be
spelled as a float literal for the fp8 constructor. Other fp8 formats
(e4m3*, e3m4, e8m0, finite fnuz incl. e5m2fnuz) have no inf and keep the
fatal. Adds e5m2 to the infinity test.
35 lines
728 B
Python
35 lines
728 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)
|
|
_test_infinity(T.float8_e5m2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|