From 67a4113562892c0bc211d708f95d373757d221f6 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 31 Oct 2026 12:44:03 -0700 Subject: [PATCH 01/12] Implements kernels for median and partition functions --- dpnp/__init__.py | 2 + dpnp/backend/CMakeLists.txt | 1 - dpnp/backend/include/dpnp_iface.hpp | 32 - dpnp/backend/include/dpnp_iface_fptr.hpp | 19 +- dpnp/backend/kernels/dpnp_krnl_sorting.cpp | 160 -- dpnp/backend/src/dpnp_fptr.hpp | 1 - dpnp/backend/src/dpnp_iface_fptr.cpp | 1 - dpnp/dpnp_algo/CMakeLists.txt | 5 - dpnp/dpnp_algo/dpnp_algo.pxd | 1 - dpnp/dpnp_algo/dpnp_algo.pyx | 6 - dpnp/dpnp_algo/dpnp_algo_sorting.pxi | 127 -- dpnp/dpnp_array.py | 14 +- dpnp/dpnp_iface_nanfunctions.py | 6 +- dpnp/dpnp_iface_sorting.py | 247 ++- dpnp/dpnp_utils/dpnp_algo_utils.pyx | 2 +- dpnp/dpnp_utils/dpnp_utils_statistics.py | 128 +- .../include/kernels/sorting/radix_select.hpp | 1333 ++++++++++++++--- .../include/kernels/sorting/radix_utils.hpp | 32 + dpnp/libtensor/source/CMakeLists.txt | 4 + dpnp/libtensor/source/sorting/partition.cpp | 267 ++++ dpnp/libtensor/source/sorting/partition.hpp | 42 + dpnp/libtensor/source/statistics/median.cpp | 243 +++ dpnp/libtensor/source/statistics/median.hpp | 38 + .../source/statistics/statistics_py.cpp | 2 + dpnp/libtensor/source/tensor_sorting.cpp | 2 + dpnp/tests/test_indexing.py | 3 +- dpnp/tests/test_nanfunctions.py | 49 +- dpnp/tests/test_sort.py | 211 +++ dpnp/tests/test_statistics.py | 65 +- .../cupy/sorting_tests/test_sort.py | 1 - 30 files changed, 2367 insertions(+), 677 deletions(-) delete mode 100644 dpnp/backend/kernels/dpnp_krnl_sorting.cpp delete mode 100644 dpnp/dpnp_algo/dpnp_algo_sorting.pxi create mode 100644 dpnp/libtensor/source/sorting/partition.cpp create mode 100644 dpnp/libtensor/source/sorting/partition.hpp create mode 100644 dpnp/libtensor/source/statistics/median.cpp create mode 100644 dpnp/libtensor/source/statistics/median.hpp 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..14475051de80 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,76 @@ 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" + ) + + _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 +273,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 +311,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,24 +336,11 @@ 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 + if use_select: assert axis == -1 - res = _calc_nanmedian(a_sorted, out=out) + res = _calc_median_by_select(a, ignore_nan, 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 = _calc_median_by_sort(a, axis, out, overwrite_input, ignore_nan) if keepdims: res_shape = list(a_shape) 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..8aafcafa18ca 100644 --- a/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp +++ b/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp @@ -326,6 +326,38 @@ radix_key_t ordered_radix_key(const T &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 { + // 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); + } + } +} + //----------- // projections //----------- diff --git a/dpnp/libtensor/source/CMakeLists.txt b/dpnp/libtensor/source/CMakeLists.txt index f4d203bb893c..5d458e67a7db 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} ) @@ -237,6 +240,7 @@ set(_tensor_reductions_impl_sources set(_tensor_sorting_impl_sources ${CMAKE_CURRENT_SOURCE_DIR}/tensor_sorting.cpp ${_sorting_sources} + ${_dpnp_sorting_sources} ) set(_tensor_linalg_impl_sources ${CMAKE_CURRENT_SOURCE_DIR}/tensor_linalg.cpp diff --git a/dpnp/libtensor/source/sorting/partition.cpp b/dpnp/libtensor/source/sorting/partition.cpp new file mode 100644 index 000000000000..1b100f4fad32 --- /dev/null +++ b/dpnp/libtensor/source/sorting/partition.cpp @@ -0,0 +1,267 @@ +//***************************************************************************** +// 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 (n_iters == 0 || n_values == 0) { + return {sycl::event(), sycl::event()}; + } + 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"); + } + + 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..eb83a2d898fa --- /dev/null +++ b/dpnp/libtensor/source/statistics/median.cpp @@ -0,0 +1,243 @@ +//***************************************************************************** +// 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(&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) { + return {sycl::event(), sycl::event()}; + } + if (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)); + } + 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..045d7bc3cd89 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) @@ -542,6 +542,53 @@ 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 + ) + a[4, :] = -0.0 + a[4, ::3] = numpy.nan + ia = dpnp.array(a) + + expected = numpy.nanmedian(a, axis=-1) + result = dpnp.nanmedian(ia, axis=-1) + assert_dtype_allclose(result, 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("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..15ff7a7ee44d 100644 --- a/dpnp/tests/test_sort.py +++ b/dpnp/tests/test_sort.py @@ -275,6 +275,217 @@ 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)) + + +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..1d0c3008e824 100644 --- a/dpnp/tests/test_statistics.py +++ b/dpnp/tests/test_statistics.py @@ -1004,8 +1004,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 +1020,69 @@ 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) + + @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) + + class TestPtp: @pytest.mark.parametrize("axis", [None, 0, 1]) @pytest.mark.parametrize( 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): From 78261be7f2d03418bd16f9efcf6e1e1efe449f1d Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sun, 1 Nov 2026 22:38:21 -0800 Subject: [PATCH 02/12] Report the coverage of dpnp.partition's kernel partition.cpp is a dpnp source, built with the coverage flags, but gen_coverage.py treated _tensor_sorting_impl as libtensor-only and left it out of the default report. List it first in the module like the other dpnp sources, so the linker keeps its instrumented copies of shared inline functions. --- dpnp/libtensor/source/CMakeLists.txt | 2 +- scripts/gen_coverage.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dpnp/libtensor/source/CMakeLists.txt b/dpnp/libtensor/source/CMakeLists.txt index 5d458e67a7db..447b4d91a23f 100644 --- a/dpnp/libtensor/source/CMakeLists.txt +++ b/dpnp/libtensor/source/CMakeLists.txt @@ -238,9 +238,9 @@ set(_tensor_reductions_impl_sources ${_reduction_sources} ) set(_tensor_sorting_impl_sources + ${_dpnp_sorting_sources} ${CMAKE_CURRENT_SOURCE_DIR}/tensor_sorting.cpp ${_sorting_sources} - ${_dpnp_sorting_sources} ) set(_tensor_linalg_impl_sources ${CMAKE_CURRENT_SOURCE_DIR}/tensor_linalg.cpp 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.", ) From 5ec773c3528d147654800870be1d0cf2983510d3 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 13:51:50 -0800 Subject: [PATCH 03/12] Instrument only the host code of the radix select sources The device coverage counters of partition.cpp and median.cpp reach the sort and top_k kernels through the shared inline helpers, about 500 device globals per kernel image. Their init events make the OpenCL CPU runtime use a freed command, so the coverage tests segfaulted at random in top_k, sort and unique_all. --- dpnp/libtensor/source/CMakeLists.txt | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/dpnp/libtensor/source/CMakeLists.txt b/dpnp/libtensor/source/CMakeLists.txt index 447b4d91a23f..7d3d8a0ff9ff 100644 --- a/dpnp/libtensor/source/CMakeLists.txt +++ b/dpnp/libtensor/source/CMakeLists.txt @@ -338,6 +338,11 @@ 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 instantiating the libtensor radix select kernels +set(_dpnp_radix_select_sources + ${_dpnp_sorting_sources} + ${CMAKE_CURRENT_SOURCE_DIR}/statistics/median.cpp +) foreach(_src_fn ${_dpnp_sources}) get_source_file_property(_cmpl_options_prop ${_src_fn} COMPILE_OPTIONS) set(_combined_options_prop @@ -345,7 +350,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 @@ -362,11 +367,15 @@ 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 sources are host only too: the device counters, ~500 + # device globals per kernel, also reach the shared sort and top_k kernels, + # and their init events crash the OpenCL CPU runtime at random 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 From 1c8d1b0cc0b039640d7bde438ddc648e5a9248c3 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 13:51:50 -0800 Subject: [PATCH 04/12] Avoid an all-NaN row in TestNanMedianSelect::test_basic With n=1 the signed-zero row was all NaN, so NumPy and dpnp warned. test_all_nan_rows covers that case. --- dpnp/tests/test_nanfunctions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dpnp/tests/test_nanfunctions.py b/dpnp/tests/test_nanfunctions.py index 045d7bc3cd89..2b9157d81552 100644 --- a/dpnp/tests/test_nanfunctions.py +++ b/dpnp/tests/test_nanfunctions.py @@ -557,8 +557,9 @@ def test_basic(self, n, dtype): 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, ::3] = numpy.nan + a[4, 1::3] = numpy.nan ia = dpnp.array(a) expected = numpy.nanmedian(a, axis=-1) From c76904058d32f1be35d28462a712904d1fb5c1f7 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 23:48:32 -0800 Subject: [PATCH 05/12] Check the shape of out in the select path of median and nanmedian out[...] = res broadcast the medians into an out of the wrong shape, e.g. (2, 3) or (1, 3) for an axis=0 median of a (4, 3) array, where the sort path raised ValueError. The shape is checked against the result with keepdims applied, as NumPy accepts such an out. NumPy's nanmedian broadcasts too, but dpnp keeps raising as before and as documented. --- dpnp/dpnp_utils/dpnp_utils_statistics.py | 27 +++++++++++++++--------- dpnp/tests/test_nanfunctions.py | 25 ++++++++++++++++++++++ dpnp/tests/test_statistics.py | 24 +++++++++++++++++++++ 3 files changed, 66 insertions(+), 10 deletions(-) diff --git a/dpnp/dpnp_utils/dpnp_utils_statistics.py b/dpnp/dpnp_utils/dpnp_utils_statistics.py index 14475051de80..b16fa0872640 100644 --- a/dpnp/dpnp_utils/dpnp_utils_statistics.py +++ b/dpnp/dpnp_utils/dpnp_utils_statistics.py @@ -64,10 +64,10 @@ def _calc_median(a, axis, out=None): return res -def _calc_median_by_select(a, ignore_nan, out=None): +def _calc_median_by_select(a, ignore_nan, out_shape, out=None): """ Compute the median of an array along its last axis by radix select, - ignoring NaNs if `ignore_nan`. + ignoring NaNs if `ignore_nan`, with the result of shape `out_shape`. """ @@ -79,6 +79,10 @@ def _calc_median_by_select(a, ignore_nan, out=None): raise ExecutionPlacementError( "Input and output allocation queues are not compatible" ) + if out.shape != out_shape: + raise ValueError( + f"Output array of shape {out_shape} is needed, got {out.shape}." + ) _manager = dpu.SequentialOrderManager[exec_q] dep_evs = _manager.submitted_events @@ -127,7 +131,7 @@ def _calc_median_by_select(a, ignore_nan, out=None): 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) + res = dpnp_array._create_from_usm_ndarray(usm_res).reshape(out_shape) if out is None: return res out[...] = res @@ -336,16 +340,19 @@ def dpnp_median( ) axis = -1 - if use_select: - assert axis == -1 - res = _calc_median_by_select(a, ignore_nan, out=out) - else: - res = _calc_median_by_sort(a, axis, out, overwrite_input, ignore_nan) - 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 use_select: + assert axis == -1 + out_shape = res_shape if keepdims else a.shape[:-1] + res = _calc_median_by_select(a, ignore_nan, out_shape, out=out) + else: + res = _calc_median_by_sort(a, axis, out, overwrite_input, ignore_nan) + if keepdims: + res = res.reshape(res_shape) return dpnp.get_result_array(res, out) diff --git a/dpnp/tests/test_nanfunctions.py b/dpnp/tests/test_nanfunctions.py index 2b9157d81552..7d8ffb4f8ea9 100644 --- a/dpnp/tests/test_nanfunctions.py +++ b/dpnp/tests/test_nanfunctions.py @@ -589,6 +589,31 @@ def test_no_warning(self, dtype, recwarn): 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: diff --git a/dpnp/tests/test_statistics.py b/dpnp/tests/test_statistics.py index 1d0c3008e824..51e5e9dfc1f0 100644 --- a/dpnp/tests/test_statistics.py +++ b/dpnp/tests/test_statistics.py @@ -1082,6 +1082,30 @@ def test_input_unchanged(self): 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) + class TestPtp: @pytest.mark.parametrize("axis", [None, 0, 1]) From c9d65fa74e030c9063e1cb0b6f90d8484332ecfc Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 23:48:38 -0800 Subject: [PATCH 06/12] Test the queue and USM type of partition and argpartition results --- dpnp/tests/test_sycl_queue.py | 2 ++ dpnp/tests/test_usm_type.py | 2 ++ 2 files changed, 4 insertions(+) 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) From 7b427aa9a6edb4cb2b15cf3f8233341c900d693b Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 23:48:38 -0800 Subject: [PATCH 07/12] Validate the inputs of _median and _partition for an empty batch Both returned on an empty batch before checking kth, the data types, contiguity and overlap, so a call that raised for a non-empty array passed for an empty one. _median also left the all_nan flag uninitialized. The early returns now come after the checks, and _median zeroes the flag first. _median also checks that src does not overlap all_nan. --- dpnp/libtensor/source/sorting/partition.cpp | 7 ++++--- dpnp/libtensor/source/statistics/median.cpp | 18 +++++++++++++----- dpnp/tests/test_sort.py | 18 ++++++++++++++++++ dpnp/tests/test_statistics.py | 18 ++++++++++++++++++ 4 files changed, 53 insertions(+), 8 deletions(-) diff --git a/dpnp/libtensor/source/sorting/partition.cpp b/dpnp/libtensor/source/sorting/partition.cpp index 1b100f4fad32..121bfc9b55eb 100644 --- a/dpnp/libtensor/source/sorting/partition.cpp +++ b/dpnp/libtensor/source/sorting/partition.cpp @@ -151,9 +151,6 @@ std::pair check_queue({&src, &dst}, names, exec_q); check_writable({&dst}, names); - if (n_iters == 0 || n_values == 0) { - return {sycl::event(), sycl::event()}; - } if (kth >= n_values) { throw py::value_error("`kth` is out of bounds for the last axis"); } @@ -172,6 +169,10 @@ std::pair "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 = diff --git a/dpnp/libtensor/source/statistics/median.cpp b/dpnp/libtensor/source/statistics/median.cpp index eb83a2d898fa..a5ed7cf6ed3e 100644 --- a/dpnp/libtensor/source/statistics/median.cpp +++ b/dpnp/libtensor/source/statistics/median.cpp @@ -144,7 +144,7 @@ std::pair static_cast(td_ns::typenum_t::UINT32)) { throw py::value_error("`all_nan` must be a single uint32 element"); } - check_no_overlap(&dst, &flag, names); + check_no_overlap({&src, &dst}, {&flag}, names); all_nan_p = reinterpret_cast(flag.get_data()); } else { @@ -152,10 +152,7 @@ std::pair check_writable({&dst}, names); } - if (n_iters == 0) { - return {sycl::event(), sycl::event()}; - } - if (n_values == 0) { + if (n_iters != 0 && n_values == 0) { throw py::value_error("The median of an empty axis is undefined"); } @@ -179,6 +176,17 @@ std::pair 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); diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py index 15ff7a7ee44d..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 ( @@ -419,6 +421,22 @@ def test_matches_argpartition(self): 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)) diff --git a/dpnp/tests/test_statistics.py b/dpnp/tests/test_statistics.py index 51e5e9dfc1f0..c623c2c5742f 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 ( @@ -1106,6 +1107,23 @@ def test_out_wrong_shape(self, xp, axis, 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") + ht_ev, ev = statistics_ext._median( + src, True, dst, sycl_queue=src.sycl_queue, all_nan=all_nan + ) + 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]) From d53da0bba8067a1f825d4d44c4f4652ab8b465b3 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 23:48:38 -0800 Subject: [PATCH 08/12] Add the changelog entries of argpartition and the radix select --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 734643223c8a..abf93aa7edc2 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,6 +31,7 @@ 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 From 084a9a0e5771966a79fab5a699c1731bf135240a Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 23:54:59 -0800 Subject: [PATCH 09/12] Accept an out of the keepdims shape in median and nanmedian With keepdims=True the sort path passed out of the keepdims shape to dpnp.mean and _calc_nanmedian, which expect the shape without the reduced axes and raised ValueError, while NumPy accepts it. out is now checked against the keepdims shape once, and both paths write the medians to a view of it with the reduced axes indexed away, which never copies. The select path no longer needs the result shape. --- CHANGELOG.md | 1 + dpnp/dpnp_utils/dpnp_utils_statistics.py | 37 ++++++++++++++++++------ dpnp/tests/test_nanfunctions.py | 15 ++++++++++ dpnp/tests/test_statistics.py | 21 ++++++++++++++ 4 files changed, 65 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index abf93aa7edc2..7141063583c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * 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/dpnp_utils/dpnp_utils_statistics.py b/dpnp/dpnp_utils/dpnp_utils_statistics.py index b16fa0872640..03160ec06d4c 100644 --- a/dpnp/dpnp_utils/dpnp_utils_statistics.py +++ b/dpnp/dpnp_utils/dpnp_utils_statistics.py @@ -64,10 +64,10 @@ def _calc_median(a, axis, out=None): return res -def _calc_median_by_select(a, ignore_nan, out_shape, out=None): +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`, with the result of shape `out_shape`. + ignoring NaNs if `ignore_nan`. """ @@ -79,9 +79,10 @@ def _calc_median_by_select(a, ignore_nan, out_shape, out=None): raise ExecutionPlacementError( "Input and output allocation queues are not compatible" ) - if out.shape != out_shape: + if out.shape != usm_a.shape[:-1]: raise ValueError( - f"Output array of shape {out_shape} is needed, got {out.shape}." + f"Output array of shape {usm_a.shape[:-1]} is needed, " + f"got {out.shape}." ) _manager = dpu.SequentialOrderManager[exec_q] @@ -131,7 +132,7 @@ def _calc_median_by_select(a, ignore_nan, out_shape, out=None): 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).reshape(out_shape) + res = dpnp_array._create_from_usm_ndarray(usm_res) if out is None: return res out[...] = res @@ -340,19 +341,37 @@ def dpnp_median( ) axis = -1 + res_out = out if keepdims: res_shape = list(a_shape) for i in _axis: res_shape[i] = 1 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 - out_shape = res_shape if keepdims else a.shape[:-1] - res = _calc_median_by_select(a, ignore_nan, out_shape, out=out) + res = _calc_median_by_select(a, ignore_nan, out=res_out) else: - res = _calc_median_by_sort(a, axis, out, overwrite_input, ignore_nan) + 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/tests/test_nanfunctions.py b/dpnp/tests/test_nanfunctions.py index 7d8ffb4f8ea9..b4cc7ac09d18 100644 --- a/dpnp/tests/test_nanfunctions.py +++ b/dpnp/tests/test_nanfunctions.py @@ -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 diff --git a/dpnp/tests/test_statistics.py b/dpnp/tests/test_statistics.py index c623c2c5742f..91274737c5c3 100644 --- a/dpnp/tests/test_statistics.py +++ b/dpnp/tests/test_statistics.py @@ -971,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) From 530042a952bac64dc704f49a49058c64a0f77243 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Wed, 4 Nov 2026 00:32:34 -0800 Subject: [PATCH 10/12] Give -0.0 the radix key of +0.0 by comparing keys partition.cpp and median.cpp are dpnp sources, built with the default fast floating-point model of icpx, which keeps -fno-signed-zeros. The compiler could fold `val == 0 ? 0 : val` to `val`, so -0.0 kept its own key, and the median of zeros came back as -0.0, where NumPy's mean gives +0.0. Mapping the key of -0.0 to that of +0.0 with integer compares holds under any floating-point flags. --- .../include/kernels/sorting/radix_utils.hpp | 30 ++++++++++--------- dpnp/tests/test_nanfunctions.py | 15 ++++++++++ dpnp/tests/test_statistics.py | 15 ++++++++++ 3 files changed, 46 insertions(+), 14 deletions(-) diff --git a/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp b/dpnp/libtensor/include/kernels/sorting/radix_utils.hpp index 8aafcafa18ca..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,23 @@ 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(normalize_signed_zero(val))); + return KeyT(order_preserving_cast(val)); } } diff --git a/dpnp/tests/test_nanfunctions.py b/dpnp/tests/test_nanfunctions.py index b4cc7ac09d18..be97422de4f1 100644 --- a/dpnp/tests/test_nanfunctions.py +++ b/dpnp/tests/test_nanfunctions.py @@ -581,6 +581,21 @@ def test_basic(self, n, dtype): 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") diff --git a/dpnp/tests/test_statistics.py b/dpnp/tests/test_statistics.py index 91274737c5c3..201acde2fbcc 100644 --- a/dpnp/tests/test_statistics.py +++ b/dpnp/tests/test_statistics.py @@ -1086,6 +1086,21 @@ def test_special_values(self, n, dtype): 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") From c71a14a0c37745dc42c0e55c823db60845a7c4f5 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Wed, 4 Nov 2026 11:54:49 -0800 Subject: [PATCH 11/12] Order the all_nan fill of the empty-batch median test after its creation dpt.full submits its fill asynchronously, so without depends the extension's clearing fill could run first and leave the initial 7. --- dpnp/tests/test_statistics.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/dpnp/tests/test_statistics.py b/dpnp/tests/test_statistics.py index 201acde2fbcc..37f95c0910f4 100644 --- a/dpnp/tests/test_statistics.py +++ b/dpnp/tests/test_statistics.py @@ -1147,8 +1147,14 @@ 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 + src, + True, + dst, + sycl_queue=src.sycl_queue, + all_nan=all_nan, + depends=mgr.submitted_events, ) ht_ev.wait() ev.wait() From 9ab3e53fcc934a1d9e1bdfa0a058e3daef05b9a1 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Wed, 4 Nov 2026 15:38:36 -0800 Subject: [PATCH 12/12] Instrument only the host code of the whole _statistics_impl module Every kernel image of a module carries all its device coverage counters, so the median kernels still got the 768 counters of the bincount, histogram and sliding window sources after median.cpp became host only. Their device global init events crash the OpenCL CPU runtime in the kernel bundle launches of the radix select, as in nanmedian. The device counts never reach the profile, so the reported coverage is unchanged. --- dpnp/libtensor/source/CMakeLists.txt | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/dpnp/libtensor/source/CMakeLists.txt b/dpnp/libtensor/source/CMakeLists.txt index 7d3d8a0ff9ff..e88c8fb2bd47 100644 --- a/dpnp/libtensor/source/CMakeLists.txt +++ b/dpnp/libtensor/source/CMakeLists.txt @@ -338,11 +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 instantiating the libtensor radix select kernels -set(_dpnp_radix_select_sources - ${_dpnp_sorting_sources} - ${CMAKE_CURRENT_SOURCE_DIR}/statistics/median.cpp -) +# 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 @@ -368,9 +365,10 @@ 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). - # The radix select sources are host only too: the device counters, ~500 - # device globals per kernel, also reach the shared sort and top_k kernels, - # and their init events crash the OpenCL CPU runtime at random + # 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}