[LANG] Support for Bitwise Operation (#502)

* [LANG] Support for Bitwise Operation * Add test
apache · Oct 1, 2017 · 4fdef3a · 4fdef3a
1 parent 9c2fc09
commit 4fdef3a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git 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
@@ -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()