Repository: taichi-dev/taichi (language: python)

Issue:
Code generation issue with ti.frexp
Dear Taichi team, 

I hope this message finds you well. I encountered this codegen bug while writing some ray marching code using Taichi. Unfortunately turning `advanced_optimization` off does not solve the issue.

**Describe the bug**
The following 11-line kernel fails to compile and produces a `[codegen_cuda.cpp:emit_extra_unary@247] Not supported.` error even when compiling with `advanced_optimization=False`. The kernel is a custom ray marching function, which I have stripped down to the minimum. The `mip_from_pos, mip_from_dt` functions are inspired from standard C raymarching libraries.

**To Reproduce**
```py
import taichi as ti
ti.init(arch=ti.cuda, debug = True, advanced_optimization=False)

@ti.func
def mip_from_pos(x, y, z, C):
    mx = max(abs(x), abs(y), abs(z))
    _, exponent = ti.frexp(mx)
    return min(C-1, max(0, exponent))

@ti.func
def mip_from_dt(dt, H, C):
    mx = dt*H*0.5
    _, exponent = ti.frexp(mx)
    return min(C-1, max(0, exponent))

@ti.kernel
def faulty_kernel(bound:ti.f32, C:ti.u32, H:ti.u32):
    level = max(mip_from_pos(0.0, 0.0, 0.0, C), mip_from_dt(0.0, H, C))
    mip_bound = min(2 ** level, bound)
    tt = mip_bound
    
    if tt >= 0:
        tt-=1

faulty_kernel(0.5, 3, 128)
```

**Log/Screenshots**
```
$ python3 ./test_driver.py 
[Taichi] version 1.8.0, llvm 15.0.4, commit 5a994612, linux, python 3.9.18
[Taichi] Starting on arch=cuda
[E 01/25/24 01:47:58.610 43698] [codegen_cuda.cpp:emit_extra_unary@247] Not supported.


***********************************
* Taichi Compiler Stack Traceback *
***********************************
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x4805884) [0x7f6111016884]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e7e69b) [0x7f610e68f69b]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e95dfc) [0x7f610e6a6dfc]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e7ac4d) [0x7f610e68bc4d]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e95dfc) [0x7f610e6a6dfc]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1ec4575) [0x7f610e6d5575]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1ec467d) [0x7f610e6d567d]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e76f3d) [0x7f610e687f3d]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x198f8e2) [0x7f610e1a08e2]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1a898e8) [0x7f610e29a8e8]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x48bcc20) [0x7f61110cdc20]
/lib/x86_64-linux-gnu/libpthread.so.0(+0x7ea7) [0x7f6120674ea7]
/lib/x86_64-linux-gnu/libc.so.6: clone

Internal error occurred. Check out this page for possible solutions:
https://docs.taichi-lang.org/docs/install
terminate called after throwing an instance of 'std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >'
[E 01/25/24 01:47:58.611 43698] Received signal 6 (Aborted)


***********************************
* Taichi Compiler Stack Traceback *
***********************************
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x4805884) [0x7f6111016884]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1b35d0e) [0x7f610e346d0e]
/lib/x86_64-linux-gnu/libc.so.6(+0x38d60) [0x7f6120382d60]
/lib/x86_64-linux-gnu/libc.so.6: gsignal
/lib/x86_64-linux-gnu/libc.so.6: abort
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x48373df) [0x7f61110483df]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x4835a56) [0x7f6111046a56]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x4835ac1) [0x7f6111046ac1]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x4835c15) [0x7f6111046c15]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x4805960) [0x7f6111016960]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e7e69b) [0x7f610e68f69b]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e95dfc) [0x7f610e6a6dfc]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e7ac4d) [0x7f610e68bc4d]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e95dfc) [0x7f610e6a6dfc]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1ec4575) [0x7f610e6d5575]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1ec467d) [0x7f610e6d567d]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1e76f3d) [0x7f610e687f3d]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x198f8e2) [0x7f610e1a08e2]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x1a898e8) [0x7f610e29a8e8]
/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/_lib/core/taichi_python.cpython-39-x86_64-linux-gnu.so(+0x48bcc20) [0x7f61110cdc20]
/lib/x86_64-linux-gnu/libpthread.so.0(+0x7ea7) [0x7f6120674ea7]
/lib/x86_64-linux-gnu/libc.so.6: clone

Internal error occurred. Check out this page for possible solutions:
https://docs.taichi-lang.org/docs/install
```

