mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Tong WUandZhiwen Mo
11456de25f
* Update dependency version for apache-tvm-ffi in pyproject.toml to fix CI * [Math] Add `tl.infinity` operation and update Python interface for infinity handling - Implemented `infinity_op` in C++ to return infinity values for supported data types. - Registered new operation `tl.infinity` with appropriate attributes. - Updated Python interface to call the new `tl.infinity` operation instead of the previous method. * Add unit tests for `tl.infinity` operation in TileLang - Introduced a new test file `test_tilelang_language_infinity.py` to validate the behavior of the `tl.infinity` operation across multiple data types (float16, bfloat16, float32, float64). - Implemented a kernel to fill a tensor with infinity values and asserted the correctness of the output against PyTorch's `torch.inf`. * lint --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
34 lines
671 B
Python
34 lines
671 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()
|