From f2b777ef2b78650723372abd0ed3b264c41505f0 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 2 Nov 2016 15:55:38 -0500 Subject: [PATCH] Array: implement bitwise integer binary ops. This adds support for &, |, and ^. --- pyopencl/array.py | 115 ++++++++++++++++++++++++++++++++++++++++ pyopencl/elementwise.py | 22 ++++++++ test/test_array.py | 68 ++++++++++++++++++++++++ 3 files changed, 205 insertions(+) diff --git a/pyopencl/array.py b/pyopencl/array.py index 751de6f3..fe96f91e 100644 --- a/pyopencl/array.py +++ b/pyopencl/array.py @@ -392,6 +392,13 @@ class Array(object): .. automethod :: __rdiv__ .. automethod :: __pow__ + .. automethod :: __and__ + .. automethod :: __xor__ + .. automethod :: __or__ + .. automethod :: __iand__ + .. automethod :: __ixor__ + .. automethod :: __ior__ + .. automethod :: __abs__ .. UNDOC reverse() @@ -863,6 +870,21 @@ class Array(object): return self.__class__(self.context, self.shape, dtype, strides=strides, allocator=self.allocator) + @staticmethod + @elwise_kernel_runner + def _scalar_binop(out, a, b, queue=None, op=None): + return elementwise.get_array_scalar_binop_kernel( + out.context, op, out.dtype, a.dtype, + np.array(b).dtype) + + @staticmethod + @elwise_kernel_runner + def _array_binop(out, a, b, queue=None, op=None): + if a.shape != b.shape: + raise ValueError("shapes of binop arguments do not match") + return elementwise.get_array_binop_kernel( + out.context, op, out.dtype, a.dtype, b.dtype) + # }}} # {{{ operators @@ -1042,6 +1064,99 @@ class Array(object): __rtruediv__ = __rdiv__ + def __and__(self, other): + common_dtype = _get_common_dtype(self, other, self.queue) + + if not np.issubdtype(common_dtype, np.integer): + raise TypeError("Integral types only") + + if isinstance(other, Array): + result = self._new_like_me(common_dtype) + result.add_event(self._array_binop(result, self, other, op="&")) + else: + # create a new array for the result + result = self._new_like_me(common_dtype) + result.add_event( + self._scalar_binop(result, self, other, op="&")) + + return result + + __rand__ = __and__ # commutes + + def __or__(self, other): + common_dtype = _get_common_dtype(self, other, self.queue) + + if not np.issubdtype(common_dtype, np.integer): + raise TypeError("Integral types only") + + if isinstance(other, Array): + result = self._new_like_me(common_dtype) + result.add_event(self._array_binop(result, self, other, op="|")) + else: + # create a new array for the result + result = self._new_like_me(common_dtype) + result.add_event( + self._scalar_binop(result, self, other, op="|")) + + return result + + __ror__ = __or__ # commutes + + def __xor__(self, other): + common_dtype = _get_common_dtype(self, other, self.queue) + + if not np.issubdtype(common_dtype, np.integer): + raise TypeError("Integral types only") + + if isinstance(other, Array): + result = self._new_like_me(common_dtype) + result.add_event(self._array_binop(result, self, other, op="^")) + else: + # create a new array for the result + result = self._new_like_me(common_dtype) + result.add_event( + self._scalar_binop(result, self, other, op="^")) + + return result + + __rxor__ = __xor__ # commutes + + def __iand__(self, other): + common_dtype = _get_common_dtype(self, other, self.queue) + + if not np.issubdtype(common_dtype, np.integer): + raise TypeError("Integral types only") + + if isinstance(other, Array): + self.add_event(self._array_binop(self, self, other, op="&")) + else: + self.add_event( + self._scalar_binop(self, self, other, op="&")) + + def __ior__(self, other): + common_dtype = _get_common_dtype(self, other, self.queue) + + if not np.issubdtype(common_dtype, np.integer): + raise TypeError("Integral types only") + + if isinstance(other, Array): + self.add_event(self._array_binop(self, self, other, op="|")) + else: + self.add_event( + self._scalar_binop(self, self, other, op="|")) + + def __ixor__(self, other): + common_dtype = _get_common_dtype(self, other, self.queue) + + if not np.issubdtype(common_dtype, np.integer): + raise TypeError("Integral types only") + + if isinstance(other, Array): + self.add_event(self._array_binop(self, self, other, op="^")) + else: + self.add_event( + self._scalar_binop(self, self, other, op="^")) + def _zero_fill(self, queue=None, wait_for=None): queue = queue or self.queue diff --git a/pyopencl/elementwise.py b/pyopencl/elementwise.py index 9a6d794e..b373d2fc 100644 --- a/pyopencl/elementwise.py +++ b/pyopencl/elementwise.py @@ -823,6 +823,28 @@ def get_pow_kernel(context, dtype_x, dtype_y, dtype_z, name="pow_method") +@context_dependent_memoize +def get_array_scalar_binop_kernel(context, operator, dtype_res, dtype_a, dtype_b): + return get_elwise_kernel(context, [ + VectorArg(dtype_res, "out", with_offset=True), + VectorArg(dtype_a, "a", with_offset=True), + ScalarArg(dtype_b, "b"), + ], + "out[i] = a[i] %s b" % operator, + name="scalar_binop_kernel") + + +@context_dependent_memoize +def get_array_binop_kernel(context, operator, dtype_res, dtype_a, dtype_b): + return get_elwise_kernel(context, [ + VectorArg(dtype_res, "out", with_offset=True), + VectorArg(dtype_a, "a", with_offset=True), + VectorArg(dtype_b, "b", with_offset=True), + ], + "out[i] = a[i] %s b[i]" % operator, + name="binop_kernel") + + @context_dependent_memoize def get_array_scalar_comparison_kernel(context, operator, dtype_a): return get_elwise_kernel(context, [ diff --git a/test/test_array.py b/test/test_array.py index 71bd1a90..fd9a76e5 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -468,6 +468,74 @@ def test_divide_array(ctx_factory): a_divide = (b_gpu / a_gpu).get() assert (np.abs(b / a - a_divide) < 1e-3).all() + +def test_bitwise(ctx_factory): + context = ctx_factory() + queue = cl.CommandQueue(context) + + from itertools import product + + dtypes = [np.dtype(t) for t in (np.int64, np.int32, np.int16, np.int8)] + + for a_dtype, b_dtype in product(dtypes, dtypes): + l = 16 + + np.random.seed(10) + + a_dev = cl.array.to_device( + queue, + np.random.randint( + low=np.iinfo(a_dtype).min, + high=1+np.iinfo(a_dtype).max, + size=(l,), + dtype=a_dtype)) + b_dev = cl.array.to_device( + queue, + np.random.randint( + low=np.iinfo(b_dtype).min, + high=1+np.iinfo(b_dtype).max, + size=(l,), + dtype=b_dtype)) + + a = a_dev.get() + b = b_dev.get() + s = np.random.randint( + low=np.iinfo(b_dtype).min, + high=1+np.iinfo(b_dtype).max) + + import operator as o + + for op in [o.and_, o.or_, o.xor]: + res_dev = op(a_dev, b_dev) + res = op(a, b) + + assert (res_dev.get() == res).all() + + res_dev = op(a_dev, s) + res = op(a, s) + + assert (res_dev.get() == res).all() + + res_dev = op(s, b_dev) + res = op(s, b) + + assert (res_dev.get() == res).all() + + for op in [o.iand, o.ior, o.ixor]: + res_dev = a_dev.copy() + op(res_dev, b_dev) + res = a.copy() + op(res, b) + + assert (res_dev.get() == res).all() + + res_dev = a_dev.copy() + op(res_dev, s) + res = a.copy() + op(res, s) + + assert (res_dev.get() == res).all() + # }}} -- GitLab