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()