**Additional comments**
Output for `ti diagnose `

```
No LSB modules are available.
[Taichi] version 1.8.0, llvm 15.0.4, commit 5a994612, linux, python 3.9.18

*******************************************
**      Taichi Programming Language      **
*******************************************

Docs:   https://docs.taichi-lang.org/
GitHub: https://github.com/taichi-dev/taichi/
Forum:  https://forum.taichi.graphics/

Taichi system diagnose:

python: 3.9.18 | packaged by conda-forge | (main, Aug 30 2023, 03:49:32) 
[GCC 12.3.0]
system: linux
executable: /home/akshatsh/DISKMOUNT/anaconda3/envs/aks/bin/python3
platform: Linux-5.10.0-27-cloud-amd64-x86_64-with-glibc2.31
architecture: 64bit ELF
uname: uname_result(system='Linux', node='instance-nov-22-newbootdisk-olddatadisk', release='5.10.0-27-cloud-amd64', version='#1 SMP Debian 5.10.205-2 (2023-12-31)', machine='x86_64')
locale: en_US.UTF-8
PATH: /usr/local/cuda/bin:/home/akshatsh/.vscode-server/bin/8b3775030ed1a69b13e4f4c628c612102e30a681/bin/remote-cli:/usr/local/cuda/bin:/usr/local/cuda/bin:/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/bin:/opt/conda/condabin:/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games
PYTHONPATH: ['/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/bin', '/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python39.zip', '/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9', '/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/lib-dynload', '/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages']

Distributor ID:	Debian
Description:	Debian GNU/Linux 11 (bullseye)
Release:	11
Codename:	bullseye



import: <module 'taichi' from '/home/akshatsh/DISKMOUNT/anaconda3/envs/aks/lib/python3.9/site-packages/taichi/__init__.py'>

cpu: True
metal: False
opengl: True
cuda: True
vulkan: True

`glewinfo` not available: [Errno 2] No such file or directory: 'glewinfo'

Thu Jan 25 01:51:45 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.23.08              Driver Version: 545.23.08    CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       On  | 00000000:00:04.0 Off |                    0 |
| N/A   72C    P0              35W /  70W |      2MiB / 15360MiB |      1%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  Tesla T4                       On  | 00000000:00:05.0 Off |                    0 |
| N/A   43C    P8              11W /  70W |      2MiB / 15360MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+

[Taichi] version 1.8.0, llvm 15.0.4, commit 5a994612, linux, python 3.9.18

[Taichi] version 1.8.0, llvm 15.0.4, commit 5a994612, linux, python 3.9.18
[Taichi] Starting on arch=x64

[Taichi] version 1.8.0, llvm 15.0.4, commit 5a994612, linux, python 3.9.18
[Taichi] Starting on arch=opengl

[Taichi] version 1.8.0, llvm 15.0.4, commit 5a994612, linux, python 3.9.18
[Taichi] Starting on arch=cuda

[Taichi] version 1.8.0, llvm 15.0.4, commit 5a994612, linux, python 3.9.18

*******************************************
**      Taichi Programming Language      **
*******************************************

Docs:   https://docs.taichi-lang.org/
GitHub: https://github.com/taichi-dev/taichi/
Forum:  https://forum.taichi.graphics/

                                   TAICHI EXAMPLES                                    
 ──────────────────────────────────────────────────────────────────────────────────── 
  0: ad_gravity               25: karman_vortex_street    50: patterns                
  1: circle_packing_image     26: keyboard                51: pbf2d                   
  2: comet                    27: laplace                 52: physarum                
  3: cornell_box              28: laplace_equation        53: poisson_disk_sampling   
  4: diff_sph                 29: mandelbrot_zoom         54: print_offset            
  5: differential_evolution   30: marching_squares        55: rasterizer              
  6: euler                    31: mass_spring_3d_ggui     56: regression              
  7: eulerfluid2d             32: mass_spring_game        57: sdf_renderer            
  8: explicit_activation      33: mass_spring_game_ggui   58: simple_derivative       
  9: export_mesh              34: mciso_advanced          59: simple_texture          
  10: export_ply              35: mgpcg                   60: simple_uv               
  11: export_videos           36: mgpcg_advanced          61: snow_phaseField         
  12: fem128                  37: minimal                 62: stable_fluid            
  13: fem128_ggui             38: minimization            63: stable_fluid_ggui       
  14: fem99                   39: mpm128                  64: stable_fluid_graph      
  15: fractal                 40: mpm128_ggui             65: taichi_bitmasked        
  16: fractal3d_ggui          41: mpm3d                   66: taichi_dynamic          
  17: fullscreen              42: mpm3d_ggui              67: taichi_logo             
  18: game_of_life            43: mpm88                   68: taichi_ngp              
  19: gui_image_io            44: mpm88_graph             69: taichi_sparse           
  20: gui_widgets             45: mpm99                   70: texture_graph           
  21: implicit_fem            46: mpm_lagrangian_forces   71: tutorial                
  22: implicit_mass_spring    47: nbody                   72: two_stream_instability  
  23: initial_value_problem   48: odop_solar              73: vortex_rings            
  24: jacobian                49: oit_renderer            74: waterwave               
 ──────────────────────────────────────────────────────────────────────────────────── 
42
Running example minimal ...
[Taichi] Starting on arch=x64
42.0
>>> Running time: 0.22s

Consider attaching this log when maintainers ask about system information.
>>> Running time: 6.56s
```

