[LANG] Support for Bitwise Operation (#502)

* [LANG] Support for Bitwise Operation

* Add test

Author: ZihengJiang

python/tvm/expr.py

@@ -62,6 +62,24 @@ def __mod__(self, other):
     def __neg__(self):
         return self.__mul__(-1)
 
+    def __lshift__(self, other):
+        return _make.Call(self.dtype, "shift_left", [self, other], Call.PureIntrinsic, None, 0)
+
+    def __rshift__(self, other):
+        return _make.Call(self.dtype, "shift_right", [self, other], Call.PureIntrinsic, None, 0)
+
+    def __and__(self, other):
+        return _make.Call(self.dtype, "bitwise_and", [self, other], Call.PureIntrinsic, None, 0)
+
+    def __or__(self, other):
+        return _make.Call(self.dtype, "bitwise_or", [self, other], Call.PureIntrinsic, None, 0)
+
+    def __xor__(self, other):
+        return _make.Call(self.dtype, "bitwise_xor", [self, other], Call.PureIntrinsic, None, 0)
+
+    def __invert__(self):
+        return _make.Call(self.dtype, "bitwise_not", [self], Call.PureIntrinsic, None, 0)
+
     def __lt__(self, other):
         return _make.LT(self, other)
 

tests/python/unittest/test_lang_basic.py

@@ -123,6 +123,16 @@ def test_all():
         '(((%s < %s) && (%s > (%s + 1))) && (%s < (%s*2)))' % (
             x.name, y.name, y.name, z.name, x.name, z.name)
 
+def test_bitwise():
+    x = tvm.var('x')
+    y = tvm.var('y')
+    assert str(x << y) == 'shift_left(x, y)'
+    assert str(x >> y) == 'shift_right(x, y)'
+    assert str(x & y) == 'bitwise_and(x, y)'
+    assert str(x | y) == 'bitwise_or(x, y)'
+    assert str(x ^ y) == 'bitwise_xor(x, y)'
+    assert str(~x) == 'bitwise_not(x)'
+
 
 if __name__ == "__main__":
     test_cast()
@@ -137,3 +147,4 @@ def test_all():
     test_dtype()
     test_any()
     test_all()
+    test_bitwise()