From 4fdef3adfc9baf21ad5a867b82f13cd0c24386a1 Mon Sep 17 00:00:00 2001
From: ziheng <ziheng@apache.org>
Date: Sun, 1 Oct 2017 10:32:02 -0700
Subject: [PATCH] [LANG] Support for Bitwise Operation (#502)

* [LANG] Support for Bitwise Operation

* Add test
---
 python/tvm/expr.py                       | 18 ++++++++++++++++++
 tests/python/unittest/test_lang_basic.py | 11 +++++++++++
 2 files changed, 29 insertions(+)

diff --git a/python/tvm/expr.py b/python/tvm/expr.py
index b1d5b51b47a8..361f40a2fc35 100644
--- a/python/tvm/expr.py
+++ b/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)
 
diff --git a/tests/python/unittest/test_lang_basic.py b/tests/python/unittest/test_lang_basic.py
index 4b3706bda7a3..dd982313daba 100644
--- a/tests/python/unittest/test_lang_basic.py
+++ b/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()