[Bug] [lang] Fix some missing operator overrides like __rlshift__ (#1978

)

* Fix some missing operator magic functions

* OFT fix tb not displayed.

python/taichi/lang/common_ops.py

@@ -114,26 +114,51 @@ def __and__(self, other):
         _taichi_skip_traceback = 1
         return ti.bit_and(self, other)
 
+    def __rand__(self, other):
+        import taichi as ti
+        _taichi_skip_traceback = 1
+        return ti.bit_and(other, self)
+
     def __or__(self, other):
         import taichi as ti
         _taichi_skip_traceback = 1
         return ti.bit_or(self, other)
 
+    def __ror__(self, other):
+        import taichi as ti
+        _taichi_skip_traceback = 1
+        return ti.bit_or(other, self)
+
     def __xor__(self, other):
         import taichi as ti
         _taichi_skip_traceback = 1
         return ti.bit_xor(self, other)
 
+    def __rxor__(self, other):
+        import taichi as ti
+        _taichi_skip_traceback = 1
+        return ti.bit_xor(other, self)
+
     def __lshift__(self, other):
         import taichi as ti
         _taichi_skip_traceback = 1
         return ti.bit_shl(self, other)
 
+    def __rlshift__(self, other):
+        import taichi as ti
+        _taichi_skip_traceback = 1
+        return ti.bit_shl(other, self)
+
     def __rshift__(self, other):
         import taichi as ti
         _taichi_skip_traceback = 1
         return ti.bit_sar(self, other)
 
+    def __rrshift__(self, other):
+        import taichi as ti
+        _taichi_skip_traceback = 1
+        return ti.bit_sar(other, self)
+
     def logical_and(self, other):
         import taichi as ti
         _taichi_skip_traceback = 1
@@ -229,6 +254,18 @@ def __imod__(self, other):
         self.assign(ti.mod(self, other))
         return self
 
+    def __ilshift__(self, other):
+        import taichi as ti
+        _taichi_skip_traceback = 1
+        self.assign(ti.bit_shl(self, other))
+        return self
+
+    def __irshift__(self, other):
+        import taichi as ti
+        _taichi_skip_traceback = 1
+        self.assign(ti.bit_shr(self, other))
+        return self
+
     def assign(self, other):
         import taichi as ti
         _taichi_skip_traceback = 1
@@ -254,6 +291,10 @@ def augassign(self, x, op):
             self |= x
         elif op == 'BitXor':
             self ^= x
+        elif op == 'RShift':
+            self >>= x
+        elif op == 'LShift':
+            self <<= x
         else:
             assert False, op
 

taichi/transforms/type_check.cpp

@@ -65,9 +65,10 @@ class TypeCheck : public IRVisitor {
     TI_ASSERT(stmt->width() == 1);
     if (stmt->val->ret_type != stmt->dest->ret_type.ptr_removed()) {
       // TODO: make sure the ptr_removed type is indeed a numerical type
-      TI_WARN("[{}] Atomic add ({} to {}) may lose precision.", stmt->name(),
+      TI_WARN("[{}] Atomic add ({} to {}) may lose precision, at", stmt->name(),
               data_type_name(stmt->val->ret_type),
               data_type_name(stmt->dest->ret_type.ptr_removed()));
+      TI_WARN("\n{}", stmt->tb);
       stmt->val = insert_type_cast_before(stmt, stmt->val,
                                           stmt->dest->ret_type.ptr_removed());
     }