ENH: add bitwise_invert

.github/workflows/array_api.yml

@@ -49,4 +49,4 @@ jobs:
           # for hypothesis-driven test case generation
           pytest $GITHUB_WORKSPACE/pre_compile_tools/pre_compile_ufuncs.py -s
           # only run a subset of the conformance tests to get started
-          pytest array_api_tests/meta/test_broadcasting.py array_api_tests/meta/test_equality_mapping.py array_api_tests/meta/test_signatures.py array_api_tests/meta/test_special_cases.py array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor
+          pytest array_api_tests/meta/test_broadcasting.py array_api_tests/meta/test_equality_mapping.py array_api_tests/meta/test_signatures.py array_api_tests/meta/test_special_cases.py array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_invert

pykokkos/__init__.py

@@ -61,7 +61,8 @@
                                  round,
                                  trunc,
                                  ceil,
-                                 floor)
+                                 floor,
+                                 bitwise_invert)
 from pykokkos.lib.info import iinfo, finfo
 from pykokkos.lib.create import (zeros,
                                  ones,

pykokkos/interface/views.py

@@ -392,6 +392,12 @@ def __array__(self, dtype=None):
         return self.data
 
 
+    def __invert__(self):
+        # scope to avoid circular import issues
+        from pykokkos.lib.ufuncs import bitwise_invert
+        return bitwise_invert(self)
+
+
     @staticmethod
     def _get_dtype_name(type_name: str) -> str:
         """

pykokkos/lib/ufunc_workunits.py

@@ -1,6 +1,115 @@
 import pykokkos as pk
 
 
+@pk.workunit
+def bitwise_invert_impl_1d_bool(tid: int, view: pk.View1D[pk.uint8], out: pk.View1D[pk.uint8]):
+    # need special handling until we have genuine
+    # bool support I think
+    if view[tid] == 0:
+        out[tid] = 1
+    else:
+        out[tid] = 0
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_bool(tid: int, view: pk.View2D[pk.uint8], out: pk.View2D[pk.uint8]):
+    # need special handling until we have genuine
+    # bool support I think
+    for i in range(view.extent(1)):
+        if view[tid][i] == 0:
+            out[tid][i] = 1
+        else:
+            out[tid][i] = 0
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_int8(tid: int, view: pk.View2D[pk.int8], out: pk.View2D[pk.int8]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_int16(tid: int, view: pk.View2D[pk.int16], out: pk.View2D[pk.int16]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_int32(tid: int, view: pk.View2D[pk.int32], out: pk.View2D[pk.int32]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_int64(tid: int, view: pk.View2D[pk.int64], out: pk.View2D[pk.int64]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_uint8(tid: int, view: pk.View2D[pk.uint8], out: pk.View2D[pk.uint8]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_uint16(tid: int, view: pk.View2D[pk.uint16], out: pk.View2D[pk.uint16]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_uint32(tid: int, view: pk.View2D[pk.uint32], out: pk.View2D[pk.uint32]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_2d_uint64(tid: int, view: pk.View2D[pk.uint64], out: pk.View2D[pk.uint64]):
+    for i in range(view.extent(1)):
+        out[tid][i] = ~view[tid][i]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_uint8(tid: int, view: pk.View1D[pk.uint8], out: pk.View1D[pk.uint8]):
+    out[tid] = ~view[tid]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_uint16(tid: int, view: pk.View1D[pk.uint16], out: pk.View1D[pk.uint16]):
+    out[tid] = ~view[tid]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_uint32(tid: int, view: pk.View1D[pk.uint32], out: pk.View1D[pk.uint32]):
+    out[tid] = ~view[tid]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_uint64(tid: int, view: pk.View1D[pk.uint64], out: pk.View1D[pk.uint64]):
+    out[tid] = ~view[tid]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_int8(tid: int, view: pk.View1D[pk.int8], out: pk.View1D[pk.int8]):
+    out[tid] = ~view[tid]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_int16(tid: int, view: pk.View1D[pk.int16], out: pk.View1D[pk.int16]):
+    out[tid] = ~view[tid]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_int32(tid: int, view: pk.View1D[pk.int32], out: pk.View1D[pk.int32]):
+    out[tid] = ~view[tid]
+
+
+@pk.workunit
+def bitwise_invert_impl_1d_int64(tid: int, view: pk.View1D[pk.int64], out: pk.View1D[pk.int64]):
+    out[tid] = ~view[tid]
+
+
 @pk.workunit
 def floor_impl_1d_double(tid: int, view: pk.View1D[pk.double], out: pk.View1D[pk.double]):
     out[tid] = floor(view[tid])

pykokkos/lib/ufuncs.py

@@ -2679,3 +2679,41 @@ def floor(view):
                              out=out,
                              view=view)
     return out
+
+
+def bitwise_invert(view):
+    """
+    Inverts (flips) each bit for each element of the input view.
+
+    Parameters
+    ----------
+    view : pykokkos view
+           Should have an integer or boolean data type.
+
+    Returns
+    -------
+    out: pykokkos view
+         A view containing the element-wise results. The returned view
+         must have the same type as the input view.
+    """
+    dtype = view.dtype
+    ndims = len(view.shape)
+    dtype_str = str(dtype)
+    out = pk.View(view.shape, dtype=dtype)
+    if ndims > 2:
+        raise NotImplementedError("only up to 2D views currently supported for bitwise_invert() ufunc.")
+
+    if view.shape == ():
+        tid = 1
+    else:
+        tid = view.shape[0]
+    if view.size == 0:
+        return view
+    _ufunc_kernel_dispatcher(tid=tid,
+                             dtype=dtype,
+                             ndims=ndims,
+                             op="bitwise_invert",
+                             sub_dispatcher=pk.parallel_for,
+                             out=out,
+                             view=view)
+    return out