• Tong WU's avatar
    [Feature] Add `tl.infinity` operator for infinity handling of bfloat16 (#1175) · 11456de2
    Tong WU authored
    
    
    * 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: default avatarZhiwen Mo <zm125@ic.ac.uk>
    11456de2
test_tilelang_language_infinity.py 671 Bytes