Please feel free to reach out if you require any additional clarifications or logs from my end. Thank you for taking the time to review this report!


Discussion / hints:

Fix patch:
diff --git a/taichi/transforms/simplify.cpp b/taichi/transforms/simplify.cpp
index 8890394637b36..550ea8220303b 100644
--- a/taichi/transforms/simplify.cpp
+++ b/taichi/transforms/simplify.cpp
@@ -133,6 +133,16 @@ class BasicBlockSimplify : public IRVisitor {
     }
   }
 
+  void visit(UnaryOpStmt *stmt) override {
+    if (stmt->op_type == UnaryOpType::abs) {
+      auto operand_type = stmt->operand->ret_type;
+      if (is_integral(operand_type) && is_unsigned(operand_type)) {
+        // abs(u) -> u
+        stmt->replace_usages_with(stmt->operand);
+        modifier.erase(stmt);
+      }
+    }
+  }
   template <typename T>
   static bool identical_vectors(const std::vector<T> &a,
                                 const std::vector<T> &b) {


Test patch:
diff --git a/tests/python/test_abs.py b/tests/python/test_abs.py
index 8ff4bc35aeddd..ab86b392507f6 100644
--- a/tests/python/test_abs.py
+++ b/tests/python/test_abs.py
@@ -78,3 +78,13 @@ def foo(x: ti.i64) -> ti.i64:
 
     for x in [-(2**40), 0, 2**40]:
         assert foo(x) == abs(x)
+
+
+@test_utils.test()
+def test_abs_u32():
+    @ti.kernel
+    def foo(x: ti.u32) -> ti.u32:
+        return abs(x)
+
+    for x in [0, 2**20]:
+        assert foo(x) == abs(x)