diff --git a/CHANGELOG.md b/CHANGELOG.md index 734643223c8a..7141063583c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Added support for boolean input in `dpnp.positive` and the `+` operator of `dpnp.ndarray` [#3128](https://github.com/IntelPython/dpnp/pull/3128) * Added `dpnp.top_k` to return the `k` largest or smallest elements and their indices along an axis, with a `sorted` keyword to return them in sorted order [#3133](https://github.com/IntelPython/dpnp/pull/3133) +* Added `dpnp.argpartition` and `dpnp.ndarray.argpartition` to return the indices that would partition an array [#3138](https://github.com/IntelPython/dpnp/pull/3138) ### Changed @@ -18,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Merged the `ufunc` and `indexing` extensions into `dpnp.libtensor._tensor_elementwise_impl` and `dpnp.libtensor._tensor_impl`, while `statistics` and `window` extensions became part of libtensor [#3131](https://github.com/IntelPython/dpnp/pull/3131) * Moved `blas`, `fft`, `lapack`, `sparse` and `vm` extensions into `dpnp/libtensor/extensions` [#3131](https://github.com/IntelPython/dpnp/pull/3131) * Moved the libtensor headers shipped with dpnp from `dpnp/tensor/libtensor/include` to `dpnp/libtensor/include`, and changed `dpnp.get_include` to return that directory [#3131](https://github.com/IntelPython/dpnp/pull/3131) +* Changed `dpnp.partition`, `dpnp.median` and `dpnp.nanmedian` to use a radix select instead of a full sort for all but complex data types [#3138](https://github.com/IntelPython/dpnp/pull/3138) ### Deprecated @@ -29,12 +31,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Removed * Removed the unused `FourZeroOffsets_Indexer`, `FixedDimStridedIndexer` and `ThreeOffsets_FixedDimStridedIndexer` from the libtensor header `utils/offset_utils.hpp` [#3131](https://github.com/IntelPython/dpnp/pull/3131) +* Removed the legacy implementation of `dpnp.partition` from the backend [#3138](https://github.com/IntelPython/dpnp/pull/3138) ### Fixed * Fixed `flags[key]` of `dpnp.ndarray` and `dpnp.tensor.usm_ndarray` to raise `KeyError` for an unknown key [#3141](https://github.com/IntelPython/dpnp/pull/3141) * Released the GIL around the oneMKL VM calls in the VM extension, which deadlocked on an in-order queue of a CPU device [#3156](https://github.com/IntelPython/dpnp/pull/3156) * Fixed `dpnp.place` and `dpnp.tensor.place` to order the write after the cast of `vals` to the array's data type, which could put wrong values into the array on an out-of-order queue [#3164](https://github.com/IntelPython/dpnp/pull/3164) +* Fixed `dpnp.median` and `dpnp.nanmedian` raising `ValueError` for an `out` array of the result shape with `keepdims=True` [#3138](https://github.com/IntelPython/dpnp/pull/3138) ### Security diff --git a/dpnp/__init__.py b/dpnp/__init__.py index ce3b196e070d..e5cf18c13383 100644 --- a/dpnp/__init__.py +++ b/dpnp/__init__.py @@ -511,6 +511,7 @@ where, ) from .dpnp_iface_sorting import ( + argpartition, argsort, partition, sort, @@ -988,6 +989,7 @@ __all__ += [ "argmax", "argmin", + "argpartition", "argwhere", "argsort", "count_nonzero", diff --git a/dpnp/backend/CMakeLists.txt b/dpnp/backend/CMakeLists.txt index 7a474b29170c..dfb2109508f0 100644 --- a/dpnp/backend/CMakeLists.txt +++ b/dpnp/backend/CMakeLists.txt @@ -31,7 +31,6 @@ set(DPNP_SRC kernels/dpnp_krnl_arraycreation.cpp kernels/dpnp_krnl_common.cpp kernels/dpnp_krnl_random.cpp - kernels/dpnp_krnl_sorting.cpp src/dpnp_iface_fptr.cpp src/memory_sycl.cpp src/queue_sycl.cpp diff --git a/dpnp/backend/include/dpnp_iface.hpp b/dpnp/backend/include/dpnp_iface.hpp index b7627101aab4..77e9eae0d922 100644 --- a/dpnp/backend/include/dpnp_iface.hpp +++ b/dpnp/backend/include/dpnp_iface.hpp @@ -103,38 +103,6 @@ INP_DLLEXPORT void dpnp_memory_memcpy_c(DPCTLSyclQueueRef q_ref, INP_DLLEXPORT void dpnp_memory_memcpy_c(void *dst, const void *src, size_t size_in_bytes); -/** - * @ingroup BACKEND_API - * @brief Return a partitioned copy of an array. - * - * @param [in] q_ref Reference to SYCL queue. - * @param [in] array Input array. - * @param [in] array2 Copy input array. - * @param [out] result Result array. - * @param [in] kth Element index to partition by. - * @param [in] shape Shape of input array. - * @param [in] ndim Number of elements in shape. - * @param [in] dep_event_vec_ref Reference to vector of SYCL events. - */ -template -INP_DLLEXPORT DPCTLSyclEventRef - dpnp_partition_c(DPCTLSyclQueueRef q_ref, - void *array, - void *array2, - void *result, - const size_t kth, - const shape_elem_type *shape, - const size_t ndim, - const DPCTLEventVectorRef dep_event_vec_ref); - -template -INP_DLLEXPORT void dpnp_partition_c(void *array, - void *array2, - void *result, - const size_t kth, - const shape_elem_type *shape, - const size_t ndim); - /** * @ingroup BACKEND_API * @brief implementation of creating filled with value array function diff --git a/dpnp/backend/include/dpnp_iface_fptr.hpp b/dpnp/backend/include/dpnp_iface_fptr.hpp index 868a31bdaac5..7311f3518d9e 100644 --- a/dpnp/backend/include/dpnp_iface_fptr.hpp +++ b/dpnp/backend/include/dpnp_iface_fptr.hpp @@ -64,17 +64,14 @@ enum class DPNPFuncName : size_t DPNP_FN_NONE, /**< Very first element of the enumeration */ DPNP_FN_INITVAL, /**< Used in numpy ones, ones_like, zeros, zeros_like impls */ - DPNP_FN_INITVAL_EXT, /**< Used in numpy ones, ones_like, zeros, zeros_like - impls */ - DPNP_FN_ONES, /**< Used in numpy.ones() impl */ - DPNP_FN_ONES_LIKE, /**< Used in numpy.ones_like() impl */ - DPNP_FN_PARTITION, /**< Used in numpy.partition() impl */ - DPNP_FN_PARTITION_EXT, /**< Used in numpy.partition() impl, requires extra - parameters */ - DPNP_FN_RNG_BETA, /**< Used in numpy.random.beta() impl */ - DPNP_FN_RNG_BETA_EXT, /**< Used in numpy.random.beta() impl, requires extra - parameters */ - DPNP_FN_RNG_BINOMIAL, /**< Used in numpy.random.binomial() impl */ + DPNP_FN_INITVAL_EXT, /**< Used in numpy ones, ones_like, zeros, zeros_like + impls */ + DPNP_FN_ONES, /**< Used in numpy.ones() impl */ + DPNP_FN_ONES_LIKE, /**< Used in numpy.ones_like() impl */ + DPNP_FN_RNG_BETA, /**< Used in numpy.random.beta() impl */ + DPNP_FN_RNG_BETA_EXT, /**< Used in numpy.random.beta() impl, requires extra + parameters */ + DPNP_FN_RNG_BINOMIAL, /**< Used in numpy.random.binomial() impl */ DPNP_FN_RNG_BINOMIAL_EXT, /**< Used in numpy.random.binomial() impl, requires extra parameters */ DPNP_FN_RNG_CHISQUARE, /**< Used in numpy.random.chisquare() impl */ diff --git a/dpnp/backend/kernels/dpnp_krnl_sorting.cpp b/dpnp/backend/kernels/dpnp_krnl_sorting.cpp deleted file mode 100644 index a1495dfb0279..000000000000 --- a/dpnp/backend/kernels/dpnp_krnl_sorting.cpp +++ /dev/null @@ -1,160 +0,0 @@ -//***************************************************************************** -// Copyright (c) 2016, Intel Corporation -// All rights reserved. -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// - Redistributions of source code must retain the above copyright notice, -// this list of conditions and the following disclaimer. -// - Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// - Neither the name of the copyright holder nor the names of its contributors -// may be used to endorse or promote products derived from this software -// without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF -// THE POSSIBILITY OF SUCH DAMAGE. -//***************************************************************************** - -#include - -#include "dpnp_fptr.hpp" -#include "dpnpc_memory_adapter.hpp" -#include "queue_sycl.hpp" -#include - -template -class dpnp_partition_c_kernel; - -template -DPCTLSyclEventRef dpnp_partition_c(DPCTLSyclQueueRef q_ref, - void *array1_in, - void *array2_in, - void *result1, - const size_t kth, - const shape_elem_type *shape_, - const size_t ndim, - const DPCTLEventVectorRef dep_event_vec_ref) -{ - // avoid warning unused variable - (void)dep_event_vec_ref; - - DPCTLSyclEventRef event_ref = nullptr; - - if ((array1_in == nullptr) || (array2_in == nullptr) || - (result1 == nullptr)) { - return event_ref; - } - - if (ndim < 1) { - return event_ref; - } - - const size_t size = std::accumulate(shape_, shape_ + ndim, 1, - std::multiplies()); - size_t size_ = size / shape_[ndim - 1]; - - if (size_ == 0) { - return event_ref; - } - - sycl::queue q = *(reinterpret_cast(q_ref)); - - _DataType *arr = static_cast<_DataType *>(array1_in); - _DataType *result = static_cast<_DataType *>(result1); - - auto policy = oneapi::dpl::execution::make_device_policy< - dpnp_partition_c_kernel<_DataType>>(q); - - // fill the result array with data from input one - q.memcpy(result, arr, size * sizeof(_DataType)).wait(); - - // note, a loop for a multidemension input array (size_ > 1) is an - // experimental and it isn't tested properly as for now - for (size_t i = 0; i < size_; i++) { - _DataType *bufptr = result + i * shape_[0]; - - // for every slice it makes a partial sorting such that: - // 1. result[0 <= i < kth] <= result[kth] - // 2. result[kth <= i < size] >= result[kth] - // event-blocking call, no need for wait() - std::nth_element(policy, bufptr, bufptr + kth, bufptr + size, - dpnp_less_comp()); - } - return event_ref; -} - -template -void dpnp_partition_c(void *array1_in, - void *array2_in, - void *result1, - const size_t kth, - const shape_elem_type *shape_, - const size_t ndim) -{ - DPCTLSyclQueueRef q_ref = reinterpret_cast(&DPNP_QUEUE); - DPCTLEventVectorRef dep_event_vec_ref = nullptr; - DPCTLSyclEventRef event_ref = - dpnp_partition_c<_DataType>(q_ref, array1_in, array2_in, result1, kth, - shape_, ndim, dep_event_vec_ref); - DPCTLEvent_WaitAndThrow(event_ref); - DPCTLEvent_Delete(event_ref); -} - -template -void (*dpnp_partition_default_c)(void *, - void *, - void *, - const size_t, - const shape_elem_type *, - const size_t) = dpnp_partition_c<_DataType>; - -template -DPCTLSyclEventRef (*dpnp_partition_ext_c)(DPCTLSyclQueueRef, - void *, - void *, - void *, - const size_t, - const shape_elem_type *, - const size_t, - const DPCTLEventVectorRef) = - dpnp_partition_c<_DataType>; - -void func_map_init_sorting(func_map_t &fmap) -{ - fmap[DPNPFuncName::DPNP_FN_PARTITION][eft_INT][eft_INT] = { - eft_INT, (void *)dpnp_partition_default_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION][eft_LNG][eft_LNG] = { - eft_LNG, (void *)dpnp_partition_default_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION][eft_FLT][eft_FLT] = { - eft_FLT, (void *)dpnp_partition_default_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION][eft_DBL][eft_DBL] = { - eft_DBL, (void *)dpnp_partition_default_c}; - - fmap[DPNPFuncName::DPNP_FN_PARTITION_EXT][eft_BLN][eft_BLN] = { - eft_BLN, (void *)dpnp_partition_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION_EXT][eft_INT][eft_INT] = { - eft_INT, (void *)dpnp_partition_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION_EXT][eft_LNG][eft_LNG] = { - eft_LNG, (void *)dpnp_partition_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION_EXT][eft_FLT][eft_FLT] = { - eft_FLT, (void *)dpnp_partition_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION_EXT][eft_DBL][eft_DBL] = { - eft_DBL, (void *)dpnp_partition_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PARTITION_EXT][eft_C64][eft_C64] = { - eft_C64, (void *)dpnp_partition_ext_c>}; - fmap[DPNPFuncName::DPNP_FN_PARTITION_EXT][eft_C128][eft_C128] = { - eft_C128, (void *)dpnp_partition_ext_c>}; - - return; -} diff --git a/dpnp/backend/src/dpnp_fptr.hpp b/dpnp/backend/src/dpnp_fptr.hpp index 15d6b7ab20ee..5e2d79b09c2d 100644 --- a/dpnp/backend/src/dpnp_fptr.hpp +++ b/dpnp/backend/src/dpnp_fptr.hpp @@ -161,6 +161,5 @@ class dpnp_less_comp void func_map_init_arraycreation(func_map_t &fmap); void func_map_init_linalg(func_map_t &fmap); void func_map_init_random(func_map_t &fmap); -void func_map_init_sorting(func_map_t &fmap); #endif // BACKEND_FPTR_H diff --git a/dpnp/backend/src/dpnp_iface_fptr.cpp b/dpnp/backend/src/dpnp_iface_fptr.cpp index ab1769c7e56c..7e319fb07246 100644 --- a/dpnp/backend/src/dpnp_iface_fptr.cpp +++ b/dpnp/backend/src/dpnp_iface_fptr.cpp @@ -101,7 +101,6 @@ static func_map_t func_map_init() func_map_init_arraycreation(fmap); func_map_init_linalg(fmap); func_map_init_random(fmap); - func_map_init_sorting(fmap); return fmap; }; diff --git a/dpnp/dpnp_algo/CMakeLists.txt b/dpnp/dpnp_algo/CMakeLists.txt index 3f8196bc64f8..2ddf3f10eb3e 100644 --- a/dpnp/dpnp_algo/CMakeLists.txt +++ b/dpnp/dpnp_algo/CMakeLists.txt @@ -27,13 +27,8 @@ # THE POSSIBILITY OF SUCH DAMAGE. # ***************************************************************************** -set(dpnp_algo_pyx_deps ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_sorting.pxi) - build_dpnp_cython_ext_with_backend( dpnp_algo ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo.pyx dpnp/dpnp_algo ) - -add_custom_target(_dpnp_algo_deps DEPENDS ${dpnp_algo_pyx_deps}) -add_dependencies(dpnp_algo _dpnp_algo_deps) diff --git a/dpnp/dpnp_algo/dpnp_algo.pxd b/dpnp/dpnp_algo/dpnp_algo.pxd index 2ac2a0a04531..2d1766e97bfb 100644 --- a/dpnp/dpnp_algo/dpnp_algo.pxd +++ b/dpnp/dpnp_algo/dpnp_algo.pxd @@ -32,7 +32,6 @@ # need this namespace for Enum import cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName": cdef enum DPNPFuncName "DPNPFuncName": - DPNP_FN_PARTITION_EXT DPNP_FN_RNG_BETA_EXT DPNP_FN_RNG_BINOMIAL_EXT DPNP_FN_RNG_CHISQUARE_EXT diff --git a/dpnp/dpnp_algo/dpnp_algo.pyx b/dpnp/dpnp_algo/dpnp_algo.pyx index 9d52599b73b6..effd9395686a 100644 --- a/dpnp/dpnp_algo/dpnp_algo.pyx +++ b/dpnp/dpnp_algo/dpnp_algo.pyx @@ -40,18 +40,12 @@ and the rest of the library import dpnp -cimport dpctl as c_dpctl - cimport dpnp.dpnp_utils as utils -from dpnp.dpnp_algo cimport shape_elem_type, shape_type_c __all__ = [ ] -include "dpnp_algo_sorting.pxi" - - """ Internal functions """ diff --git a/dpnp/dpnp_algo/dpnp_algo_sorting.pxi b/dpnp/dpnp_algo/dpnp_algo_sorting.pxi deleted file mode 100644 index 516c080a924a..000000000000 --- a/dpnp/dpnp_algo/dpnp_algo_sorting.pxi +++ /dev/null @@ -1,127 +0,0 @@ -# cython: language_level=3 -# cython: linetrace=True -# -*- coding: utf-8 -*- -# ***************************************************************************** -# Copyright (c) 2016, Intel Corporation -# All rights reserved. -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# - Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# - Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# - Neither the name of the copyright holder nor the names of its contributors -# may be used to endorse or promote products derived from this software -# without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF -# THE POSSIBILITY OF SUCH DAMAGE. -# ***************************************************************************** - -"""Module Backend (Sorting part) - -This module contains interface functions between C backend layer -and the rest of the library - -""" - -# NO IMPORTs here. All imports must be placed into main "dpnp_algo.pyx" file - -__all__ += [ - "dpnp_partition", -] - - -ctypedef c_dpctl.DPCTLSyclEventRef( - *fptr_dpnp_partition_t)( - c_dpctl.DPCTLSyclQueueRef, - void *, - void *, - void *, - const size_t, - const shape_elem_type *, - const size_t, - const c_dpctl.DPCTLEventVectorRef, -) - - -cpdef utils.dpnp_descriptor dpnp_partition( - utils.dpnp_descriptor arr, int kth, - axis=-1, kind="introselect", order=None, -): - cdef shape_type_c shape1 = arr.shape - - cdef size_t kth_ = ( - kth if kth >= 0 else (arr.ndim + kth) - ) - cdef DPNPFuncType param1_type = ( - dpnp_dtype_to_DPNPFuncType(arr.dtype) - ) - - cdef DPNPFuncData kernel_data = ( - get_dpnp_function_ptr( - DPNP_FN_PARTITION_EXT, - param1_type, param1_type, - ) - ) - - cdef utils.dpnp_descriptor arr2 = ( - dpnp.get_dpnp_descriptor( - arr.get_pyobj().copy(), - copy_when_nondefault_queue=False, - ) - ) - - arr_obj = arr.get_array() - - cdef utils.dpnp_descriptor result = ( - utils.create_output_descriptor( - arr.shape, - kernel_data.return_type, - None, - device=arr_obj.sycl_device, - usm_type=arr_obj.usm_type, - sycl_queue=arr_obj.sycl_queue, - ) - ) - - result_sycl_queue = result.get_array().sycl_queue - - cdef c_dpctl.SyclQueue q = ( - result_sycl_queue - ) - cdef c_dpctl.DPCTLSyclQueueRef q_ref = ( - q.get_queue_ref() - ) - - cdef fptr_dpnp_partition_t func = ( - kernel_data.ptr - ) - - cdef c_dpctl.DPCTLSyclEventRef event_ref = func( - q_ref, - arr.get_data(), - arr2.get_data(), - result.get_data(), - kth_, - shape1.data(), - arr.ndim, - NULL, - ) - - with nogil: - c_dpctl.DPCTLEvent_WaitAndThrow(event_ref) - c_dpctl.DPCTLEvent_Delete(event_ref) - - return result diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py index c7f9fe035943..bc6c9843862e 100644 --- a/dpnp/dpnp_array.py +++ b/dpnp/dpnp_array.py @@ -891,7 +891,19 @@ def argmin(self, /, axis=None, out=None, *, keepdims=False): return dpnp.argmin(self, axis=axis, out=out, keepdims=keepdims) - # 'argpartition', + def argpartition(self, /, kth, axis=-1, kind="introselect", order=None): + """ + Return indices that would partition this array. + + Refer to :obj:`dpnp.argpartition` for full documentation. + + See Also + -------- + :obj:`dpnp.argpartition` : Equivalent function. + + """ + + return dpnp.argpartition(self, kth, axis=axis, kind=kind, order=order) def argsort( self, axis=-1, kind=None, order=None, *, descending=False, stable=None diff --git a/dpnp/dpnp_iface_nanfunctions.py b/dpnp/dpnp_iface_nanfunctions.py index f634064b0161..144e42ac10f0 100644 --- a/dpnp/dpnp_iface_nanfunctions.py +++ b/dpnp/dpnp_iface_nanfunctions.py @@ -736,8 +736,10 @@ def nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=False): """ dpnp.check_supported_arrays_type(a) - ignore_nan = False - if dpnp.issubdtype(a.dtype, dpnp.inexact): + # selecting the medians of real data types ignores NaNs at no extra cost, + # while the sort used for complex data types only has to if there are any + ignore_nan = dpnp.issubdtype(a.dtype, dpnp.floating) + if dpnp.issubdtype(a.dtype, dpnp.complexfloating): mask = dpnp.isnan(a) if dpnp.any(mask): ignore_nan = True diff --git a/dpnp/dpnp_iface_sorting.py b/dpnp/dpnp_iface_sorting.py index 6ee1e75094ed..9032501e1bb8 100644 --- a/dpnp/dpnp_iface_sorting.py +++ b/dpnp/dpnp_iface_sorting.py @@ -39,17 +39,20 @@ """ +# pylint: disable=protected-access +# pylint: disable=no-name-in-module + import operator from collections.abc import Sequence from typing import NamedTuple +import dpctl.utils as dpu + import dpnp +import dpnp.libtensor._tensor_impl as ti +import dpnp.libtensor._tensor_sorting_impl as sorting_ext import dpnp.tensor as dpt -# pylint: disable=no-name-in-module -from .dpnp_algo import ( - dpnp_partition, -) from .dpnp_array import dpnp_array from .dpnp_utils import ( map_dtype_to_device, @@ -57,6 +60,112 @@ from .tensor._numpy_helper import normalize_axis_index +def _normalize_kth(kth, length, check_bounds=True): + """Validate `kth` and return it as a list of non-negative integers.""" + + if isinstance(kth, int): + kth = (kth,) + elif not isinstance(kth, Sequence): + raise TypeError( + f"kth must be int or sequence of ints, but got {type(kth)}" + ) + elif not all(isinstance(k, int) for k in kth): + raise TypeError("kth is a sequence, but not all elements are integers") + + kth = list(kth) + if not check_bounds: + return kth + for i, k in enumerate(kth): + if k < 0: + k += length + if not 0 <= k < length: + raise ValueError(f"kth(={kth[i]}) out of bounds {length}") + kth[i] = k + return kth + + +def _wrap_partition_argpartition(a, kth, axis, kind, order, return_indices): + """ + Partition `a` along `axis` about a single `kth` by radix select, falling + back to a full sort for several `kth` and for complex data types. + + """ + + dpnp.check_supported_arrays_type(a) + + if kind != "introselect": + raise NotImplementedError( + "`kind` keyword argument is only supported with its default value." + ) + if order is not None: + raise NotImplementedError( + "`order` keyword argument is only supported with its default value." + ) + + if axis is None: + a = dpnp.ravel(a) + axis = -1 + + nd = a.ndim + axis = normalize_axis_index(axis, nd) + # an empty array is partitioned whatever kth is + kth = _normalize_kth(kth, a.shape[axis], check_bounds=a.size > 0) + + if len(kth) == 0 and not return_indices: + # nothing is to be put in its sorted position + return dpnp.copy(a) + if ( + a.size == 0 + or len(kth) != 1 + or dpnp.issubdtype(a.dtype, dpnp.complexfloating) + ): + sort_fn = dpnp.argsort if return_indices else dpnp.sort + return sort_fn(a, axis=axis) + + usm_a = dpnp.get_usm_ndarray(a) + if axis != nd - 1: + perm = [i for i in range(nd) if i != axis] + [axis] + usm_a = dpt.permute_dims(usm_a, perm) + + exec_q = usm_a.sycl_queue + _manager = dpu.SequentialOrderManager[exec_q] + dep_evs = _manager.submitted_events + + # the rows are partitioned along the last axis of C-contiguous arrays + if not usm_a.flags.c_contiguous: + usm_tmp = dpt.empty_like(usm_a, order="C") + ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( + src=usm_a, dst=usm_tmp, sycl_queue=exec_q, depends=dep_evs + ) + _manager.add_event_pair(ht_ev, copy_ev) + usm_a = usm_tmp + dep_evs = [copy_ev] + + if return_indices: + partition_fn = sorting_ext._argpartition + res_dt = dpt.int64 + else: + partition_fn = sorting_ext._partition + res_dt = usm_a.dtype + + usm_res = dpt.empty( + usm_a.shape, + dtype=res_dt, + usm_type=usm_a.usm_type, + order="C", + sycl_queue=exec_q, + ) + ht_ev, part_ev = partition_fn( + usm_a, kth[0], usm_res, sycl_queue=exec_q, depends=dep_evs + ) + _manager.add_event_pair(ht_ev, part_ev) + + if axis != nd - 1: + inv_perm = sorted(range(nd), key=lambda d: perm[d]) + usm_res = dpt.permute_dims(usm_res, inv_perm) + return dpnp_array._create_from_usm_ndarray(usm_res) + + # pylint:disable=missing-class-docstring class TopKResult(NamedTuple): values: dpnp.ndarray @@ -101,6 +210,79 @@ def _wrap_sort_argsort( return dpnp_array._create_from_usm_ndarray(usm_res) +def argpartition(a, kth, axis=-1, kind="introselect", order=None): + """ + Perform an indirect partition along the given axis using the algorithm + specified by the `kind` keyword. It returns an array of indices of the + same shape as `a` that index data along the given axis in partitioned + order. + + For full documentation refer to :obj:`numpy.argpartition`. + + Parameters + ---------- + a : {dpnp.ndarray, usm_ndarray} + Array to sort. + kth : {int, sequence of ints} + Element index to partition by. The k-th element will be in its final + sorted position and all smaller elements will be moved before it and + all larger elements behind it. The order of all elements in the + partitions is undefined. If provided with a sequence of k-th it will + partition all of them into their sorted position at once. + axis : {None, int}, optional + Axis along which to sort. If ``None``, the array is flattened before + sorting. The default is ``-1``, which sorts along the last axis. + + Default: ``-1``. + + Returns + ------- + index_array : dpnp.ndarray + Array of indices that partition `a` along the specified axis. If `a` + is one-dimensional, ``a[index_array]`` yields a partitioned `a`. More + generally, ``dpnp.take_along_axis(a, index_array, axis=axis)`` always + yields the partitioned `a`, irrespective of dimensionality. + + Limitations + ----------- + Parameters `kind` and `order` are only supported with its default value. + Otherwise ``NotImplementedError`` exception will be raised. + + See Also + -------- + :obj:`dpnp.partition` : Describes partition algorithms used. + :obj:`dpnp.ndarray.partition` : Inplace partition. + :obj:`dpnp.argsort` : Full indirect sort. + :obj:`dpnp.take_along_axis` : Apply ``index_array`` from argpartition to + an array as if by calling partition. + + Examples + -------- + One dimensional array: + + >>> import dpnp as np + >>> x = np.array([3, 4, 2, 1]) + >>> x[np.argpartition(x, 3)] + array([2, 1, 3, 4]) # may vary + >>> x[np.argpartition(x, (1, 3))] + array([1, 2, 3, 4]) # may vary + + Multi-dimensional array: + + >>> x = np.array([[3, 4, 2], [1, 3, 1]]) + >>> index_array = np.argpartition(x, kth=1, axis=-1) + >>> # below is the same as np.partition(x, kth=1) + >>> np.take_along_axis(x, index_array, axis=-1) + array([[2, 3, 4], + [1, 1, 3]]) + + """ + + return _wrap_partition_argpartition( + a, kth, axis, kind, order, return_indices=True + ) + + def argsort( a, axis=-1, kind=None, order=None, *, descending=False, stable=None ): @@ -277,60 +459,9 @@ def partition(a, kth, axis=-1, kind="introselect", order=None): """ - dpnp.check_supported_arrays_type(a) - - if kind != "introselect": - raise NotImplementedError( - "`kind` keyword argument is only supported with its default value." - ) - if order is not None: - raise NotImplementedError( - "`order` keyword argument is only supported with its default value." - ) - - if axis is None: - a = dpnp.ravel(a) - axis = -1 - - nd = a.ndim - axis = normalize_axis_index(axis, nd) - length = a.shape[axis] - - if isinstance(kth, int): - kth = (kth,) - elif not isinstance(kth, Sequence): - raise TypeError( - f"kth must be int or sequence of ints, but got {type(kth)}" - ) - elif not all(isinstance(k, int) for k in kth): - raise TypeError("kth is a sequence, but not all elements are integers") - - nkth = len(kth) - if nkth == 0 or a.size == 0: - return dpnp.copy(a) - - # validate kth - kth = list(kth) - for i in range(nkth): - if kth[i] < 0: - kth[i] += length - - if not 0 <= kth[i] < length: - raise ValueError(f"kth(={kth[i]}) out of bounds {length}") - - dt = a.dtype - if ( - nd > 1 - or nkth > 1 - or dpnp.issubdtype(dt, dpnp.unsignedinteger) - or dt in (dpnp.int8, dpnp.int16) - or dpnp.is_cuda_backend(a.get_array()) - ): - # sort is a faster path in case of ndim > 1 - return dpnp.sort(a, axis=axis) - - desc = dpnp.get_dpnp_descriptor(a, copy_when_nondefault_queue=False) - return dpnp_partition(desc, kth[0], axis, kind, order).get_pyobj() + return _wrap_partition_argpartition( + a, kth, axis, kind, order, return_indices=False + ) def sort(a, axis=-1, kind=None, order=None, *, descending=False, stable=None): diff --git a/dpnp/dpnp_utils/dpnp_algo_utils.pyx b/dpnp/dpnp_utils/dpnp_algo_utils.pyx index 5a3fdd397501..ff5d0f6eabcd 100644 --- a/dpnp/dpnp_utils/dpnp_algo_utils.pyx +++ b/dpnp/dpnp_utils/dpnp_algo_utils.pyx @@ -429,7 +429,7 @@ cdef dpnp_descriptor create_output_descriptor( @cython.profile(False) -cpdef inline tuple _object_to_tuple(object obj): +cpdef tuple _object_to_tuple(object obj): """ Converts Python object into tuple """ diff --git a/dpnp/dpnp_utils/dpnp_utils_statistics.py b/dpnp/dpnp_utils/dpnp_utils_statistics.py index 98d10bdd4afd..03160ec06d4c 100644 --- a/dpnp/dpnp_utils/dpnp_utils_statistics.py +++ b/dpnp/dpnp_utils/dpnp_utils_statistics.py @@ -29,7 +29,13 @@ import math import warnings +import dpctl.utils as dpu + import dpnp + +# pylint: disable=no-name-in-module +import dpnp.libtensor._statistics_impl as statistics_ext +import dpnp.libtensor._tensor_impl as ti import dpnp.tensor as dpt from dpnp.dpnp_array import dpnp_array from dpnp.exceptions import ExecutionPlacementError @@ -58,6 +64,81 @@ def _calc_median(a, axis, out=None): return res +def _calc_median_by_select(a, ignore_nan, out=None): + """ + Compute the median of an array along its last axis by radix select, + ignoring NaNs if `ignore_nan`. + + """ + + usm_a = dpnp.get_usm_ndarray(a) + exec_q = usm_a.sycl_queue + if out is not None: + dpnp.check_supported_arrays_type(out) + if dpt.get_execution_queue((exec_q, out.sycl_queue)) is None: + raise ExecutionPlacementError( + "Input and output allocation queues are not compatible" + ) + if out.shape != usm_a.shape[:-1]: + raise ValueError( + f"Output array of shape {usm_a.shape[:-1]} is needed, " + f"got {out.shape}." + ) + + _manager = dpu.SequentialOrderManager[exec_q] + dep_evs = _manager.submitted_events + + # the medians are selected along the last axis of a C-contiguous array + if not usm_a.flags.c_contiguous: + usm_tmp = dpt.empty_like(usm_a, order="C") + ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( + src=usm_a, dst=usm_tmp, sycl_queue=exec_q, depends=dep_evs + ) + _manager.add_event_pair(ht_ev, copy_ev) + usm_a = usm_tmp + dep_evs = [copy_ev] + + # the medians of floating data types keep their type, as with numpy + if dpnp.issubdtype(usm_a.dtype, dpnp.floating): + res_dt = usm_a.dtype + else: + res_dt = dpnp.default_float_type(sycl_queue=exec_q) + + usm_res = dpt.empty( + usm_a.shape[:-1], + dtype=res_dt, + usm_type=usm_a.usm_type, + order="C", + sycl_queue=exec_q, + ) + # set to whether a row has NaNs only, the one case warned about + usm_all_nan = None + if ignore_nan: + usm_all_nan = dpt.empty( + 1, dtype=dpt.uint32, usm_type=usm_a.usm_type, sycl_queue=exec_q + ) + + ht_ev, median_ev = statistics_ext._median( + usm_a, + ignore_nan, + usm_res, + sycl_queue=exec_q, + all_nan=usm_all_nan, + depends=dep_evs, + ) + _manager.add_event_pair(ht_ev, median_ev) + + # reading a scalar through asnumpy is several times cheaper than int() + if ignore_nan and dpt.asnumpy(usm_all_nan)[0]: + warnings.warn("All-NaN slice encountered", RuntimeWarning, stacklevel=4) + + res = dpnp_array._create_from_usm_ndarray(usm_res) + if out is None: + return res + out[...] = res + return out + + def _calc_nanmedian(a, out=None): """Compute the median of an array along a specified axis, ignoring NaNs.""" mask = dpnp.isnan(a) @@ -197,6 +278,29 @@ def dpnp_cov( return c.squeeze() +def _calc_median_by_sort(a, axis, out, overwrite_input, ignore_nan): + """Compute the median of an array along a specified axis by sorting.""" + + if overwrite_input: + if isinstance(a, dpt.usm_ndarray): + # dpnp.ndarray.sort only works with dpnp_array + a = dpnp_array._create_from_usm_ndarray(a) + a.sort(axis=axis) + a_sorted = a + else: + a_sorted = dpnp.sort(a, axis=axis) + + if ignore_nan: + # sorting puts NaNs at the end + assert axis == -1 + return _calc_nanmedian(a_sorted, out=out) + + # We can't pass keepdims and use it in dpnp.mean and dpnp.any + # because of the reshape hack that might have been used in + # `_flatten_array_along_axes`. + return _calc_median(a_sorted, axis=axis, out=out) + + def dpnp_median( a, axis=None, @@ -212,8 +316,16 @@ def dpnp_median( _axis = range(a_ndim) if axis is None else axis _axis = normalize_axis_tuple(_axis, a_ndim) + # the medians of real data types are selected rather than sorted + use_select = a.size > 0 and not dpnp.issubdtype( + a.dtype, dpnp.complexfloating + ) + if a.size == 0 or not dpnp.issubdtype(a.dtype, dpnp.inexact): + # there are no NaNs to ignore + ignore_nan = False + if len(_axis) == 1: - if ignore_nan: + if ignore_nan or use_select: a = dpnp.moveaxis(a, _axis[0], -1) axis = -1 else: @@ -229,29 +341,37 @@ def dpnp_median( ) axis = -1 - if overwrite_input: - if isinstance(a, dpt.usm_ndarray): - # dpnp.ndarray.sort only works with dpnp_array - a = dpnp_array._create_from_usm_ndarray(a) - a.sort(axis=axis) - a_sorted = a - else: - a_sorted = dpnp.sort(a, axis=axis) - - if ignore_nan: - # sorting puts NaNs at the end - assert axis == -1 - res = _calc_nanmedian(a_sorted, out=out) - else: - # We can't pass keepdims and use it in dpnp.mean and dpnp.any - # because of the reshape hack that might have been used in - # `_flatten_array_along_axes`. - res = _calc_median(a_sorted, axis=axis, out=out) - + res_out = out if keepdims: res_shape = list(a_shape) for i in _axis: res_shape[i] = 1 - res = res.reshape(tuple(res_shape)) + res_shape = tuple(res_shape) + + if out is not None: + dpnp.check_supported_arrays_type(out) + if out.shape != res_shape: + raise ValueError( + f"Output array of shape {res_shape} is needed, " + f"got {out.shape}." + ) + # the medians go to a view of `out` without the reduced axes + res_out = out[ + tuple(0 if i in _axis else slice(None) for i in range(a_ndim)) + ] + + if use_select: + assert axis == -1 + res = _calc_median_by_select(a, ignore_nan, out=res_out) + else: + res = _calc_median_by_sort( + a, axis, res_out, overwrite_input, ignore_nan + ) + + if out is None: + if keepdims: + res = res.reshape(res_shape) + return dpnp.get_result_array(res) - return dpnp.get_result_array(res, out) + dpnp.get_result_array(res, res_out) + return dpnp.get_result_array(out) diff --git a/dpnp/libtensor/include/kernels/sorting/radix_select.hpp b/dpnp/libtensor/include/kernels/sorting/radix_select.hpp index 82084e8b2f55..5c4d13ab9e0a 100644 --- a/dpnp/libtensor/include/kernels/sorting/radix_select.hpp +++ b/dpnp/libtensor/include/kernels/sorting/radix_select.hpp @@ -29,7 +29,8 @@ //===----------------------------------------------------------------------===// /// /// \file -/// This file defines radix select kernels for tensor topk operation. +/// This file defines radix select kernels for tensor topk, partition, +/// argpartition and median operations. //===----------------------------------------------------------------------===// #pragma once @@ -105,21 +106,386 @@ kernels_bundle get_kernels_bundle(sycl::queue &exec_q) return kernels_bundle{std::move(kb), max_wg_size}; } +//----------------------------------------------------------------------- +// radix select: what is written for each row +//----------------------------------------------------------------------- + +enum class select_output +{ + // the k smallest elements, in unspecified order + selected, + // the whole row, the elements ordering before the k-th smallest ahead of + // those equal to it, ahead of the others + whole_row, + // the lower and the upper median + median, +}; + +/*! @brief Sink which discards the elements passed to it */ +struct NoOut +{ + void operator()(std::uint64_t, std::size_t) const {} +}; + +/*! @brief Stores the value and the index of element `i` of a row as the + * `j`-th selected element of the row */ +template +struct SelectedOut +{ + const ValueT *arg; + ValueT *vals; + IndexT *inds; + + void operator()(std::uint64_t j, std::size_t i) const + { + vals[j] = arg[i]; + inds[j] = static_cast(i); + } +}; + +/*! @brief Stores the value of element `i` of a row at position `j` */ +template +struct ValueOut +{ + const ValueT *arg; + ValueT *vals; + + void operator()(std::uint64_t j, std::size_t i) const { vals[j] = arg[i]; } +}; + +/*! @brief Stores the index of element `i` of a row at position `j` */ +template +struct IndexOut +{ + IndexT *inds; + + void operator()(std::uint64_t j, std::size_t i) const + { + inds[j] = static_cast(i); + } +}; + +/*! @brief Passes element `i` to `out` at position `offset + j` */ +template +struct ShiftedOut +{ + OutT out; + std::uint64_t offset; + + void operator()(std::uint64_t j, std::size_t i) const + { + out(offset + j, i); + } +}; + +template +ShiftedOut shifted(const OutT &out, std::uint64_t offset) +{ + return {out, offset}; +} + +inline NoOut shifted(const NoOut &, std::uint64_t) { return {}; } + +/*! @brief Passes candidate `i` of a row to `out` as the element of the row it + * lists */ +template +struct CandidateOut +{ + const std::uint32_t *cand; + OutT out; + + void operator()(std::uint64_t j, std::size_t i) const { out(j, cand[i]); } +}; + +template +CandidateOut via_candidates(const std::uint32_t *cand, const OutT &out) +{ + return {cand, out}; +} + +inline NoOut via_candidates(const std::uint32_t *, const NoOut &) { return {}; } + +/*! @brief Lists element `i` of a row as its `j`-th candidate */ +struct CandidateListOut +{ + std::uint32_t *cand; + + void operator()(std::uint64_t j, std::size_t i) const + { + cand[j] = static_cast(i); + } +}; + +/*! @brief The `k` smallest elements of each row and their indices, written + * into `(n_iters, k)` arrays */ +template +struct TopKOp +{ + using value_type = ValueT; + static constexpr select_output output = select_output::selected; + + ValueT *vals; + IndexT *inds; + std::size_t k; + + std::uint64_t rank(std::uint64_t) const { return k; } + + SelectedOut row_out(std::size_t iter_id, + const ValueT *row_arg) const + { + return {row_arg, vals + iter_id * k, inds + iter_id * k}; + } +}; + +/*! @brief Each row partitioned about its `k`-th smallest element, written + * into an `(n_iters, n_values)` array */ +template +struct PartitionOp +{ + using value_type = ValueT; + static constexpr select_output output = select_output::whole_row; + + ValueT *vals; + std::size_t n_values; + std::size_t k; + + std::uint64_t rank(std::uint64_t) const { return k; } + + ValueOut row_out(std::size_t iter_id, const ValueT *row_arg) const + { + return {row_arg, vals + iter_id * n_values}; + } +}; + +/*! @brief The indices which partition each row about its `k`-th smallest + * element, written into an `(n_iters, n_values)` array */ +template +struct ArgPartitionOp +{ + using value_type = ValueT; + static constexpr select_output output = select_output::whole_row; + + IndexT *inds; + std::size_t n_values; + std::size_t k; + + std::uint64_t rank(std::uint64_t) const { return k; } + + IndexOut row_out(std::size_t iter_id, const ValueT *) const + { + return {inds + iter_id * n_values}; + } +}; + +/*! @brief The median of each row, the mean of its lower and upper medians, + * written as `ResT` into an `(n_iters,)` array; it is NaN for a row with NaNs + * unless `ignore_nan`, and for a row of NaNs only, which then sets + * `*all_nan` if it is not null */ +template +struct MedianOp +{ + using value_type = ValueT; + static constexpr select_output output = select_output::median; + + ResT *out; + std::uint32_t *all_nan; + std::size_t n_values; + bool ignore_nan; + + std::uint64_t n_valid(std::uint64_t n_nan) const + { + return n_values - n_nan; + } + + /*! @brief 1-based rank of the upper median, 0 if the median is NaN */ + std::uint64_t rank(std::uint64_t n_nan) const + { + if (n_nan > 0 && !ignore_nan) { + return 0; + } + const std::uint64_t n = n_valid(n_nan); + return (n == 0) ? 0 : n / 2 + 1; + } + + /*! @brief 0-based rank of the lower median */ + std::uint64_t lo_rank(std::uint64_t n_nan) const + { + return (n_valid(n_nan) - 1) / 2; + } + + /*! @brief Whether the lower median is the largest element ordering + * before the upper one, of 1-based rank `k`, rather than equal to it */ + bool lo_before(std::uint64_t k, + std::uint64_t n_less, + std::uint64_t n_nan) const + { + return (lo_rank(n_nan) + 1 != k) && (n_less + 1 == k); + } + + NoOut row_out(std::size_t, const ValueT *) const { return {}; } + + void write(std::size_t iter_id, ValueT lo, ValueT hi) const + { + // a single median is written as is, its mean with itself could + // overflow + if (lo == hi) { + out[iter_id] = static_cast(hi); + } + else if constexpr (std::is_same_v) { + // as numpy, the mean of half floats is taken in single precision + out[iter_id] = static_cast( + (static_cast(lo) + static_cast(hi)) / 2); + } + else { + out[iter_id] = + (static_cast(lo) + static_cast(hi)) / ResT(2); + } + } + + void write_nan(std::size_t iter_id) const + { + if constexpr (!std::is_integral_v) { + out[iter_id] = std::numeric_limits::quiet_NaN(); + // with NaNs ignored, the median is NaN for a row of NaNs only + if (ignore_nan && all_nan != nullptr) { + sycl::atomic_ref + flag(*all_nan); + flag.store(1); + } + } + } + + /*! @brief Writes the median of a row given the keys of its lower and + * upper medians */ + template + void write_median_keys(std::size_t iter_id, + WideKeyT lo_key, + WideKeyT hi_key) const + { + using KeyT = radix_utils::radix_key_t; + write(iter_id, + radix_utils::ordered_radix_value(KeyT(lo_key)), + radix_utils::ordered_radix_value(KeyT(hi_key))); + } + + /*! @brief Writes the median of a row given the largest keys of its + * elements ordering before the upper median, and of those equal to it */ + template + void write_keys(std::size_t iter_id, + std::uint64_t k, + std::uint64_t n_less, + std::uint64_t n_nan, + WideKeyT less_max, + WideKeyT eq_max) const + { + write_median_keys( + iter_id, (lo_before(k, n_less, n_nan)) ? less_max : eq_max, eq_max); + } +}; + +/*! @brief Whether the first pass counts the NaNs of the rows */ +template +constexpr bool counts_nan() +{ + return OpT::output == select_output::median && + !std::is_integral_v; +} + +/*! @brief Key shared by the NaNs, 0 if the rows' NaNs are not counted */ +template +radix_utils::radix_key_t nan_key() +{ + using ValueT = typename OpT::value_type; + if constexpr (counts_nan()) { + return radix_utils::ordered_radix_key( + std::numeric_limits::quiet_NaN()); + } + else { + return 0; + } +} + +/*! @brief Whether the selection of a row ends with a bucket of `count` + * elements holding the sought one at rank `k_rem` */ +template +bool bucket_resolves(std::uint64_t count, std::uint64_t k_rem) +{ + if constexpr (OpT::output == select_output::selected) { + // the whole bucket is among the k smallest + return count == k_rem; + } + else { + // the sought element itself is needed, the bucket must hold it only + // or be resolved to the last digit + return count == 1; + } +} + +/*! @brief Whether the selection stops once the keys sharing the prefix are + * all equal; the k smallest elements are resolved by counts alone, while the + * sought element otherwise needs a bucket of one, which runs of equal values + * never reach before the last pass. Equal keys never split, so whether they + * are is only checked in a pass whose bucket keeps all of the keys which had + * the prefix, which spread out keys rarely do */ +template +constexpr bool stops_on_equal() +{ + return OpT::output != select_output::selected; +} + +/*! @brief Sink for the elements ordering after the sought one, which only a + * partition of the whole row writes */ +template +auto greater_out(const OutT &out) +{ + if constexpr (OpT::output == select_output::whole_row) { + return out; + } + else { + return NoOut{}; + } +} + +/*! @brief Unsigned type holding a key, at least 32 bits wide for reductions */ +template +using wide_key_t = std:: + conditional_t<(sizeof(KeyT) < sizeof(std::uint32_t)), std::uint32_t, KeyT>; + //----------------------------------------------------------------------- // radix select: work-group level building blocks //----------------------------------------------------------------------- +/*! @brief What a work-item saw of the keys whose digits it counted */ +template +struct digit_stats +{ + // number of the keys equal to the NaN key, if counted + std::uint32_t n_nan; + // bitwise or and and of the keys, if tracked, else no bits and all of + // them, so the keys are all equal if and only if these are + wide_key_t key_or; + wide_key_t key_and; +}; + /*! @brief Adds to the zeroed local histogram `hist` the digits at `shift` of - * the keys in `[begin, end)` whose prefix under `mask` is `desired` */ -template -void count_digits(const sycl::nd_item<1> &ndit, - const KeyFnT &key_fn, - std::size_t begin, - std::size_t end, - KeyT desired, - KeyT mask, - std::uint32_t shift, - const HistAccT &hist) + * the keys in `[begin, end)` whose prefix under `mask` is `desired`, and + * returns how many of the work-item's keys, if `count_nan`, equal `nan_key`, + * and, if `track_bits`, their bitwise or and and */ +template +digit_stats count_digits(const sycl::nd_item<1> &ndit, + const KeyFnT &key_fn, + std::size_t begin, + std::size_t end, + KeyT desired, + KeyT mask, + std::uint32_t shift, + const HistAccT &hist, + KeyT nan_key = KeyT{0}) { using AtomicT = sycl::atomic_ref &ndit, const std::size_t lid = ndit.get_local_linear_id(); const std::size_t wg_size = ndit.get_local_range(0); + using WideKeyT = wide_key_t; + digit_stats stats{0, WideKeyT(0), + WideKeyT(std::numeric_limits::max())}; for (std::size_t i0 = begin; i0 < end; i0 += wg_size) { const std::size_t i = i0 + lid; @@ -139,6 +508,15 @@ void count_digits(const sycl::nd_item<1> &ndit, const KeyT key = key_fn(i); match = ((key & mask) == desired); bucket_id = radix_utils::get_bucket_id(key, shift); + if constexpr (count_nan) { + stats.n_nan += (match && key == nan_key); + } + if constexpr (track_bits) { + if (match) { + stats.key_or |= WideKeyT(key); + stats.key_and &= WideKeyT(key); + } + } } // a sub-group whose keys all land in one bucket, as for runs of @@ -154,6 +532,7 @@ void count_digits(const sycl::nd_item<1> &ndit, AtomicT(hist[bucket_id]).fetch_add(std::uint32_t(1)); } } + return stats; } /*! @brief Finds the bucket of the complete local histogram `hist` holding @@ -211,13 +590,16 @@ bucket_info find_bucket(const sycl::nd_item<1> &ndit, inline constexpr std::size_t max_tile_size = (std::size_t(1) << 16) - 1; /*! @brief Passes the elements `i` of `[begin, end)` whose prefix under `mask` - * is less than `desired` to `less_out(j, i)`, and the first `n_ties` of the - * row equal to it to `tie_out(j, i)`, `j` being their rank in index order */ + * is less than `desired` to `less_out(j, i)`, the first `n_ties` of the row + * equal to it to `tie_out(j, i)`, and the others to `greater_out(j, i)`, `j` + * being their rank in index order among the row's elements of their kind; + * sinks of type `NoOut` are skipped */ template + typename TieOutT, + typename GreaterOutT = NoOut> void gather_selected(const sycl::nd_item<1> &ndit, const KeyFnT &key_fn, std::size_t begin, @@ -229,8 +611,12 @@ void gather_selected(const sycl::nd_item<1> &ndit, std::uint64_t n_ties_before, // ties before begin std::uint64_t n_ties, const LessOutT &less_out, - const TieOutT &tie_out) + const TieOutT &tie_out, + const GreaterOutT &greater_out = {}) { + static constexpr bool keeps_less = !std::is_same_v; + static constexpr bool keeps_greater = !std::is_same_v; + // an element's kind and its rank within a tile are counted in the low // (less) and high (tie) halves of one integer static constexpr std::uint32_t less_flag = 1; @@ -253,8 +639,11 @@ void gather_selected(const sycl::nd_item<1> &ndit, std::uint64_t less_done = 0; std::uint64_t ties_done = 0; for (std::size_t t0 = begin; t0 < end; t0 += tile_size) { - if (less_done >= n_less && n_ties_before + ties_done >= n_ties) { - break; + if constexpr (!keeps_greater) { + if ((!keeps_less || less_done >= n_less) && + n_ties_before + ties_done >= n_ties) { + break; + } } const std::size_t sg_begin = t0 + sg_tile_offset; @@ -291,21 +680,28 @@ void gather_selected(const sycl::nd_item<1> &ndit, #pragma unroll for (std::uint32_t j = 0; j < elems_per_wi; ++j) { - if (flags[j] == 0) { - continue; - } const std::size_t i = sg_begin + j * sg_size + lane_id; const std::uint32_t r = sg_offset + ranks[j]; + // numbers of the row's less and tie elements before i + const std::uint64_t less_rank = + less_offset + less_done + (r & half_mask); + const std::uint64_t tie_rank = + n_ties_before + ties_done + (r >> 16); if (flags[j] == less_flag) { - less_out(less_offset + less_done + (r & half_mask), i); + if constexpr (keeps_less) { + less_out(less_rank, i); + } } - else { - const std::uint64_t tie_rank = - n_ties_before + ties_done + (r >> 16); + else if (flags[j] == tie_flag) { if (tie_rank < n_ties) { tie_out(tie_rank, i); } } + else if constexpr (keeps_greater) { + if (i < end) { + greater_out(i - less_rank - tie_rank, i); + } + } } less_done += (tile_count & half_mask); @@ -313,6 +709,41 @@ void gather_selected(const sycl::nd_item<1> &ndit, } } +/*! @brief Largest keys, over the work-group, of the elements in `[begin, end)` + * whose prefix under `mask` is less than `desired`, and of those for which it + * equals it; 0 when there are none */ +template +std::pair, wide_key_t> + max_keys_about(const sycl::nd_item<1> &ndit, + const KeyFnT &key_fn, + std::size_t begin, + std::size_t end, + KeyT desired, + KeyT mask) +{ + using WideKeyT = wide_key_t; + + const std::size_t lid = ndit.get_local_linear_id(); + const std::size_t wg_size = ndit.get_local_range(0); + + WideKeyT less_max = 0; + WideKeyT eq_max = 0; + for (std::size_t i = begin + lid; i < end; i += wg_size) { + const KeyT key = key_fn(i); + const KeyT prefix = key & mask; + if (prefix < desired) { + less_max = std::max(less_max, WideKeyT(key)); + } + else if (prefix == desired) { + eq_max = std::max(eq_max, WideKeyT(key)); + } + } + + const auto &wg = ndit.get_group(); + return {sycl::reduce_over_group(wg, less_max, sycl::maximum()), + sycl::reduce_over_group(wg, eq_max, sycl::maximum())}; +} + template struct RowKey { @@ -337,83 +768,126 @@ struct CandidateKey } }; -/*! @brief Stores the value and the index of element `i` of a row as the - * `j`-th selected element of the row */ -template -struct SelectedOut +/*! @brief Outcome of selecting the element of rank `k` of a row */ +template +struct row_selection { - const ValueT *arg; - ValueT *vals; - IndexT *inds; - - void operator()(std::uint64_t j, std::size_t i) const - { - vals[j] = arg[i]; - inds[j] = static_cast(i); - } + // the elements whose key has the prefix `desired` under `mask`, the + // sought one among them, number `n_eq`; when the selection resolves the + // sought element they all compare equal to it, or it is the only one + KeyT desired; + KeyT mask; + // 1-based rank of the sought element, 0 if there is none + std::uint64_t k; + // number of elements whose prefix orders before `desired` + std::uint64_t n_less; + std::uint64_t n_eq; + std::uint64_t n_nan; }; -/*! @brief Stores the value and the index of candidate `i` of a row as the - * `j`-th selected element of the row */ -template -struct CandidateSelectedOut +/*! @brief Selects, by the work-group, the element of rank `op.rank(n_nan)` + * of a row of `n_values` elements */ +template +row_selection select_over_group(const sycl::nd_item<1> &ndit, + const OpT &op, + const KeyFnT &key_fn, + std::size_t n_values, + const HistAccT &hist) { - const ValueT *arg; - const std::uint32_t *cand; - ValueT *vals; - IndexT *inds; + static constexpr std::uint32_t n_passes = n_radix_passes(); + static constexpr std::uint32_t key_bits = + radix_utils::number_of_bits_in_type(); + static constexpr bool count_nan = counts_nan(); + static constexpr bool stop_on_equal = stops_on_equal(); + using WideKeyT = wide_key_t; - void operator()(std::uint64_t j, std::size_t i) const - { - const std::uint32_t c = cand[i]; - vals[j] = arg[c]; - inds[j] = static_cast(c); - } -}; + const auto &wg = ndit.get_group(); + const std::size_t lid = ndit.get_local_linear_id(); + const std::size_t wg_size = ndit.get_local_range(0); -/*! @brief Lists element `i` of a row as its `j`-th candidate */ -struct CandidateListOut -{ - std::uint32_t *cand; + row_selection sel{KeyT{0}, KeyT{0}, 0, 0, 0, 0}; + std::uint64_t k_rem = 0; + for (std::uint32_t pass = 0; pass < n_passes; ++pass) { + const std::uint32_t shift = key_bits - (pass + 1) * radix_bits; - void operator()(std::uint64_t j, std::size_t i) const - { - cand[j] = static_cast(i); + for (std::size_t b = lid; b < radix_states; b += wg_size) { + hist[b] = 0; + } + sycl::group_barrier(wg); + + const digit_stats wi_stats = + count_digits( + ndit, key_fn, 0, n_values, sel.desired, sel.mask, shift, hist, + nan_key()); + sycl::group_barrier(wg); + + if (pass == 0) { + if constexpr (count_nan) { + sel.n_nan = + sycl::reduce_over_group(wg, std::uint64_t(wi_stats.n_nan), + sycl::plus()); + } + sel.k = op.rank(sel.n_nan); + if (sel.k == 0) { + return sel; + } + k_rem = sel.k; + } + + const bucket_info info = find_bucket(ndit, hist, k_rem); + + // the keys having the prefix, among which the sought one is, are all + // equal, so the bucket of this pass resolves it + bool all_equal = false; + if constexpr (stop_on_equal) { + if (info.count == ((pass == 0) ? n_values : sel.n_eq)) { + all_equal = + (sycl::reduce_over_group(wg, wi_stats.key_or, + sycl::bit_or()) == + sycl::reduce_over_group(wg, wi_stats.key_and, + sycl::bit_and())); + } + } + sel.desired |= static_cast(KeyT(info.bucket_id) << shift); + sel.mask |= static_cast(KeyT(radix_mask) << shift); + k_rem -= info.before; + sel.n_eq = info.count; + + if (all_equal || bucket_resolves(info.count, k_rem)) { + break; + } } -}; + sel.n_less = sel.k - k_rem; + return sel; +} //----------------------------------------------------------------------- // radix select: one work-group per row //----------------------------------------------------------------------- -template +template class radix_select_one_group_krn; template + typename OpT> sycl::event radix_select_one_group_submit(sycl::queue &exec_q, std::size_t n_iters, std::size_t n_values, - std::size_t k, const ValueT *arg_ptr, - ValueT *vals_ptr, - IndexT *inds_ptr, + const OpT &op, std::size_t wg_size, const std::vector &depends) { using KeyT = radix_utils::radix_key_t; using KernelName = - radix_select_one_group_krn; - - static constexpr std::uint32_t n_passes = n_radix_passes(); - static constexpr std::uint32_t key_bits = - radix_utils::number_of_bits_in_type(); + radix_select_one_group_krn; const kernels_bundle bundle = get_kernels_bundle(exec_q); wg_size = std::min(wg_size, bundle.max_wg_size); @@ -433,48 +907,40 @@ sycl::event cgh.parallel_for(ndRange, [=](sycl::nd_item<1> ndit) { const std::size_t iter_id = ndit.get_group(0); - const std::size_t lid = ndit.get_local_linear_id(); - const RowKey key_fn{arg_ptr + - iter_id * n_values}; + const ValueT *row_arg = arg_ptr + iter_id * n_values; + const RowKey key_fn{row_arg}; - KeyT desired{0}; - KeyT mask{0}; - std::uint64_t k_rem = k; - for (std::uint32_t pass = 0; pass < n_passes; ++pass) { - const std::uint32_t shift = key_bits - (pass + 1) * radix_bits; + const row_selection sel = + select_over_group(ndit, op, key_fn, + n_values, hist); - for (std::size_t b = lid; b < radix_states; b += wg_size) { - hist[b] = 0; + if constexpr (OpT::output == select_output::median) { + if (sel.k == 0) { + if (ndit.get_local_linear_id() == 0) { + op.write_nan(iter_id); + } + return; } - sycl::group_barrier(ndit.get_group()); - - count_digits(ndit, key_fn, 0, n_values, desired, mask, shift, - hist); - sycl::group_barrier(ndit.get_group()); - - const bucket_info info = find_bucket(ndit, hist, k_rem); - - desired |= static_cast(KeyT(info.bucket_id) << shift); - mask |= static_cast(KeyT(radix_mask) << shift); - k_rem -= info.before; - - // all keys with the prefix are selected, the rest of the - // digits does not matter - if (info.count == k_rem) { - break; + const auto [less_max, eq_max] = max_keys_about( + ndit, key_fn, 0, n_values, sel.desired, sel.mask); + if (ndit.get_local_linear_id() == 0) { + op.write_keys(iter_id, sel.k, sel.n_less, sel.n_nan, + less_max, eq_max); } } - - const std::uint64_t n_less = k - k_rem; - using OutT = SelectedOut; - const ValueT *row_arg = arg_ptr + iter_id * n_values; - ValueT *row_vals = vals_ptr + iter_id * k; - IndexT *row_inds = inds_ptr + iter_id * k; - gather_selected( - ndit, key_fn, 0, n_values, desired, mask, 0, n_less, 0, k_rem, - OutT{row_arg, row_vals, row_inds}, - OutT{row_arg, row_vals + n_less, row_inds + n_less}); + else { + const auto row_out = op.row_out(iter_id, row_arg); + const std::uint64_t n_ties = + (OpT::output == select_output::whole_row) + ? sel.n_eq + : sel.k - sel.n_less; + gather_selected( + ndit, key_fn, 0, n_values, sel.desired, sel.mask, 0, + sel.n_less, 0, n_ties, row_out, + shifted(row_out, sel.n_less), + greater_out(shifted(row_out, sel.n_less + sel.n_eq))); + } }); }); } @@ -483,34 +949,108 @@ sycl::event // radix select: one sub-group per row //----------------------------------------------------------------------- -template +template class radix_select_sub_group_krn; +/*! @brief Selection of a row of `n` elements by the sub-group `sg`, ranking + * each element by counting from memory the elements ordering before it; for a + * sub-group too small to hold the row in registers */ +template +void select_sub_group_row_from_memory(const sycl::sub_group &sg, + const OpT &op, + std::size_t iter_id, + const ValueT *row_arg, + std::uint32_t n) +{ + using KeyT = radix_utils::radix_key_t; + + const std::uint32_t lane_id = sg.get_local_linear_id(); + const std::uint32_t sg_size = sg.get_local_linear_range(); + + const auto key_at = [=](std::uint32_t i) { + return radix_utils::ordered_radix_key(row_arg[i]); + }; + const auto rank_of = [=](std::uint32_t i, KeyT ki) { + std::uint32_t rank = 0; + for (std::uint32_t j = 0; j < n; ++j) { + const KeyT kj = key_at(j); + rank += (kj < ki) || (kj == ki && j < i); + } + return rank; + }; + + if constexpr (OpT::output == select_output::median) { + std::uint32_t n_nan = 0; + if constexpr (counts_nan()) { + const KeyT nan = nan_key(); + std::uint32_t wi_nan = 0; + for (std::uint32_t i = lane_id; i < n; i += sg_size) { + wi_nan += (key_at(i) == nan); + } + n_nan = sycl::reduce_over_group(sg, wi_nan, + sycl::plus()); + } + const std::uint64_t k = op.rank(n_nan); + if (k == 0) { + if (lane_id == 0) { + op.write_nan(iter_id); + } + return; + } + // the ranks are distinct, so each median is held by a single lane + using WideKeyT = wide_key_t; + const std::uint64_t lo = op.lo_rank(n_nan); + WideKeyT lo_key = 0; + WideKeyT hi_key = 0; + for (std::uint32_t i = lane_id; i < n; i += sg_size) { + const KeyT ki = key_at(i); + const std::uint32_t rank = rank_of(i, ki); + if (rank == lo) { + lo_key = ki; + } + if (rank + 1 == k) { + hi_key = ki; + } + } + lo_key = sycl::reduce_over_group(sg, lo_key, sycl::maximum()); + hi_key = sycl::reduce_over_group(sg, hi_key, sycl::maximum()); + if (lane_id == 0) { + op.write_median_keys(iter_id, lo_key, hi_key); + } + } + else { + const std::uint64_t n_out = + (OpT::output == select_output::whole_row) ? n : op.rank(0); + const auto row_out = op.row_out(iter_id, row_arg); + for (std::uint32_t i = lane_id; i < n; i += sg_size) { + const std::uint32_t rank = rank_of(i, key_at(i)); + if (rank < n_out) { + row_out(rank, i); + } + } + } +} + /*! @brief Selection with a sub-group per row of at most `max_chunks` times the * sub-group size, each element is ranked by counting the elements ordering * before it and written at its rank, so the selection comes out sorted */ template + typename OpT> sycl::event radix_select_sub_group_submit(sycl::queue &exec_q, std::size_t n_iters, std::size_t n_values, - std::size_t k, const ValueT *arg_ptr, - ValueT *vals_ptr, - IndexT *inds_ptr, + const OpT &op, std::size_t min_sg_size, std::size_t wg_size, const std::vector &depends) { using KeyT = radix_utils::radix_key_t; using KernelName = - radix_select_sub_group_krn; + radix_select_sub_group_krn; const kernels_bundle bundle = get_kernels_bundle(exec_q); wg_size = std::min(wg_size, bundle.max_wg_size); @@ -544,29 +1084,13 @@ sycl::event for (std::size_t iter_id = sg_id; iter_id < n_iters; iter_id += n_sub_groups) { const ValueT *row_arg = arg_ptr + iter_id * n_values; - ValueT *row_vals = vals_ptr + iter_id * k; - IndexT *row_inds = inds_ptr + iter_id * k; // a sub-group smaller than the dispatch assumed, such as a // partial one, cannot hold the row in registers and ranks it // from memory instead if (n_chunks > max_chunks) { - for (std::uint32_t i = lane_id; i < n; i += sg_size) { - const KeyT ki = - radix_utils::ordered_radix_key( - row_arg[i]); - std::uint32_t rank = 0; - for (std::uint32_t j = 0; j < n; ++j) { - const KeyT kj = - radix_utils::ordered_radix_key( - row_arg[j]); - rank += (kj < ki) || (kj == ki && j < i); - } - if (rank < k) { - row_vals[rank] = row_arg[i]; - row_inds[rank] = static_cast(i); - } - } + select_sub_group_row_from_memory( + sg, op, iter_id, row_arg, n); continue; } @@ -607,12 +1131,63 @@ sycl::event } } + if constexpr (OpT::output == select_output::median) { + std::uint32_t n_nan = 0; + if constexpr (counts_nan()) { + const KeyT nan = nan_key(); + std::uint32_t wi_nan = 0; #pragma unroll - for (std::uint32_t c = 0; c < max_chunks; ++c) { - const std::uint32_t i = c * sg_size + lane_id; - if (i < n && ranks[c] < k) { - row_vals[ranks[c]] = row_arg[i]; - row_inds[ranks[c]] = static_cast(i); + for (std::uint32_t c = 0; c < max_chunks; ++c) { + const std::uint32_t i = c * sg_size + lane_id; + wi_nan += (i < n && keys[c] == nan); + } + n_nan = sycl::reduce_over_group( + sg, wi_nan, sycl::plus()); + } + const std::uint64_t k = op.rank(n_nan); + if (k == 0) { + if (lane_id == 0) { + op.write_nan(iter_id); + } + continue; + } + // NaNs rank last, and the ranks are distinct, so + // each median is held by a single lane + using WideKeyT = wide_key_t; + const std::uint64_t lo = op.lo_rank(n_nan); + WideKeyT lo_key = 0; + WideKeyT hi_key = 0; +#pragma unroll + for (std::uint32_t c = 0; c < max_chunks; ++c) { + const std::uint32_t i = c * sg_size + lane_id; + if (i < n) { + if (ranks[c] == lo) { + lo_key = keys[c]; + } + if (ranks[c] + 1 == k) { + hi_key = keys[c]; + } + } + } + lo_key = sycl::reduce_over_group(sg, lo_key, + sycl::maximum()); + hi_key = sycl::reduce_over_group(sg, hi_key, + sycl::maximum()); + if (lane_id == 0) { + op.write_median_keys(iter_id, lo_key, hi_key); + } + } + else { + const std::uint64_t n_out = + (OpT::output == select_output::whole_row) ? n + : op.rank(0); + const auto row_out = op.row_out(iter_id, row_arg); +#pragma unroll + for (std::uint32_t c = 0; c < max_chunks; ++c) { + const std::uint32_t i = c * sg_size + lane_id; + if (i < n && ranks[c] < n_out) { + row_out(ranks[c], i); + } } } } @@ -630,36 +1205,43 @@ struct row_state { KeyT desired; KeyT mask; + // 1-based rank of the sought element, set by the first pass, 0 if there + // is none + std::uint64_t k; // rank of the sought element among the keys having the prefix std::uint64_t k_rem; - // when nonzero, the passes after the first one go over this many - // candidates, the elements which share the prefix of the first pass + // number of keys having the prefix + std::uint64_t n_eq; + std::uint64_t n_nan; + // when nonzero, the passes after `filter_pass` go over this many + // candidates, the elements which share the prefix of that pass std::uint64_t n_cand; - // number of elements selected by the first pass + // number of elements ordering before the candidates std::uint64_t n_less_first; + std::uint32_t filter_pass; std::uint32_t done; // number of work-groups of the row which completed the current pass std::uint32_t n_arrived; }; -template +template class radix_select_init_krn; -template +template class radix_select_count_krn; -template +template class radix_select_filter_krn; -template +template class radix_select_gather_krn; +template +class radix_select_median_segment_krn; + +template +class radix_select_median_row_krn; + /*! @brief Range `[begin, end)` of segment `segment_id` of `n_segments` * covering `[0, nelems)` */ inline std::pair segment_range(std::size_t segment_id, @@ -709,38 +1291,46 @@ class wait_on_throw template + typename OpT> sycl::event radix_select_multi_group_impl(sycl::queue &exec_q, std::size_t n_iters, std::size_t n_values, - std::size_t k, const ValueT *arg_ptr, - ValueT *vals_ptr, - IndexT *inds_ptr, + const OpT &op, std::size_t n_segments, std::size_t wg_size, bool filter, const std::vector &depends) { using KeyT = radix_utils::radix_key_t; + using WideKeyT = wide_key_t; using StateT = row_state; using RowKeyT = RowKey; using CandKeyT = CandidateKey; - using CountKernelName = - radix_select_count_krn; - using FilterKernelName = - radix_select_filter_krn; - using GatherKernelName = - radix_select_gather_krn; static constexpr std::uint32_t n_passes = n_radix_passes(); static constexpr std::uint32_t key_bits = radix_utils::number_of_bits_in_type(); + static constexpr bool count_nan = counts_nan(); + static constexpr bool stop_on_equal = stops_on_equal(); + static constexpr bool is_median = (OpT::output == select_output::median); + + using CountKernelName = radix_select_count_krn; + using FilterKernelName = + radix_select_filter_krn; + using GatherKernelName = + radix_select_gather_krn; + using MedianSegmentKernelName = + radix_select_median_segment_krn; + // the kernel ending the selection over the rows' segments + using LastKernelName = + std::conditional_t; // the kernels share the work-groups of the rows' segments const kernels_bundle bundle = - get_kernels_bundle( + get_kernels_bundle( exec_q); wg_size = std::min(wg_size, bundle.max_wg_size); @@ -751,9 +1341,9 @@ sycl::event wg_size * elems_per_wi > max_tile_size) { throw std::runtime_error("Invalid parameters for radix select"); } - // when filtering, the elements sharing the prefix of the first pass are - // listed, by 32-bit indices within their row, for the later passes to go - // over + // when filtering, the elements sharing the prefix of the first pass to + // leave few enough of them are listed, by 32-bit indices within their + // row, for the later passes to go over filter = filter && (n_passes > 1) && (n_values <= std::numeric_limits::max()); // rows with more candidates than this are not filtered @@ -773,14 +1363,24 @@ sycl::event std::uint32_t *segment_hist_ptr = hist_owner.get(); // per work-group counts of less and tie elements, and their exclusive - // scans over the work-groups of a row + // scans over the work-groups of a row, followed by the per work-group + // counts of NaNs if they are counted, and the bitwise or and and of the + // keys having the prefix if equal keys stop the selection + static constexpr std::size_t n_counts = + 4 + (count_nan ? 1 : 0) + (stop_on_equal ? 2 : 0); auto counts_owner = dpnp::tensor::alloc_utils::smart_malloc_device( - 4 * n_all_segments, exec_q); + n_counts * n_all_segments, exec_q); std::uint64_t *less_count_ptr = counts_owner.get(); std::uint64_t *less_offset_ptr = less_count_ptr + n_all_segments; std::uint64_t *tie_count_ptr = less_offset_ptr + n_all_segments; std::uint64_t *tie_offset_ptr = tie_count_ptr + n_all_segments; + std::uint64_t *nan_count_ptr = tie_offset_ptr + n_all_segments; + std::uint64_t *key_or_ptr = + (stop_on_equal) ? nan_count_ptr + (count_nan ? n_all_segments : 0) + : nullptr; + std::uint64_t *key_and_ptr = + (stop_on_equal) ? key_or_ptr + n_all_segments : nullptr; // candidates of each row, in index order auto cand_owner = @@ -788,6 +1388,13 @@ sycl::event (filter) ? n_iters * cand_cap : 1, exec_q); std::uint32_t *cand_ptr = cand_owner.get(); + // the largest keys of the elements ordering before the upper median and + // of those equal to it, for each work-group + auto partial_owner = + dpnp::tensor::alloc_utils::smart_malloc_device( + (is_median) ? 2 * n_all_segments : 1, exec_q); + WideKeyT *partial_ptr = partial_owner.get(); + // the last kernel submitted, each one depends on the one before; if a // submission throws, it is waited for before the temporaries are freed sycl::event last_ev{}; @@ -796,10 +1403,11 @@ sycl::event last_ev = exec_q.submit([&](sycl::handler &cgh) { cgh.depends_on(depends); - using KernelName = radix_select_init_krn; + using KernelName = radix_select_init_krn; cgh.parallel_for( sycl::range<1>(n_iters), [=](sycl::id<1> id) { - state_ptr[id[0]] = StateT{KeyT{0}, KeyT{0}, k, 0, 0, 0, 0}; + state_ptr[id[0]] = + StateT{KeyT{0}, KeyT{0}, 0, 0, 0, 0, 0, 0, 0, 0, 0}; }); }); @@ -808,7 +1416,8 @@ sycl::event for (std::uint32_t pass = 0; pass < n_passes; ++pass) { const std::uint32_t shift = key_bits - (pass + 1) * radix_bits; const bool last_pass = (pass + 1 == n_passes); - const bool filter_after = filter && (pass == 0); + // the last pass leaves nothing to filter for + const bool filter_after = filter && !last_pass; last_ev = exec_q.submit([&](sycl::handler &cgh) { cgh.depends_on(last_ev); @@ -840,18 +1449,41 @@ sycl::event sycl::group_barrier(wg); const ValueT *row_arg = arg_ptr + iter_id * n_values; + digit_stats wi_stats; if (n_cand) { const auto [begin, end] = segment_range(segment_id, n_segments, n_cand); - count_digits( + wi_stats = count_digits( ndit, CandKeyT{row_arg, cand_ptr + iter_id * cand_cap}, begin, end, desired, mask, shift, hist); } else { const auto [begin, end] = segment_range(segment_id, n_segments, n_values); - count_digits(ndit, RowKeyT{row_arg}, begin, end, desired, - mask, shift, hist); + wi_stats = count_digits( + ndit, RowKeyT{row_arg}, begin, end, desired, mask, + shift, hist, nan_key()); + } + if constexpr (count_nan) { + // the first pass goes over the row + if (pass == 0) { + const std::uint64_t n_nan = sycl::reduce_over_group( + wg, std::uint64_t(wi_stats.n_nan), + sycl::plus()); + if (lid == 0) { + nan_count_ptr[group_id] = n_nan; + } + } + } + if constexpr (stop_on_equal) { + const WideKeyT key_or = sycl::reduce_over_group( + wg, wi_stats.key_or, sycl::bit_or()); + const WideKeyT key_and = sycl::reduce_over_group( + wg, wi_stats.key_and, sycl::bit_and()); + if (lid == 0) { + key_or_ptr[group_id] = key_or; + key_and_ptr[group_id] = key_and; + } } sycl::group_barrier(wg); @@ -880,7 +1512,34 @@ sycl::event sycl::atomic_fence(sycl::memory_order::acquire, sycl::memory_scope::device); - // the last work-group of the row resolves this pass' digit + // the last work-group of the row resolves this pass' digit; + // the state is read ahead of the barriers below, as work-item + // 0 updates it after them + std::uint64_t k = state.k; + std::uint64_t n_nan = state.n_nan; + std::uint64_t k_rem = state.k_rem; + const std::uint64_t n_eq = state.n_eq; + const std::uint32_t filter_pass = state.filter_pass; + if (pass == 0) { + if constexpr (count_nan) { + const std::uint64_t *row_nan_count = + nan_count_ptr + iter_id * n_segments; + n_nan = sycl::joint_reduce( + wg, row_nan_count, row_nan_count + n_segments, + std::uint64_t(0), sycl::plus()); + } + k = op.rank(n_nan); + k_rem = k; + if (k == 0) { + if (lid == 0) { + state.n_nan = n_nan; + state.done = 1; + state.n_arrived = 0; + } + return; + } + } + for (std::size_t b = lid; b < radix_states; b += wg_size) { std::uint64_t s = 0; for (std::size_t j = 0; j < n_segments; ++j) { @@ -890,16 +1549,38 @@ sycl::event } sycl::group_barrier(wg); - const std::uint64_t k_rem = state.k_rem; const bucket_info info = find_bucket(ndit, row_hist, k_rem); const std::uint64_t new_k_rem = k_rem - info.before; - const bool done = last_pass || (info.count == new_k_rem); + + // the keys having the prefix, among which the sought one is, + // are all equal, so the bucket of this pass resolves it + bool all_equal = false; + if constexpr (stop_on_equal) { + if (info.count == ((pass == 0) ? n_values : n_eq)) { + const std::uint64_t *row_key_or = + key_or_ptr + iter_id * n_segments; + const std::uint64_t *row_key_and = + key_and_ptr + iter_id * n_segments; + all_equal = + (sycl::joint_reduce( + wg, row_key_or, row_key_or + n_segments, + std::uint64_t(0), + sycl::bit_or()) == + sycl::joint_reduce( + wg, row_key_and, row_key_and + n_segments, + ~std::uint64_t(0), + sycl::bit_and())); + } + } + const bool done = last_pass || all_equal || + bucket_resolves(info.count, new_k_rem); // the elements of this pass' bucket become the candidates - const bool to_filter = - filter_after && !done && (info.count <= cand_cap); + const bool to_filter = filter_after && !done && (n_cand == 0) && + (info.count <= cand_cap); // the counts restart when the segments switch to candidates - const bool restart = (pass == 0) || (n_cand && pass == 1); + const bool restart = + (pass == 0) || (n_cand && pass == filter_pass + 1); std::uint64_t *row_less_count = less_count_ptr + iter_id * n_segments; std::uint64_t *row_less_offset = @@ -939,20 +1620,24 @@ sycl::event static_cast(KeyT(info.bucket_id) << shift); state.mask = mask | static_cast(KeyT(radix_mask) << shift); + state.k = k; state.k_rem = new_k_rem; + state.n_eq = info.count; + state.n_nan = n_nan; state.done = done; state.n_arrived = 0; if (to_filter) { state.n_cand = info.count; - state.n_less_first = info.before; + state.n_less_first = k - new_k_rem; + state.filter_pass = pass; } } }); }); if (filter_after) { - // writes out the less elements of the first pass and lists its - // candidates + // writes out the elements resolved by the rows filtered after this + // pass and lists their candidates last_ev = exec_q.submit([&](sycl::handler &cgh) { cgh.depends_on(last_ev); cgh.use_kernel_bundle(bundle.kb); @@ -965,72 +1650,159 @@ sycl::event group_id - iter_id * n_segments; const StateT state = state_ptr[iter_id]; - if (state.n_cand == 0) { + if (state.n_cand == 0 || state.filter_pass != pass) { return; } + const ValueT *row_arg = arg_ptr + iter_id * n_values; + const auto row_out = op.row_out(iter_id, row_arg); const auto [begin, end] = segment_range(segment_id, n_segments, n_values); gather_selected( - ndit, RowKeyT{arg_ptr + iter_id * n_values}, begin, - end, state.desired, state.mask, - less_offset_ptr[group_id], less_count_ptr[group_id], - tie_offset_ptr[group_id], state.n_cand, - SelectedOut{ - arg_ptr + iter_id * n_values, - vals_ptr + iter_id * k, inds_ptr + iter_id * k}, - CandidateListOut{cand_ptr + iter_id * cand_cap}); + ndit, RowKeyT{row_arg}, begin, end, state.desired, + state.mask, less_offset_ptr[group_id], + less_count_ptr[group_id], tie_offset_ptr[group_id], + state.n_cand, row_out, + CandidateListOut{cand_ptr + iter_id * cand_cap}, + greater_out(shifted( + row_out, state.n_less_first + state.n_cand))); }); }); } } - last_ev = exec_q.submit([&](sycl::handler &cgh) { - cgh.depends_on(last_ev); - cgh.use_kernel_bundle(bundle.kb); + if constexpr (is_median) { + last_ev = exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(last_ev); + cgh.use_kernel_bundle(bundle.kb); - cgh.parallel_for(ndRange, [=](sycl::nd_item<1> ndit) { - const std::size_t group_id = ndit.get_group(0); - const std::size_t iter_id = group_id / n_segments; - const std::size_t segment_id = group_id - iter_id * n_segments; + cgh.parallel_for( + ndRange, [=](sycl::nd_item<1> ndit) { + const std::size_t group_id = ndit.get_group(0); + const std::size_t iter_id = group_id / n_segments; + const std::size_t segment_id = + group_id - iter_id * n_segments; - const StateT state = state_ptr[iter_id]; + const StateT state = state_ptr[iter_id]; + if (state.k == 0) { + return; + } - const ValueT *row_arg = arg_ptr + iter_id * n_values; - // the first pass wrote out its less elements already - ValueT *row_vals = vals_ptr + iter_id * k + state.n_less_first; - IndexT *row_inds = inds_ptr + iter_id * k + state.n_less_first; - const std::uint64_t n_less = k - state.n_less_first - state.k_rem; - - if (state.n_cand) { - using OutT = CandidateSelectedOut; - const std::uint32_t *row_cand = cand_ptr + iter_id * cand_cap; - const auto [begin, end] = - segment_range(segment_id, n_segments, state.n_cand); - gather_selected( - ndit, CandKeyT{row_arg, row_cand}, begin, end, - state.desired, state.mask, less_offset_ptr[group_id], - less_count_ptr[group_id], tie_offset_ptr[group_id], - state.k_rem, OutT{row_arg, row_cand, row_vals, row_inds}, - OutT{row_arg, row_cand, row_vals + n_less, - row_inds + n_less}); - } - else { - using OutT = SelectedOut; - const auto [begin, end] = - segment_range(segment_id, n_segments, n_values); - gather_selected( - ndit, RowKeyT{row_arg}, begin, end, state.desired, - state.mask, less_offset_ptr[group_id], - less_count_ptr[group_id], tie_offset_ptr[group_id], - state.k_rem, OutT{row_arg, row_vals, row_inds}, - OutT{row_arg, row_vals + n_less, row_inds + n_less}); - } + // the elements ordering before the candidates are gone over + // only if the lower median is among them + const std::uint64_t n_less = state.k - state.k_rem; + const bool in_cand = + state.n_cand && + (n_less > state.n_less_first || + !op.lo_before(state.k, n_less, state.n_nan)); + + const ValueT *row_arg = arg_ptr + iter_id * n_values; + std::pair maxima; + if (in_cand) { + const auto [begin, end] = + segment_range(segment_id, n_segments, state.n_cand); + maxima = max_keys_about( + ndit, + CandKeyT{row_arg, cand_ptr + iter_id * cand_cap}, + begin, end, state.desired, state.mask); + } + else { + const auto [begin, end] = + segment_range(segment_id, n_segments, n_values); + maxima = max_keys_about(ndit, RowKeyT{row_arg}, begin, + end, state.desired, state.mask); + } + if (ndit.get_local_linear_id() == 0) { + partial_ptr[2 * group_id] = maxima.first; + partial_ptr[2 * group_id + 1] = maxima.second; + } + }); }); - }); + + last_ev = exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(last_ev); + + using KernelName = radix_select_median_row_krn; + cgh.parallel_for( + sycl::range<1>(n_iters), [=](sycl::id<1> id) { + const std::size_t iter_id = id[0]; + const StateT state = state_ptr[iter_id]; + if (state.k == 0) { + op.write_nan(iter_id); + return; + } + + const WideKeyT *row_partial = + partial_ptr + 2 * iter_id * n_segments; + WideKeyT less_max = 0; + WideKeyT eq_max = 0; + for (std::size_t j = 0; j < n_segments; ++j) { + less_max = std::max(less_max, row_partial[2 * j]); + eq_max = std::max(eq_max, row_partial[2 * j + 1]); + } + op.write_keys(iter_id, state.k, state.k - state.k_rem, + state.n_nan, less_max, eq_max); + }); + }); + } + else { + last_ev = exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(last_ev); + cgh.use_kernel_bundle(bundle.kb); + + cgh.parallel_for( + ndRange, [=](sycl::nd_item<1> ndit) { + const std::size_t group_id = ndit.get_group(0); + const std::size_t iter_id = group_id / n_segments; + const std::size_t segment_id = + group_id - iter_id * n_segments; + + const StateT state = state_ptr[iter_id]; + + const ValueT *row_arg = arg_ptr + iter_id * n_values; + // the filter wrote out the elements before the candidates + // already + const auto row_out = shifted(op.row_out(iter_id, row_arg), + state.n_less_first); + const std::uint64_t n_less = + state.k - state.n_less_first - state.k_rem; + const std::uint64_t n_ties = + (OpT::output == select_output::whole_row) ? state.n_eq + : state.k_rem; + const auto tie_out = shifted(row_out, n_less); + const auto greater = + greater_out(shifted(row_out, n_less + state.n_eq)); + + if (state.n_cand) { + const std::uint32_t *row_cand = + cand_ptr + iter_id * cand_cap; + const auto [begin, end] = + segment_range(segment_id, n_segments, state.n_cand); + gather_selected( + ndit, CandKeyT{row_arg, row_cand}, begin, end, + state.desired, state.mask, + less_offset_ptr[group_id], less_count_ptr[group_id], + tie_offset_ptr[group_id], n_ties, + via_candidates(row_cand, row_out), + via_candidates(row_cand, tie_out), + via_candidates(row_cand, greater)); + } + else { + const auto [begin, end] = + segment_range(segment_id, n_segments, n_values); + gather_selected( + ndit, RowKeyT{row_arg}, begin, end, state.desired, + state.mask, less_offset_ptr[group_id], + less_count_ptr[group_id], tie_offset_ptr[group_id], + n_ties, row_out, tie_out, greater); + } + }); + }); + } return dpnp::tensor::alloc_utils::async_smart_free( - exec_q, {last_ev}, state_owner, hist_owner, counts_owner, cand_owner); + exec_q, {last_ev}, state_owner, hist_owner, counts_owner, cand_owner, + partial_owner); } //----------------------------------------------------------------------- @@ -1046,14 +1818,12 @@ inline constexpr std::size_t sub_group_max_n = 64; inline constexpr std::uint32_t sub_group_max_chunks = 8; inline constexpr std::uint32_t sub_group_few_chunks = 4; -template +template sycl::event radix_select_dispatch(sycl::queue &exec_q, std::size_t n_iters, std::size_t n_values, - std::size_t k, const ValueT *arg_ptr, - ValueT *vals_ptr, - IndexT *inds_ptr, + const OpT &op, const std::vector &depends) { const auto &dev = exec_q.get_device(); @@ -1078,13 +1848,13 @@ sycl::event radix_select_dispatch(sycl::queue &exec_q, if (n_values <= sub_group_few_chunks * min_sg_size) { return radix_select_sub_group_submit( - exec_q, n_iters, n_values, k, arg_ptr, vals_ptr, inds_ptr, - min_sg_size, wg_size, depends); + exec_q, n_iters, n_values, arg_ptr, op, min_sg_size, wg_size, + depends); } return radix_select_sub_group_submit( - exec_q, n_iters, n_values, k, arg_ptr, vals_ptr, inds_ptr, - min_sg_size, wg_size, depends); + exec_q, n_iters, n_values, arg_ptr, op, min_sg_size, wg_size, + depends); } // enough work-groups to occupy the device @@ -1111,8 +1881,8 @@ sycl::event radix_select_dispatch(sycl::queue &exec_q, // candidates static constexpr bool filter = (sizeof(KeyT) >= 4); return radix_select_multi_group_impl( - exec_q, n_iters, n_values, k, arg_ptr, vals_ptr, inds_ptr, - n_segments, wg_size, filter, depends); + exec_q, n_iters, n_values, arg_ptr, op, n_segments, wg_size, filter, + depends); } // short rows leave most of a large work-group idle std::size_t row_wg_size = 64; @@ -1121,8 +1891,7 @@ sycl::event radix_select_dispatch(sycl::queue &exec_q, } row_wg_size = std::min(row_wg_size, wg_size); return radix_select_one_group_submit( - exec_q, n_iters, n_values, k, arg_ptr, vals_ptr, inds_ptr, row_wg_size, - depends); + exec_q, n_iters, n_values, arg_ptr, op, row_wg_size, depends); } /*! @brief Writes the `k` smallest (largest if not `is_ascending`) elements of @@ -1144,12 +1913,80 @@ sycl::event radix_select_impl(sycl::queue &exec_q, throw std::runtime_error("Invalid value of k for radix select"); } + const TopKOp op{vals_ptr, inds_ptr, k}; if (is_ascending) { return radix_select_dispatch( - exec_q, n_iters, n_values, k, arg_ptr, vals_ptr, inds_ptr, depends); + exec_q, n_iters, n_values, arg_ptr, op, depends); } return radix_select_dispatch( - exec_q, n_iters, n_values, k, arg_ptr, vals_ptr, inds_ptr, depends); + exec_q, n_iters, n_values, arg_ptr, op, depends); +} + +/*! @brief Writes each row of C-contiguous `(n_iters, n_values)` array into + * `(n_iters, n_values)` array rearranged so that its element at `kth` is the + * one a sort would put there, preceded by those ordering before it and + * followed by the others; each of these three groups keeps the order of the + * row, NaNs order last, and -0.0 and +0.0 compare equal */ +template +sycl::event radix_partition_impl(sycl::queue &exec_q, + std::size_t n_iters, + std::size_t n_values, + std::size_t kth, + const ValueT *arg_ptr, + ValueT *vals_ptr, + const std::vector &depends) +{ + if (kth >= n_values) { + throw std::runtime_error("Invalid value of kth for radix partition"); + } + + const PartitionOp op{vals_ptr, n_values, kth + 1}; + return radix_select_dispatch( + exec_q, n_iters, n_values, arg_ptr, op, depends); +} + +/*! @brief Writes the indices by which `radix_partition_impl` rearranges each + * row of C-contiguous `(n_iters, n_values)` array into `(n_iters, n_values)` + * array */ +template +sycl::event radix_argpartition_impl(sycl::queue &exec_q, + std::size_t n_iters, + std::size_t n_values, + std::size_t kth, + const ValueT *arg_ptr, + IndexT *inds_ptr, + const std::vector &depends) +{ + if (kth >= n_values) { + throw std::runtime_error("Invalid value of kth for radix partition"); + } + + const ArgPartitionOp op{inds_ptr, n_values, kth + 1}; + return radix_select_dispatch( + exec_q, n_iters, n_values, arg_ptr, op, depends); +} + +/*! @brief Writes the median of each row of C-contiguous `(n_iters, + * n_values)` array, as `ResT`, into `(n_iters,)` array; it is NaN for a row + * with NaNs, or with NaNs only if `ignore_nan`, which then sets `*all_nan_ptr` + * unless it is null */ +template +sycl::event radix_median_impl(sycl::queue &exec_q, + std::size_t n_iters, + std::size_t n_values, + bool ignore_nan, + const ValueT *arg_ptr, + ResT *out_ptr, + std::uint32_t *all_nan_ptr, + const std::vector &depends) +{ + if (n_values == 0) { + throw std::runtime_error("Median of an empty row is undefined"); + } + + const MedianOp op{out_ptr, all_nan_ptr, n_values, ignore_nan}; + return radix_select_dispatch( + exec_q, n_iters, n_values, arg_ptr, op, depends); } } // namespace dpnp::tensor::kernels::radix_select_details diff --git a/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp b/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp index 7febd8a09ae8..6bcf28ed8d4c 100644 --- a/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp +++ b/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp @@ -297,18 +297,6 @@ struct uint_of_size<8> template using radix_key_t = typename uint_of_size::type; -template -T normalize_signed_zero(T val) -{ - if constexpr (std::is_floating_point_v || - std::is_same_v) { - return (val == T(0)) ? T(0) : val; - } - else { - return val; - } -} - /*! @brief Order-preserving unsigned key of `val`, unlike * `order_preserving_cast` equal exactly when the values compare equal, so that * -0.0 and +0.0 as well as all NaNs share a key */ @@ -320,9 +308,55 @@ radix_key_t ordered_radix_key(const T &val) // by reference, see order_preserving_cast for bool return KeyT(order_preserving_cast(val)); } + else if constexpr (std::is_floating_point_v || + std::is_same_v) { + static constexpr KeyT sign_bit = + KeyT(1) << (number_of_bits_in_type() - 1); + // the keys of -0.0 and +0.0, the descending keys being the inverses of + // the ascending ones + static constexpr KeyT neg_zero_key = + (is_ascending) ? static_cast(~sign_bit) : sign_bit; + static constexpr KeyT pos_zero_key = + (is_ascending) ? sign_bit : static_cast(~sign_bit); + + // compared as keys: -fno-signed-zeros lets a test on the value fold + const KeyT key = KeyT(order_preserving_cast(val)); + return (key == neg_zero_key) ? pos_zero_key : key; + } + else { + return KeyT(order_preserving_cast(val)); + } +} + +/*! @brief Value of type `T` whose `ordered_radix_key` is `key`, NaNs and zeros + * come back as the quiet NaN and +0.0 */ +template +T ordered_radix_value(radix_key_t key) +{ + using KeyT = radix_key_t; + static constexpr KeyT sign_bit = KeyT(1) + << (number_of_bits_in_type() - 1); + + if constexpr (std::is_same_v) { + return (is_ascending) ? (key != 0) : (key == 0); + } else { - return KeyT( - order_preserving_cast(normalize_signed_zero(val))); + // the descending keys are the bitwise inverses of the ascending ones + const KeyT asc_key = (is_ascending) ? key : static_cast(~key); + if constexpr (std::is_integral_v && std::is_unsigned_v) { + return T(asc_key); + } + else if constexpr (std::is_integral_v) { + return sycl::bit_cast(static_cast(asc_key ^ sign_bit)); + } + else { + // non-negative values have the sign bit flipped, negative ones + // all their bits + const KeyT bits = (asc_key & sign_bit) + ? static_cast(asc_key ^ sign_bit) + : static_cast(~asc_key); + return sycl::bit_cast(bits); + } } } diff --git a/dpnp/libtensor/source/CMakeLists.txt b/dpnp/libtensor/source/CMakeLists.txt index f4d203bb893c..e88c8fb2bd47 100644 --- a/dpnp/libtensor/source/CMakeLists.txt +++ b/dpnp/libtensor/source/CMakeLists.txt @@ -72,11 +72,13 @@ set(_dpnp_indexing_sources ${CMAKE_CURRENT_SOURCE_DIR}/choose.cpp ${CMAKE_CURRENT_SOURCE_DIR}/putmask.cpp ) +set(_dpnp_sorting_sources ${CMAKE_CURRENT_SOURCE_DIR}/sorting/partition.cpp) set(_statistics_impl_sources ${CMAKE_CURRENT_SOURCE_DIR}/statistics/bincount.cpp ${CMAKE_CURRENT_SOURCE_DIR}/statistics/histogram.cpp ${CMAKE_CURRENT_SOURCE_DIR}/statistics/histogramdd.cpp ${CMAKE_CURRENT_SOURCE_DIR}/statistics/histogram_common.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/statistics/median.cpp ${CMAKE_CURRENT_SOURCE_DIR}/statistics/sliding_dot_product1d.cpp ${CMAKE_CURRENT_SOURCE_DIR}/statistics/sliding_window1d.cpp ${CMAKE_CURRENT_SOURCE_DIR}/statistics/statistics_py.cpp @@ -88,6 +90,7 @@ set(_window_impl_sources set(_dpnp_sources ${_dpnp_elementwise_sources} ${_dpnp_indexing_sources} + ${_dpnp_sorting_sources} ${_statistics_impl_sources} ${_window_impl_sources} ) @@ -235,6 +238,7 @@ set(_tensor_reductions_impl_sources ${_reduction_sources} ) set(_tensor_sorting_impl_sources + ${_dpnp_sorting_sources} ${CMAKE_CURRENT_SOURCE_DIR}/tensor_sorting.cpp ${_sorting_sources} ) @@ -334,6 +338,8 @@ endforeach() # An inline function with FP math that both instantiate gets the flags of the # copy the linker keeps (the dpnp one, see above), so keep such helpers out of # shared headers or give them the same flags in both groups. +# dpnp sources of the modules with the libtensor radix select kernels +set(_dpnp_radix_select_sources ${_dpnp_sorting_sources} ${_statistics_impl_sources}) foreach(_src_fn ${_dpnp_sources}) get_source_file_property(_cmpl_options_prop ${_src_fn} COMPILE_OPTIONS) set(_combined_options_prop @@ -341,7 +347,7 @@ foreach(_src_fn ${_dpnp_sources}) "${_clang_prefix}-fno-approx-func" "${_clang_prefix}-fno-finite-math-only" ) - if(DPNP_GENERATE_COVERAGE) + if(DPNP_GENERATE_COVERAGE AND NOT _src_fn IN_LIST _dpnp_radix_select_sources) list( APPEND _combined_options_prop -fprofile-instr-generate @@ -358,11 +364,16 @@ endforeach() if(DPNP_GENERATE_COVERAGE) # libtensor code the dpnp functions use, measured like the dpnp sources. # Host only: their device coverage records break sycl-link - # (a source property, unlike target options, keeps the repeated -Xarch_host) + # (a source property, unlike target options, keeps the repeated -Xarch_host). + # The radix select modules are host only too: every kernel image of a module + # gets all its device counters as device globals, and their init events + # crash the OpenCL CPU runtime at random in the kernel bundle launches. + # The device counts never reach the profile anyway set_property( SOURCE ${_static_lib_sources} ${CMAKE_CURRENT_SOURCE_DIR}/elementwise_functions/elementwise_functions_type_utils.cpp + ${_dpnp_radix_select_sources} APPEND PROPERTY COMPILE_OPTIONS diff --git a/dpnp/libtensor/source/sorting/partition.cpp b/dpnp/libtensor/source/sorting/partition.cpp new file mode 100644 index 000000000000..121bfc9b55eb --- /dev/null +++ b/dpnp/libtensor/source/sorting/partition.cpp @@ -0,0 +1,268 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** + +#include +#include +#include +#include + +#include + +#include +#include + +#include "dpnp4pybind11.hpp" + +#include "partition.hpp" + +// dpnp tensor headers +#include "kernels/sorting/radix_select.hpp" +#include "utils/output_validation.hpp" +#include "utils/type_dispatch.hpp" +#include "utils/type_utils.hpp" + +// utils extension headers +#include "ext/common.hpp" +#include "ext/validation_utils.hpp" + +namespace py = pybind11; +namespace td_ns = dpnp::tensor::type_dispatch; + +using dpnp::tensor::usm_ndarray; + +using ext::validation::array_names; +using ext::validation::check_c_contig; +using ext::validation::check_has_dtype; +using ext::validation::check_no_overlap; +using ext::validation::check_num_dims; +using ext::validation::check_queue; +using ext::validation::check_same_dtype; +using ext::validation::check_writable; + +namespace dpnp::tensor::py_internal +{ +using ext::common::init_dispatch_vector; + +typedef sycl::event (*partition_fn_ptr_t)(sycl::queue &, + const std::size_t, // n_iters + const std::size_t, // n_values + const std::size_t, // kth + const char *, // src + char *, // dst + const std::vector &); + +static partition_fn_ptr_t partition_dispatch_vector[td_ns::num_types]; +static partition_fn_ptr_t argpartition_dispatch_vector[td_ns::num_types]; + +namespace +{ +namespace rs_ns = dpnp::tensor::kernels::radix_select_details; + +// types with an order-preserving radix key +template +constexpr bool use_radix_select = !dpnp::tensor::type_utils::is_complex_v; + +template +sycl::event partition_call(sycl::queue &exec_q, + const std::size_t n_iters, + const std::size_t n_values, + const std::size_t kth, + const char *src_p, + char *dst_p, + const std::vector &depends) +{ + return rs_ns::radix_partition_impl( + exec_q, n_iters, n_values, kth, reinterpret_cast(src_p), + reinterpret_cast(dst_p), depends); +} + +template +sycl::event argpartition_call(sycl::queue &exec_q, + const std::size_t n_iters, + const std::size_t n_values, + const std::size_t kth, + const char *src_p, + char *dst_p, + const std::vector &depends) +{ + using IndexT = std::int64_t; + return rs_ns::radix_argpartition_impl( + exec_q, n_iters, n_values, kth, reinterpret_cast(src_p), + reinterpret_cast(dst_p), depends); +} + +/*! @brief Partitions the rows of `src` along its last axis into `dst` of the + * same shape, the implementation being taken from `dispatch_vector` */ +std::pair + partition_rows(const partition_fn_ptr_t *dispatch_vector, + const usm_ndarray &src, + const std::size_t kth, + const usm_ndarray &dst, + sycl::queue &exec_q, + const std::vector &depends) +{ + array_names names = {{&src, "src"}, {&dst, "dst"}}; + + const int nd = src.get_ndim(); + if (nd == 0) { + throw py::value_error("The input array must have at least one " + "dimension"); + } + check_num_dims(&dst, nd, names); + + const py::ssize_t *src_shape = src.get_shape_raw(); + const py::ssize_t *dst_shape = dst.get_shape_raw(); + std::size_t n_iters(1); + for (int i = 0; i < nd; ++i) { + if (src_shape[i] != dst_shape[i]) { + throw py::value_error("`src` and `dst` shapes must match"); + } + if (i + 1 < nd) { + n_iters *= static_cast(src_shape[i]); + } + } + const std::size_t n_values = static_cast(src_shape[nd - 1]); + + check_queue({&src, &dst}, names, exec_q); + check_writable({&dst}, names); + + if (kth >= n_values) { + throw py::value_error("`kth` is out of bounds for the last axis"); + } + + check_c_contig({&src, &dst}, names); + check_no_overlap(&src, &dst, names); + dpnp::tensor::validation::AmpleMemory::throw_if_not_ample( + dst, n_iters * n_values); + + auto types = td_ns::usm_ndarray_types(); + const int src_typeid = types.typenum_to_lookup_id(src.get_typenum()); + + auto fn = dispatch_vector[src_typeid]; + if (fn == nullptr) { + throw py::value_error("No partition implementation is available for " + "the data type of the input array"); + } + + if (n_iters == 0) { + return {sycl::event(), sycl::event()}; + } + + sycl::event comp_ev = fn(exec_q, n_iters, n_values, kth, src.get_data(), + dst.get_data(), depends); + sycl::event ht_ev = + dpnp::utils::keep_args_alive(exec_q, {src, dst}, {comp_ev}); + + return std::make_pair(ht_ev, comp_ev); +} +} // namespace + +std::pair + py_partition(const usm_ndarray &src, + const std::size_t kth, + const usm_ndarray &dst, + sycl::queue &exec_q, + const std::vector &depends) +{ + array_names names = {{&src, "src"}, {&dst, "dst"}}; + check_same_dtype(&src, &dst, names); + + return partition_rows(partition_dispatch_vector, src, kth, dst, exec_q, + depends); +} + +std::pair + py_argpartition(const usm_ndarray &src, + const std::size_t kth, + const usm_ndarray &dst, + sycl::queue &exec_q, + const std::vector &depends) +{ + array_names names = {{&src, "src"}, {&dst, "dst"}}; + check_has_dtype(&dst, td_ns::typenum_t::INT64, names); + + return partition_rows(argpartition_dispatch_vector, src, kth, dst, exec_q, + depends); +} + +template +struct PartitionFactory +{ + fnT get() + { + if constexpr (use_radix_select) { + return partition_call; + } + else { + return nullptr; + } + } +}; + +template +struct ArgPartitionFactory +{ + fnT get() + { + if constexpr (use_radix_select) { + return argpartition_call; + } + else { + return nullptr; + } + } +}; + +static void populate_partition_dispatch_vectors() +{ + init_dispatch_vector( + partition_dispatch_vector); + init_dispatch_vector( + argpartition_dispatch_vector); +} + +void init_partition(py::module_ &m) +{ + populate_partition_dispatch_vectors(); + + m.def("_partition", &py_partition, + "Partitions each row of C-contiguous `src` along its last axis " + "about the element at `kth` of the sorted row, writing into `dst` " + "of the same shape and data type", + py::arg("src"), py::arg("kth"), py::arg("dst"), py::arg("sycl_queue"), + py::arg("depends") = py::list()); + + m.def("_argpartition", &py_argpartition, + "Writes into int64 `dst` of the shape of C-contiguous `src` the " + "indices which partition each row of `src` along its last axis " + "about the element at `kth` of the sorted row", + py::arg("src"), py::arg("kth"), py::arg("dst"), py::arg("sycl_queue"), + py::arg("depends") = py::list()); +} + +} // namespace dpnp::tensor::py_internal diff --git a/dpnp/libtensor/source/sorting/partition.hpp b/dpnp/libtensor/source/sorting/partition.hpp new file mode 100644 index 000000000000..7fb7f25c826e --- /dev/null +++ b/dpnp/libtensor/source/sorting/partition.hpp @@ -0,0 +1,42 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** +// +// This file defines functions of dpnp.libtensor._tensor_sorting_impl extension +// +//***************************************************************************** + +#pragma once + +#include + +namespace py = pybind11; + +namespace dpnp::tensor::py_internal +{ +void init_partition(py::module_ &m); +} // namespace dpnp::tensor::py_internal diff --git a/dpnp/libtensor/source/statistics/median.cpp b/dpnp/libtensor/source/statistics/median.cpp new file mode 100644 index 000000000000..a5ed7cf6ed3e --- /dev/null +++ b/dpnp/libtensor/source/statistics/median.cpp @@ -0,0 +1,251 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include "dpnp4pybind11.hpp" + +#include "median.hpp" + +// dpnp tensor headers +#include "kernels/sorting/radix_select.hpp" +#include "utils/output_validation.hpp" +#include "utils/type_dispatch.hpp" +#include "utils/type_utils.hpp" + +// utils extension headers +#include "ext/common.hpp" +#include "ext/validation_utils.hpp" + +namespace td_ns = dpnp::tensor::type_dispatch; + +using dpnp::tensor::usm_ndarray; + +using ext::validation::array_names; +using ext::validation::check_c_contig; +using ext::validation::check_no_overlap; +using ext::validation::check_num_dims; +using ext::validation::check_queue; +using ext::validation::check_writable; + +namespace statistics::median +{ +using ext::common::init_dispatch_table; + +typedef sycl::event (*median_fn_ptr_t)(sycl::queue &, + const std::size_t, // n_iters + const std::size_t, // n_values + const bool, // ignore_nan + const char *, // src + char *, // dst + std::uint32_t *, // all_nan + const std::vector &); + +static median_fn_ptr_t median_dispatch_table[td_ns::num_types] + [td_ns::num_types]; + +namespace +{ +template +sycl::event median_call(sycl::queue &exec_q, + const std::size_t n_iters, + const std::size_t n_values, + const bool ignore_nan, + const char *src_p, + char *dst_p, + std::uint32_t *all_nan_p, + const std::vector &depends) +{ + using dpnp::tensor::kernels::radix_select_details::radix_median_impl; + return radix_median_impl(exec_q, n_iters, n_values, ignore_nan, + reinterpret_cast(src_p), + reinterpret_cast(dst_p), + all_nan_p, depends); +} +} // namespace + +/*! @brief Writes the median of each row of C-contiguous `src` along its last + * axis into C-contiguous `dst` of the shape of `src` without its last axis; + * the single element of `all_nan`, when it is given, is set to whether + * `ignore_nan` and a row has NaNs only */ +std::pair + py_median(const usm_ndarray &src, + const bool ignore_nan, + const usm_ndarray &dst, + sycl::queue &exec_q, + const std::optional &all_nan, + const std::vector &depends) +{ + array_names names = {{&src, "src"}, {&dst, "dst"}}; + + const int nd = src.get_ndim(); + if (nd == 0) { + throw py::value_error("The input array must have at least one " + "dimension"); + } + check_num_dims(&dst, nd - 1, names); + + const py::ssize_t *src_shape = src.get_shape_raw(); + const py::ssize_t *dst_shape = dst.get_shape_raw(); + std::size_t n_iters(1); + for (int i = 0; i + 1 < nd; ++i) { + if (src_shape[i] != dst_shape[i]) { + throw py::value_error("`dst` shape must match `src` shape " + "without its last axis"); + } + n_iters *= static_cast(src_shape[i]); + } + const std::size_t n_values = static_cast(src_shape[nd - 1]); + + std::uint32_t *all_nan_p = nullptr; + if (all_nan.has_value()) { + const usm_ndarray &flag = all_nan.value(); + names.emplace(&flag, "all_nan"); + check_queue({&src, &dst, &flag}, names, exec_q); + check_writable({&dst, &flag}, names); + auto types = td_ns::usm_ndarray_types(); + if (flag.get_size() != 1 || + types.typenum_to_lookup_id(flag.get_typenum()) != + static_cast(td_ns::typenum_t::UINT32)) { + throw py::value_error("`all_nan` must be a single uint32 element"); + } + check_no_overlap({&src, &dst}, {&flag}, names); + all_nan_p = reinterpret_cast(flag.get_data()); + } + else { + check_queue({&src, &dst}, names, exec_q); + check_writable({&dst}, names); + } + + if (n_iters != 0 && n_values == 0) { + throw py::value_error("The median of an empty axis is undefined"); + } + + check_c_contig({&src, &dst}, names); + check_no_overlap(&src, &dst, names); + dpnp::tensor::validation::AmpleMemory::throw_if_not_ample(dst, n_iters); + + auto types = td_ns::usm_ndarray_types(); + const int src_typeid = types.typenum_to_lookup_id(src.get_typenum()); + const int dst_typeid = types.typenum_to_lookup_id(dst.get_typenum()); + + auto fn = median_dispatch_table[src_typeid][dst_typeid]; + if (fn == nullptr) { + throw py::value_error("No median implementation is available for " + "the data types of the input and the output " + "arrays"); + } + + std::vector median_deps(depends); + if (all_nan_p != nullptr) { + median_deps.push_back( + exec_q.fill(all_nan_p, 0, 1, depends)); + } + if (n_iters == 0) { + // no row has NaNs only + if (all_nan_p == nullptr) { + return {sycl::event(), sycl::event()}; + } + sycl::event fill_ev = median_deps.back(); + sycl::event ht_ev = + dpnp::utils::keep_args_alive(exec_q, {all_nan.value()}, {fill_ev}); + return std::make_pair(ht_ev, fill_ev); + } + + sycl::event comp_ev = + fn(exec_q, n_iters, n_values, ignore_nan, src.get_data(), + dst.get_data(), all_nan_p, median_deps); + sycl::event ht_ev = + (all_nan.has_value()) + ? dpnp::utils::keep_args_alive(exec_q, {src, dst, all_nan.value()}, + {comp_ev}) + : dpnp::utils::keep_args_alive(exec_q, {src, dst}, {comp_ev}); + + return std::make_pair(ht_ev, comp_ev); +} + +/*! @brief Whether the median of `T` is written as `ResT`: floating types + * keep their type, the others are averaged in double or, on devices without + * it, in single precision */ +template +constexpr bool is_median_type_pair() +{ + using dpnp::tensor::type_utils::is_complex_v; + // complex types have no order-preserving radix key + if constexpr (is_complex_v || is_complex_v) { + return false; + } + else if constexpr (std::is_floating_point_v || + std::is_same_v) { + return std::is_same_v; + } + else { + return std::is_same_v || std::is_same_v; + } +} + +template +struct MedianFactory +{ + fnT get() + { + if constexpr (is_median_type_pair()) { + return median_call; + } + else { + return nullptr; + } + } +}; + +void populate_median(py::module_ m) +{ + init_dispatch_table(median_dispatch_table); + + m.def("_median", &py_median, + "Writes the median of each row of C-contiguous `src` along its last " + "axis into `dst` of the shape of `src` without its last axis; it is " + "NaN for a row with NaNs, or, if `ignore_nan`, with NaNs only; the " + "single element of `all_nan`, if given, is set to whether any row " + "is of the latter", + py::arg("src"), py::arg("ignore_nan"), py::arg("dst"), + py::arg("sycl_queue"), py::arg("all_nan") = py::none(), + py::arg("depends") = py::list()); +} + +} // namespace statistics::median diff --git a/dpnp/libtensor/source/statistics/median.hpp b/dpnp/libtensor/source/statistics/median.hpp new file mode 100644 index 000000000000..f0deb56d8a3b --- /dev/null +++ b/dpnp/libtensor/source/statistics/median.hpp @@ -0,0 +1,38 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** + +#pragma once + +#include + +namespace py = pybind11; + +namespace statistics::median +{ +void populate_median(py::module_ m); +} // namespace statistics::median diff --git a/dpnp/libtensor/source/statistics/statistics_py.cpp b/dpnp/libtensor/source/statistics/statistics_py.cpp index f8492d1fb070..a0bbe62adf58 100644 --- a/dpnp/libtensor/source/statistics/statistics_py.cpp +++ b/dpnp/libtensor/source/statistics/statistics_py.cpp @@ -35,6 +35,7 @@ #include "bincount.hpp" #include "histogram.hpp" #include "histogramdd.hpp" +#include "median.hpp" #include "sliding_dot_product1d.hpp" PYBIND11_MODULE(_statistics_impl, m, py::mod_gil_not_used()) @@ -43,4 +44,5 @@ PYBIND11_MODULE(_statistics_impl, m, py::mod_gil_not_used()) statistics::histogram::populate_histogram(m); statistics::sliding_window1d::populate_sliding_dot_product1d(m); statistics::histogram::populate_histogramdd(m); + statistics::median::populate_median(m); } diff --git a/dpnp/libtensor/source/tensor_sorting.cpp b/dpnp/libtensor/source/tensor_sorting.cpp index ec6646f74c6d..6306fd5e5095 100644 --- a/dpnp/libtensor/source/tensor_sorting.cpp +++ b/dpnp/libtensor/source/tensor_sorting.cpp @@ -38,6 +38,7 @@ #include "sorting/isin.hpp" #include "sorting/merge_argsort.hpp" #include "sorting/merge_sort.hpp" +#include "sorting/partition.hpp" #include "sorting/radix_argsort.hpp" #include "sorting/radix_sort.hpp" #include "sorting/searchsorted.hpp" @@ -52,4 +53,5 @@ PYBIND11_MODULE(_tensor_sorting_impl, m, py::mod_gil_not_used()) dpnp::tensor::py_internal::init_radix_sort_functions(m); dpnp::tensor::py_internal::init_radix_argsort_functions(m); dpnp::tensor::py_internal::init_topk_functions(m); + dpnp::tensor::py_internal::init_partition(m); } diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py index 9eab8cfc198f..b80c3a6739be 100644 --- a/dpnp/tests/test_indexing.py +++ b/dpnp/tests/test_indexing.py @@ -1089,8 +1089,7 @@ class TestTakeAlongAxis: pytest.param( _add_keepdims(dpnp.max), _add_keepdims(dpnp.argmax), {} ), - # TODO: unmute, once `dpnp.argpartition` is implemented - # pytest.param(dpnp.partition, dpnp.argpartition, {"kth": 2}), + pytest.param(dpnp.partition, dpnp.argpartition, {"kth": 2}), ], ) def test_argequivalent(self, func, argfunc, kwargs): diff --git a/dpnp/tests/test_nanfunctions.py b/dpnp/tests/test_nanfunctions.py index e9f2d6c7a17e..be97422de4f1 100644 --- a/dpnp/tests/test_nanfunctions.py +++ b/dpnp/tests/test_nanfunctions.py @@ -493,8 +493,8 @@ def test_overwrite_input(self, axis): ib = ia.copy() expected = numpy.nanmedian(b, axis=axis, overwrite_input=True) result = dpnp.nanmedian(ib, axis=axis, overwrite_input=True) + # overwrite_input only allows the input to be modified assert not numpy.all(a == b) - assert not dpnp.all(ia == ib) assert_dtype_allclose(result, expected) @@ -526,6 +526,21 @@ def test_out(self, dtype, axis, out_shape): assert result is out_dp assert_dtype_allclose(result, expected) + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + @pytest.mark.parametrize("axis", [None, 1, (0, 2)]) + def test_out_keepdims(self, dtype, axis): + a = generate_random_numpy_array((2, 3, 4), dtype=dtype) + a[0, 1, ::2] = a[1, 2, 1] = numpy.nan + ia = dpnp.array(a) + out_shape = numpy.nanmedian(a, axis=axis, keepdims=True).shape + + out_np = numpy.empty_like(a, shape=out_shape) + out_dp = dpt.empty(out_shape, dtype=dtype) + expected = numpy.nanmedian(a, axis=axis, keepdims=True, out=out_np) + result = dpnp.nanmedian(ia, axis=axis, keepdims=True, out=out_dp) + assert result.get_array() is out_dp + assert_dtype_allclose(result, expected) + def test_error(self): a = dpnp.arange(6.0).reshape(2, 3) a[0, 0] = a[-1, -1] = numpy.nan @@ -542,6 +557,94 @@ def test_error(self): dpnp.nanmedian(a, axis=1, out=res) +class TestNanMedianSelect: + # row lengths reaching the sub-group, the work-group and the multi-group + # kernels of radix select + @pytest.mark.parametrize( + "n", [1, 2, 3, 63, 64, 65, 1000, 1001, 2**16 + 2, 2**17 + 1] + ) + @pytest.mark.parametrize("dtype", get_float_dtypes(no_float16=False)) + def test_basic(self, n, dtype): + rnd = numpy.random.RandomState(81) + a = rnd.standard_normal(size=(5, n)).astype(dtype) + # a varying number of NaNs, to give both even and odd valid counts + for i in range(1, 4): + a[i, rnd.choice(n, min(i * n // 4 + i, n - 1), replace=False)] = ( + numpy.nan + ) + # signed zeros and NaNs, never all NaN (see test_all_nan_rows) + a[4, :] = -0.0 + a[4, 1::3] = numpy.nan + ia = dpnp.array(a) + + expected = numpy.nanmedian(a, axis=-1) + result = dpnp.nanmedian(ia, axis=-1) + assert_dtype_allclose(result, expected) + + # NumPy's mean gives +0.0 for zeros of either sign + @pytest.mark.parametrize("n", [2, 3, 64, 1000, 2**16 + 1]) + @pytest.mark.parametrize("dtype", get_float_dtypes(no_float16=False)) + def test_zero_sign(self, n, dtype): + a = numpy.full((3, n), -0.0, dtype=dtype) + a[:, 1::4] = numpy.nan + a[1, ::3] = 0.0 + a[2, : n // 4] = -1 + ia = dpnp.array(a) + + expected = numpy.nanmedian(a, axis=-1) + result = dpnp.nanmedian(ia, axis=-1) + assert_dtype_allclose(result, expected) + assert_array_equal(dpnp.signbit(result), numpy.signbit(expected)) + + @pytest.mark.parametrize("n", [3, 1000, 2**16 + 1]) + def test_all_nan_rows(self, n): + a = generate_random_numpy_array((3, n), dtype="f4") + a[1] = numpy.nan + ia = dpnp.array(a) + + with pytest.warns(RuntimeWarning, match="All-NaN slice"): + result = dpnp.nanmedian(ia, axis=-1) + with pytest.warns(RuntimeWarning, match="All-NaN slice"): + expected = numpy.nanmedian(a, axis=-1) + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize( + "dtype", get_all_dtypes(no_none=True, no_float16=False, no_complex=True) + ) + def test_no_warning(self, dtype, recwarn): + a = generate_random_numpy_array((3, 100), dtype=dtype) + ia = dpnp.array(a) + + result = dpnp.nanmedian(ia, axis=-1) + assert_dtype_allclose(result, numpy.nanmedian(a, axis=-1)) + assert not [w for w in recwarn if "All-NaN" in str(w.message)] + + @pytest.mark.parametrize("keepdims", [False, True]) + def test_out(self, keepdims): + a = generate_random_numpy_array((4, 3)) + a[1, 2] = numpy.nan + ia = dpnp.array(a) + shape = (1, 3) if keepdims else (3,) + + out_np = numpy.empty_like(a, shape=shape) + out_dp = dpnp.empty_like(ia, shape=shape) + expected = numpy.nanmedian(a, axis=0, keepdims=keepdims, out=out_np) + result = dpnp.nanmedian(ia, axis=0, keepdims=keepdims, out=out_dp) + assert result is out_dp + assert_dtype_allclose(result, expected) + + # NumPy broadcasts the result into these, dpnp requires the exact shape + @pytest.mark.parametrize( + "axis, out_shape", [(0, (2, 3)), (0, (1, 3)), (None, (3,))] + ) + def test_out_wrong_shape(self, axis, out_shape): + a = dpnp.arange(12.0).reshape(4, 3) + a[1, 2] = dpnp.nan + out = dpnp.empty(out_shape) + with pytest.raises(ValueError, match="Output array of shape"): + dpnp.nanmedian(a, axis=axis, out=out) + + @pytest.mark.parametrize("func", ["nanprod", "nansum"]) class TestNanProdSum: @pytest.mark.parametrize("axis", [None, 0, 1, -1, 2, -2, (1, 2), (0, -2)]) diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py index 48bd2733074d..b730818fe6c3 100644 --- a/dpnp/tests/test_sort.py +++ b/dpnp/tests/test_sort.py @@ -3,6 +3,8 @@ from numpy.testing import assert_array_equal, assert_equal, assert_raises import dpnp +import dpnp.libtensor._tensor_sorting_impl as sorting_ext +import dpnp.tensor as dpt from dpnp.exceptions import AxisError from .helper import ( @@ -275,6 +277,233 @@ def test_v_scalar(self): assert_equal(result, expected) +def _check_partitioned(a, p, kth): + """Check that `p` is `a` partitioned along the last axis at `kth`.""" + a_sorted = numpy.sort(a, axis=-1) + assert_array_equal(numpy.sort(p, axis=-1), a_sorted) + + pk = p[..., kth : kth + 1] + assert_array_equal(pk, a_sorted[..., kth : kth + 1]) + if numpy.issubdtype(a.dtype, numpy.inexact): + # NaNs order last + assert ((p[..., :kth] <= pk) | numpy.isnan(pk)).all() + assert ( + (p[..., kth + 1 :] >= pk) | numpy.isnan(p[..., kth + 1 :]) + ).all() + else: + assert (p[..., :kth] <= pk).all() + assert (p[..., kth + 1 :] >= pk).all() + + +def _check_argpartitioned(a, idx, kth): + """Check that `idx` indexes `a` partitioned along the last axis at `kth`.""" + assert idx.dtype == numpy.int64 + n = a.shape[-1] + assert_array_equal( + numpy.sort(idx, axis=-1), numpy.broadcast_to(numpy.arange(n), a.shape) + ) + _check_partitioned(a, numpy.take_along_axis(a, idx, axis=-1), kth) + + +def _generate_select_data(shape, dtype, kind, seed=81): + """Generate data exercising the radix select kernels.""" + rnd = numpy.random.RandomState(seed) + dtype = numpy.dtype(dtype) + if dtype == numpy.bool_: + return rnd.randint(0, 2, size=shape).astype(dtype) + + if kind == "duplicates": + return rnd.randint(0, 4, size=shape).astype(dtype) + + if numpy.issubdtype(dtype, numpy.integer): + info = numpy.iinfo(dtype) + a = rnd.randint(info.min, info.max, size=shape, dtype=dtype) + if kind == "extremes": + a.flat[:: max(a.size // 7, 1)] = info.min + a.flat[1 :: max(a.size // 5, 1)] = info.max + return a + + a = rnd.standard_normal(size=shape).astype(dtype) + if kind == "extremes": + info = numpy.finfo(dtype) + a.flat[:: max(a.size // 7, 1)] = numpy.nan + a.flat[1 :: max(a.size // 5, 1)] = -0.0 + a.flat[2 :: max(a.size // 5, 1)] = 0.0 + a.flat[3 :: max(a.size // 11, 1)] = numpy.inf + a.flat[4 :: max(a.size // 11, 1)] = -numpy.inf + a.flat[5 :: max(a.size // 13, 1)] = info.max + a.flat[6 :: max(a.size // 13, 1)] = info.tiny / 2 + return a + + +# row lengths reaching the sub-group, the work-group and the multi-group +# kernels of radix select +_select_shapes = [ + (1,), + (2,), + (7,), + (4, 33), + (5, 64), + (3, 65), + (6, 1000), + (2, 5001), + (1, 2**16 + 3), + (3, 2**17 + 1), +] + + +class TestPartitionSelect: + @pytest.mark.parametrize("shape", _select_shapes) + @pytest.mark.parametrize("kind", ["random", "duplicates", "extremes"]) + @pytest.mark.parametrize( + "dtype", get_all_dtypes(no_none=True, no_complex=True, no_float16=False) + ) + def test_partition(self, shape, kind, dtype): + a = _generate_select_data(shape, dtype, kind) + ia = dpnp.array(a) + + n = shape[-1] + for kth in {0, n // 3, n // 2, n - 1}: + _check_partitioned(a, dpnp.asnumpy(dpnp.partition(ia, kth)), kth) + + @pytest.mark.parametrize("shape", _select_shapes) + @pytest.mark.parametrize("kind", ["random", "duplicates", "extremes"]) + @pytest.mark.parametrize( + "dtype", get_all_dtypes(no_none=True, no_complex=True, no_float16=False) + ) + def test_argpartition(self, shape, kind, dtype): + a = _generate_select_data(shape, dtype, kind) + ia = dpnp.array(a) + + n = shape[-1] + for kth in {0, n // 3, n // 2, n - 1}: + idx = dpnp.argpartition(ia, kth) + _check_argpartitioned(a, dpnp.asnumpy(idx), kth) + + @pytest.mark.parametrize("n", [50, 3000, 2**16 + 5]) + def test_signed_zeros_kept(self, n): + # -0.0 and +0.0 compare equal, but each element keeps its sign + a = numpy.zeros(n) + a[::3] = -0.0 + a[1::7] = 1.0 + ia = dpnp.array(a) + + p = dpnp.asnumpy(dpnp.partition(ia, n // 2)) + assert numpy.signbit(p).sum() == numpy.signbit(a).sum() + + idx = dpnp.asnumpy(dpnp.argpartition(ia, n // 2)) + assert_array_equal(numpy.sort(idx), numpy.arange(n)) + + @pytest.mark.parametrize("axis", [0, 1, 2, -1, None]) + def test_non_contiguous(self, axis): + a = generate_random_numpy_array((70, 3, 90), dtype="f4") + ia = dpnp.array(a)[::-1, :, ::2] + a = a[::-1, :, ::2] + + kth = 1 + p = dpnp.asnumpy(dpnp.partition(ia, kth, axis=axis)) + idx = dpnp.asnumpy(dpnp.argpartition(ia, kth, axis=axis)) + if axis is None: + a = a.ravel() + axis = -1 + _check_partitioned( + numpy.moveaxis(a, axis, -1), numpy.moveaxis(p, axis, -1), kth + ) + _check_argpartitioned( + numpy.moveaxis(a, axis, -1), numpy.moveaxis(idx, axis, -1), kth + ) + + def test_matches_argpartition(self): + a = _generate_select_data((4, 3000), numpy.float32, "duplicates") + ia = dpnp.array(a) + + p = dpnp.partition(ia, 1234) + idx = dpnp.argpartition(ia, 1234) + assert_array_equal(p, dpnp.take_along_axis(ia, idx, axis=-1)) + + # the extension validates an empty batch like a non-empty one + @pytest.mark.parametrize("fn", ["_partition", "_argpartition"]) + @pytest.mark.parametrize("shape", [(0, 5), (3, 0)]) + def test_ext_empty_kth_error(self, fn, shape): + src = dpt.empty(shape, dtype="i4") + dst = dpt.empty(shape, dtype="i4" if fn == "_partition" else "i8") + with pytest.raises(ValueError, match="out of bounds"): + getattr(sorting_ext, fn)(src, 5, dst, sycl_queue=src.sycl_queue) + + @pytest.mark.parametrize("fn", ["_partition", "_argpartition"]) + def test_ext_empty_dtype_error(self, fn): + src = dpt.empty((0, 5), dtype="c8") + dst = dpt.empty((0, 5), dtype="c8" if fn == "_partition" else "i8") + with pytest.raises(ValueError, match="No partition implementation"): + getattr(sorting_ext, fn)(src, 1, dst, sycl_queue=src.sycl_queue) + + +class TestArgpartition: + @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True)) + @pytest.mark.parametrize("kth", [0, 3, -1, (1, 5), [2, 4, 8]]) + def test_basic(self, dtype, kth): + a = generate_random_numpy_array(10, dtype) + ia = dpnp.array(a) + + idx = dpnp.argpartition(ia, kth) + p = numpy.take(a, dpnp.asnumpy(idx)) + expected = numpy.partition(a, kth) + for k in numpy.atleast_1d(kth) % a.size: + assert_array_equal(p[k], expected[k]) + assert (p[:k] <= p[k]).all() + assert (p[k + 1 :] >= p[k]).all() + + @pytest.mark.parametrize("axis", [0, 1, -1, None]) + def test_axis(self, axis): + a = generate_random_numpy_array((4, 6), dtype="i4") + ia = dpnp.array(a) + + kth = 2 + result = dpnp.argpartition(ia, kth, axis=axis) + expected = numpy.partition(a, kth, axis=axis) + assert_array_equal( + dpnp.take_along_axis(ia, result, axis=axis)[..., kth], + expected[..., kth], + ) + + def test_method(self): + a = numpy.array([3, 4, 2, 1]) + ia = dpnp.array(a) + + result = ia.argpartition(1) + assert_array_equal(ia[result][1], numpy.partition(a, 1)[1]) + + @pytest.mark.parametrize("axis", list(range(-4, 4)) + [None]) + def test_empty_array(self, axis): + a = numpy.empty((3, 2, 1, 0)) + ia = dpnp.array(a) + + result = dpnp.argpartition(ia, 0, axis=axis) + expected = numpy.argpartition(a, 0, axis=axis) + assert_equal(result, expected) + + def test_empty_kth(self): + a = numpy.array([6, 0, 4, 2]) + ia = dpnp.array(a) + + result = dpnp.argpartition(ia, []) + assert_array_equal(ia[result], dpnp.sort(ia)) + + @pytest.mark.parametrize("xp", [dpnp, numpy]) + def test_kth_errors(self, xp): + a = xp.arange(10) + assert_raises(ValueError, a.argpartition, 10) + assert_raises(ValueError, a.argpartition, -11) + assert_raises(TypeError, a.argpartition, 9.0) + assert_raises(TypeError, a.argpartition, [1, 7.0]) + assert_raises(AxisError, xp.argpartition, a, 2, axis=1) + + def test_not_implemented_kwargs(self): + a = dpnp.arange(10) + assert_raises(NotImplementedError, a.argpartition, 2, kind="nonsense") + assert_raises(NotImplementedError, a.argpartition, 2, order=[]) + + class TestPartition: @pytest.mark.parametrize("data", [[2, 1], [1, 2], [1, 1]]) @pytest.mark.parametrize("kth", [0, 1]) diff --git a/dpnp/tests/test_statistics.py b/dpnp/tests/test_statistics.py index 3d5387d3de39..37f95c0910f4 100644 --- a/dpnp/tests/test_statistics.py +++ b/dpnp/tests/test_statistics.py @@ -9,6 +9,7 @@ ) import dpnp +import dpnp.libtensor._statistics_impl as statistics_ext import dpnp.tensor as dpt from .helper import ( @@ -970,6 +971,27 @@ def test_out(self, dtype, axis, out_shape): assert result is out_dp assert_dtype_allclose(result, expected) + @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True)) + @pytest.mark.parametrize("axis", [None, 1, (0, 2)]) + def test_out_keepdims(self, dtype, axis): + a = generate_random_numpy_array((2, 3, 4), dtype=dtype) + ia = dpnp.array(a) + out_shape = numpy.median(a, axis=axis, keepdims=True).shape + + out_np = numpy.empty_like(a, shape=out_shape) + out_dp = dpnp.empty_like(ia, shape=out_shape) + expected = numpy.median(a, axis=axis, keepdims=True, out=out_np) + result = dpnp.median(ia, axis=axis, keepdims=True, out=out_dp) + assert result is out_dp + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize("dtype", [dpnp.int32, dpnp.complex64]) + def test_out_keepdims_wrong_shape(self, dtype): + ia = dpnp.ones((4, 3), dtype=dtype) + out = dpnp.empty(3, dtype=dtype) + with pytest.raises(ValueError, match="Output array of shape"): + dpnp.median(ia, axis=0, keepdims=True, out=out) + def test_0d_array(self): a = numpy.array(20) ia = dpnp.array(a) @@ -1004,8 +1026,8 @@ def test_overwrite_input(self, axis, keepdims): result = dpnp.median( ib, axis=axis, keepdims=keepdims, overwrite_input=True ) + # overwrite_input only allows the input to be modified assert not numpy.all(a == b) - assert not dpnp.all(ia == ib) assert_dtype_allclose(result, expected) @@ -1020,6 +1042,131 @@ def test_usm_ndarray(self, axis, overwrite_input): assert_dtype_allclose(result, expected) +class TestMedianSelect: + # row lengths reaching the sub-group, the work-group and the multi-group + # kernels of radix select + @pytest.mark.parametrize( + "shape", + [(1,), (2,), (4, 3), (5, 64), (3, 65), (6, 1000), (2, 1001)] + + [(1, 2**16 + 3), (3, 2**17), (2, 2**17 + 1)], + ) + @pytest.mark.parametrize("duplicates", [False, True]) + @pytest.mark.parametrize( + "dtype", get_all_dtypes(no_none=True, no_complex=True, no_float16=False) + ) + def test_basic(self, shape, duplicates, dtype): + rnd = numpy.random.RandomState(81) + if duplicates or dtype == dpnp.bool: + a = rnd.randint(0, 4, size=shape).astype(dtype) + elif dpnp.issubdtype(dtype, dpnp.integer): + # the median of two large integers must not overflow + info = numpy.iinfo(dtype) + a = rnd.randint(info.min, info.max, size=shape, dtype=dtype) + else: + a = rnd.standard_normal(size=shape).astype(dtype) + ia = dpnp.array(a) + + expected = numpy.median(a, axis=-1) + result = dpnp.median(ia, axis=-1) + assert_dtype_allclose(result, expected) + + @pytest.mark.parametrize("n", [5, 64, 1000, 2**16 + 1]) + @pytest.mark.parametrize("dtype", get_float_dtypes(no_float16=False)) + def test_special_values(self, n, dtype): + a = numpy.zeros((4, n), dtype=dtype) + a[:, ::3] = -0.0 + a[1, ::5] = numpy.inf + a[2, ::4] = -numpy.inf + a[3, : n // 2] = numpy.finfo(dtype).max + # a single NaN makes the median NaN + a[0, n // 3] = numpy.nan + ia = dpnp.array(a) + + expected = numpy.median(a, axis=-1) + result = dpnp.median(ia, axis=-1) + assert_dtype_allclose(result, expected) + + # NumPy's mean gives +0.0 for zeros of either sign + @pytest.mark.parametrize("n", [1, 2, 3, 64, 1000, 2**16 + 1]) + @pytest.mark.parametrize("dtype", get_float_dtypes(no_float16=False)) + def test_zero_sign(self, n, dtype): + a = numpy.full((3, n), -0.0, dtype=dtype) + a[1, ::2] = 0.0 + a[2, : n // 3] = -1 + a[2, n - n // 3 :] = 1 + ia = dpnp.array(a) + + expected = numpy.median(a, axis=-1) + result = dpnp.median(ia, axis=-1) + assert_dtype_allclose(result, expected) + assert_array_equal(dpnp.signbit(result), numpy.signbit(expected)) + + @pytest.mark.parametrize("axis", [0, 1, 2, (0, 2), None]) + def test_non_contiguous(self, axis): + a = generate_random_numpy_array((70, 3, 90), dtype="f4") + ia = dpnp.array(a)[::-1, :, ::2] + a = a[::-1, :, ::2] + + expected = numpy.median(a, axis=axis) + result = dpnp.median(ia, axis=axis) + assert_dtype_allclose(result, expected) + + def test_input_unchanged(self): + a = generate_random_numpy_array((3, 1001), dtype="f4") + ia = dpnp.array(a) + + result = dpnp.median(ia, axis=-1, overwrite_input=True) + assert_dtype_allclose(result, numpy.median(a, axis=-1)) + assert_array_equal(ia, a) + + @pytest.mark.parametrize("keepdims", [False, True]) + def test_out(self, keepdims): + a = generate_random_numpy_array((4, 3)) + ia = dpnp.array(a) + shape = (1, 3) if keepdims else (3,) + + out_np = numpy.empty_like(a, shape=shape) + out_dp = dpnp.empty_like(ia, shape=shape) + expected = numpy.median(a, axis=0, keepdims=keepdims, out=out_np) + result = dpnp.median(ia, axis=0, keepdims=keepdims, out=out_dp) + assert result is out_dp + assert_dtype_allclose(result, expected) + + # shapes the result broadcasts to + @pytest.mark.parametrize( + "axis, out_shape", [(0, (2, 3)), (0, (1, 3)), (None, (3,))] + ) + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_out_wrong_shape(self, xp, axis, out_shape): + a = xp.arange(12.0).reshape(4, 3) + out = xp.empty(out_shape) + with pytest.raises(ValueError): + xp.median(a, axis=axis, out=out) + + def test_ext_empty_batch_clears_all_nan(self): + src = dpt.empty((0, 5), dtype="f4") + dst = dpt.empty((0,), dtype="f4") + all_nan = dpt.full(1, 7, dtype="u4") + mgr = dpctl.utils.SequentialOrderManager[src.sycl_queue] + ht_ev, ev = statistics_ext._median( + src, + True, + dst, + sycl_queue=src.sycl_queue, + all_nan=all_nan, + depends=mgr.submitted_events, + ) + ht_ev.wait() + ev.wait() + assert dpt.asnumpy(all_nan)[0] == 0 + + def test_ext_empty_batch_dtype_error(self): + src = dpt.empty((0, 5), dtype="c8") + dst = dpt.empty((0,), dtype="c8") + with pytest.raises(ValueError, match="No median implementation"): + statistics_ext._median(src, False, dst, sycl_queue=src.sycl_queue) + + class TestPtp: @pytest.mark.parametrize("axis", [None, 0, 1]) @pytest.mark.parametrize( diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py index 40e6a9f35ad7..77dce6cfaf56 100644 --- a/dpnp/tests/test_sycl_queue.py +++ b/dpnp/tests/test_sycl_queue.py @@ -610,6 +610,8 @@ def test_logic_op_2in(op, device): [ pytest.param("searchsorted", [11, 12, 13, 14, 15], 13), pytest.param("broadcast_to", numpy.ones(7), (2, 7)), + pytest.param("partition", [[2.0, 1.0, 7.0], [4.0, 0.0, 3.0]], 1), + pytest.param("argpartition", [[2.0, 1.0, 7.0], [4.0, 0.0, 3.0]], 1), ], ) @pytest.mark.parametrize("device", valid_dev, ids=dev_ids) diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py index 50a34fa25c4a..cec6ff980973 100644 --- a/dpnp/tests/test_usm_type.py +++ b/dpnp/tests/test_usm_type.py @@ -773,6 +773,8 @@ def test_2in_2out(func, data1, data2, usm_type_x, usm_type_y): [ pytest.param("searchsorted", [11, 12, 13, 14, 15], 13), pytest.param("broadcast_to", numpy.ones(7), (2, 7)), + pytest.param("partition", [[2.0, 1.0, 7.0], [4.0, 0.0, 3.0]], 1), + pytest.param("argpartition", [[2.0, 1.0, 7.0], [4.0, 0.0, 3.0]], 1), ], ) @pytest.mark.parametrize("usm_type", list_of_usm_types) diff --git a/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py b/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py index 382393de7216..c17f6553e89d 100644 --- a/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py +++ b/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py @@ -826,7 +826,6 @@ def test_partition_invalid_negative_axis2(self): } ) ) -@pytest.mark.skip("argpartition isn't supported yet") class TestArgpartition(unittest.TestCase): def argpartition(self, a, kth, axis=-1): diff --git a/scripts/gen_coverage.py b/scripts/gen_coverage.py index ebd14955d016..e509a8e64550 100644 --- a/scripts/gen_coverage.py +++ b/scripts/gen_coverage.py @@ -50,12 +50,12 @@ _DPNP_PROFRAW = "dpnp_pytest.profraw" _TENSOR_PROFRAW = "dpnp_tensor_pytest.profraw" # pybind11 modules of dpnp/libtensor holding libtensor functions only, unlike -# _tensor_impl and _tensor_elementwise_impl which hold dpnp functions too +# _tensor_impl, _tensor_elementwise_impl and _tensor_sorting_impl which hold +# dpnp functions too _TENSOR_ONLY_OBJS = ( "_tensor_accumulation_impl.", "_tensor_linalg_impl.", "_tensor_reductions_impl.", - "_tensor_sorting_impl.", )