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

Issue:
ti.bit_shr is broken when the input is unsigned
<!-- We've collected some common issue solutions in https://docs.taichi-lang.org/docs/install. Make sure you've check them out first. Hopefully they could address your problem. -->

**Describe the bug**
ti.bit_shr is broken when the input is unsigned because we only transform it to bit_sar(unsigned) when the input is signed, and the codegen assumes that bit_shr should not exist when the ir reaches the codegen.

https://github.com/taichi-dev/taichi/blob/6a0b63ef709ae4da836492712d7f43c2c226d479/taichi/transforms/demote_operations.cpp#L152-L175

**To Reproduce**

```py
import taichi as ti
ti.init()                                                                                                                        

@ti.kernel
def func(x: ti.u32, y: ti.i32)->ti.u32:
    return ti.bit_shr(x, y)

print(func(5, 2))
```

**Log/Screenshots**
Please post the **full log** of the program (instead of just a few lines around the error message, unless the log is > 1000 lines). This will help us diagnose what's happening. For example:

```
[Taichi] version 1.6.0, llvm 15.0.4, commit 2467d918, linux, python 3.10.0
[Taichi] Starting on arch=x64
[I 03/24/23 17:04:04.280 1754876] [codegen_llvm.cpp:visit@834] 
"binary_op_type_name(op)": "bit_shr"
[E 03/24/23 17:04:04.280 1754876] [codegen_llvm.cpp:visit@835] Not supported.
```



Discussion / hints:

Fix patch:
diff --git a/taichi/transforms/demote_operations.cpp b/taichi/transforms/demote_operations.cpp
index 980ebc79b073f..30264bbadcbfd 100644
--- a/taichi/transforms/demote_operations.cpp
+++ b/taichi/transforms/demote_operations.cpp
@@ -149,19 +149,19 @@ class DemoteOperations : public BasicStmtVisitor {
         modifier.insert_before(stmt, std::move(new_matrix));
         modifier.erase(stmt);
       }
-    } else if (stmt->op_type == BinaryOpType::bit_shr &&
-               is_integral(lhs->element_type()) &&
-               is_integral(rhs->element_type()) &&
-               is_signed(lhs->element_type())) {
+    } else if (stmt->op_type == BinaryOpType::bit_shr) {
       // @ti.func
       // def bit_shr(a, b):
       //     unsigned_a = ti.cast(a, ti.uXX)
       //     shifted = ti.bit_sar(unsigned_a, b)
       //     ret = ti.cast(shifted, ti.iXX)
       //     return ret
+      TI_ASSERT(is_integral(lhs->element_type()) &&
+                is_integral(rhs->element_type()));
       auto unsigned_cast = Stmt::make<UnaryOpStmt>(UnaryOpType::cast_bits, lhs);
+      auto lhs_type = lhs->element_type();
       unsigned_cast->as<UnaryOpStmt>()->cast_type =
-          to_unsigned(lhs->element_type());
+          is_signed(lhs_type) ? to_unsigned(lhs_type) : lhs_type;
       auto shift = Stmt::make<BinaryOpStmt>(BinaryOpType::bit_sar,
                                             unsigned_cast.get(), rhs);
       auto signed_cast =


Test patch:
diff --git a/tests/python/test_bit_operations.py b/tests/python/test_bit_operations.py
index 9cea78323aa09..26fbeb0abef4c 100644
--- a/tests/python/test_bit_operations.py
+++ b/tests/python/test_bit_operations.py
@@ -60,3 +60,12 @@ def shr(a: ti.i32, b: ti.i32) -> ti.i32:
     for i in range(n):
         offset = 0x100000000 if i > 0 else 0
         assert shr(neg_test_num, i) == (neg_test_num + offset) >> i
+
+
+@test_utils.test()
+def test_bit_shr_uint():
+    @ti.kernel
+    def func(x: ti.u32, y: ti.i32) -> ti.u32:
+        return ti.bit_shr(x, y)
+
+    assert (func(5, 2) == 1)