From ebded7072cec50225aaff856ff8fddd90d648687 Mon Sep 17 00:00:00 2001 From: ZihengJiang Date: Sat, 30 Sep 2017 23:57:59 -0700 Subject: [PATCH 1/2] [LANG] Support for Bitwise Operation --- python/tvm/expr.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/python/tvm/expr.py b/python/tvm/expr.py index b1d5b51b47a8..c5b336f78bca 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, other): + return _make.Call(self.dtype, "bitwise_not", [self, other], Call.PureIntrinsic, None, 0) + def __lt__(self, other): return _make.LT(self, other) From c645507707a000bb766b125ea6c521a01ca3e718 Mon Sep 17 00:00:00 2001 From: ZihengJiang Date: Sun, 1 Oct 2017 00:06:26 -0700 Subject: [PATCH 2/2] Add test --- python/tvm/expr.py | 4 ++-- tests/python/unittest/test_lang_basic.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/python/tvm/expr.py b/python/tvm/expr.py index c5b336f78bca..361f40a2fc35 100644 --- a/python/tvm/expr.py +++ b/python/tvm/expr.py @@ -77,8 +77,8 @@ def __or__(self, other): def __xor__(self, other): return _make.Call(self.dtype, "bitwise_xor", [self, other], Call.PureIntrinsic, None, 0) - def __invert__(self, other): - return _make.Call(self.dtype, "bitwise_not", [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()