diff --git a/.github/workflows/array_api.yml b/.github/workflows/array_api.yml index 2d656c21..ba53c971 100644 --- a/.github/workflows/array_api.yml +++ b/.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 diff --git a/pykokkos/__init__.py b/pykokkos/__init__.py index 3b83a079..8753544b 100644 --- a/pykokkos/__init__.py +++ b/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, diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 7b169874..033c02a9 100644 --- a/pykokkos/interface/views.py +++ b/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: """ diff --git a/pykokkos/lib/ufunc_workunits.py b/pykokkos/lib/ufunc_workunits.py index 06f12530..3634f08e 100644 --- a/pykokkos/lib/ufunc_workunits.py +++ b/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]) diff --git a/pykokkos/lib/ufuncs.py b/pykokkos/lib/ufuncs.py index 13be86e2..9dfdec8c 100644 --- a/pykokkos/lib/ufuncs.py +++ b/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