Repository navigation
Implement kernels for median and paritition functions - #3138
ndgrigorian wants to merge 35 commits into
Conversation
|
View rendered docs @ https://intelpython.github.io/dpnp/pull/3138/index.html |
|
Array API standard conformance tests for dpnp=0.22.0dev1=np2py314h8d9cdd5_45 ran successfully. |
0f57805 to
8667696
Compare
e24ca7b to
2df1e76
Compare
8667696 to
58ce933
Compare
|
coverage: 84.577% (+0.2%) from 84.424% — feature/implement-median-partition-fns into master |
236c2e1 to
43ca63d
Compare
58ce933 to
89cc1d4
Compare
74de517 to
c7864b9
Compare
759b1f2 to
9ab3e53
Compare
There was a problem hiding this comment.
🟡 Changes recommended
Active compatibility issues remain: even-count median and nanmedian return a finite value where NumPy overflows to infinity, and argpartition rejects valid NumPy integer indices.
2 open issues
What changed in this PR
This PR extends dpnp’s SYCL radix-selection kernels to support median and partition operations, replacing the legacy partition backend and adding public argpartition APIs.
Changes:
- Generalizes radix selection for
median,nanmedian,partition, andargpartition, retaining sorting fallbacks. - Adds
argpartitionfunction and array-method APIs, and fixes median output handling withkeepdims=True. - Expands correctness and device-placement tests, removes legacy partition code, and updates build and coverage configuration.
| File | Description |
|---|---|
scripts/gen_coverage.py |
Includes the sorting extension in dpnp coverage. |
dpnp/tests/third_party/cupy/sorting_tests/test_sort.py |
Enables upstream argpartition tests. |
dpnp/tests/test_usm_type.py |
Tests partition result allocation types. |
dpnp/tests/test_sycl_queue.py |
Tests partition queue placement. |
dpnp/tests/test_statistics.py |
Tests median selection and output handling. |
dpnp/tests/test_sort.py |
Adds partition and argpartition coverage. |
dpnp/tests/test_nanfunctions.py |
Tests NaN-aware median selection. |
dpnp/tests/test_indexing.py |
Enables partition/index equivalence checks. |
dpnp/libtensor/source/tensor_sorting.cpp |
Registers partition bindings. |
dpnp/libtensor/source/statistics/statistics_py.cpp |
Registers median bindings. |
dpnp/libtensor/source/statistics/median.hpp |
Declares median binding initialization. |
dpnp/libtensor/source/statistics/median.cpp |
Implements median validation and dispatch. |
dpnp/libtensor/source/sorting/partition.hpp |
Declares partition binding initialization. |
dpnp/libtensor/source/sorting/partition.cpp |
Implements partition validation and dispatch. |
dpnp/libtensor/source/CMakeLists.txt |
Builds new bindings and adjusts coverage instrumentation. |
dpnp/libtensor/include/kernels/sorting/radix_utils.hpp |
Adds inverse radix keys and revises zero normalization. |
dpnp/libtensor/include/kernels/sorting/radix_select.hpp |
Generalizes selection kernels for partition and median. |
dpnp/dpnp_utils/dpnp_utils_statistics.py |
Routes median operations and handles output views. |
dpnp/dpnp_utils/dpnp_algo_utils.pyx |
Removes inline from _object_to_tuple. |
dpnp/dpnp_iface_sorting.py |
Adds argpartition and shared partition dispatch. |
dpnp/dpnp_iface_nanfunctions.py |
Routes floating-point nanmedian through selection. |
dpnp/dpnp_array.py |
Adds the array argpartition method. |
dpnp/dpnp_algo/dpnp_algo.pyx |
Removes the legacy sorting wrapper include. |
dpnp/dpnp_algo/dpnp_algo.pxd |
Removes the legacy partition enum declaration. |
dpnp/dpnp_algo/dpnp_algo_sorting.pxi |
Deletes the legacy Cython partition wrapper. |
dpnp/dpnp_algo/CMakeLists.txt |
Removes obsolete sorting wrapper dependencies. |
dpnp/backend/src/dpnp_iface_fptr.cpp |
Removes legacy sorting registration. |
dpnp/backend/src/dpnp_fptr.hpp |
Removes the sorting registration declaration. |
dpnp/backend/kernels/dpnp_krnl_sorting.cpp |
Deletes legacy partition kernels. |
dpnp/backend/include/dpnp_iface.hpp |
Removes legacy partition API declarations. |
dpnp/backend/include/dpnp_iface_fptr.hpp |
Removes legacy partition function identifiers. |
dpnp/backend/CMakeLists.txt |
Removes the legacy sorting build source. |
dpnp/__init__.py |
Exports argpartition. |
CHANGELOG.md |
Records new APIs, migration, and output-shape fixes. |
Review effort: Balanced
antonwolfy
left a comment
antonwolfy
left a comment
There was a problem hiding this comment.
Tested on a PTL-H iGPU (level_zero:gpu) and on opencl:cpu. The selection itself looks right: the changed test files, the CuPy sorting and statistics tests and the tensor sort and top_k tests pass on both devices with all integer types, and about 300k fuzz cases against NumPy found no wrong partition, argpartition or top_k result, also for inputs still being written by earlier device work. No new compiler or pytest warnings.
Open bot points I confirm:
- Copilot, median of two equal large middle values (
radix_select.hpp:332): reproduced, and master givesinfthere; details in that thread. - Copilot, NumPy integer
kth(dpnp_iface_sorting.py:73): reproduced fornumpy.int64(1),numpy.array(1),numpy.array([1, 3])and lists of NumPy integers. Master'spartitionrejects them too, butargpartitionis new.
I'd fix the median overflow and the overlap check before merge; the inline points below are smaller.
The internal Jenkins Windows build jobs of this head failed on 11-05 and their logs are gone. Other PRs failed in the same jobs at that time for infrastructure reasons, and the Windows 3.13 build and the GitHub Windows jobs passed, so a rerun should tell.
9ab3e53 to
8616f96
Compare
There was a problem hiding this comment.
🟡 Changes recommended
An exception while obtaining the kernel bundle or allocating temporaries can leave the submitted fill without a keep-alive event, allowing it to write to freed USM memory during cleanup.
1 open issue
2 resolved since last review
Review effort: Balanced
There was a problem hiding this comment.
Re-tested the current head on a PTL-H iGPU (Arc B390, level_zero:gpu) and on opencl:cpu, with master at the merge base built the same way next to it.
The points from the last round are fixed, and I checked each on both devices: even rows average their two middle values for every float type on the sub-group, one-group and multi-group paths, _partition and _argpartition reject src sharing memory with dst, kth behaves like NumPy's for NumPy integers, 0-d and 1-D integer arrays, booleans, floats, and empty and 2-D kth (only numpy.uint64(2**64 - 1) differs: NumPy wraps it to -1), and the docstrings and backend leftovers are gone.
The full dpnp/tests (79.5k passed) and dpnp/tests/tensor (15.5k passed) suites pass on both devices, and so do the changed test files with DPNP_TEST_ALL_INT_TYPES=1. The only pytest warning is in test_sign_nan, which master has too. The build has no compiler warnings, the legacy backend included, and neither do the GitHub Linux and Windows build logs.
Performance, geomean of master time / PR time, alternating builds (master, PR, PR, master), every result checked against NumPy:
| level_zero:gpu | opencl:cpu | |
|---|---|---|
median |
6.1x | 13.6x |
nanmedian |
12.0x | 30.8x |
partition, one kth |
5.1x | 10.4x |
argpartition vs master's argsort |
9.2x | 12.6x |
partition, several kth (still a sort) |
0.99x | 0.99x |
complex median (still a sort) |
0.94x to 1.07x | 0.94x to 1.31x |
top_k |
0.98x | 1.00x |
Shapes are 1-D 106, 106 + 1, 107 and 107 + 1, 1000x1000 along each axis and flattened, 105x16 and 16x106, for float32, int32 and int64, plus float64 and float16 for median. Small inputs gain the most: median of 10 float32 elements takes 31 us instead of 173 us on the GPU. The only thing that got slower is a first call: on opencl:cpu the first 1-D partition JIT-compiles the radix select in 468 ms, where master's backend kernel took 92 ms, and each later call saves about 16 ms. The top_k slowdown is small but consistent; details inline.
Build and footprint (clean in-place build, 8 jobs, 16 cores):
| master | PR | |
|---|---|---|
| wall time | 15:35 | 15:39 |
| sum of all build steps | 4910 s | 5005 s |
largest ld.lld (_tensor_impl) |
3.41 GiB | 3.40 GiB |
ld.lld of _tensor_sorting_impl |
2.73 GiB | 2.90 GiB |
| peak of all build processes | 5.57 GiB | 5.56 GiB |
installed .so files |
183.2 MiB | 187.8 MiB |
The extra work is off the critical path, which is the _tensor_impl link: the two new sources take 17 to 18 s each, and the _tensor_sorting_impl and _statistics_impl links take about 25 s longer each. _tensor_sorting_impl and _statistics_impl grow by 2.7 MiB each, while libdpnp_backend_c shrinks by 0.7 MiB. RSS after import dpnp is 3 MB higher.
Coverage goes from 84.40% to 84.52%. partition.cpp has 99 of 106 lines covered and median.cpp 91 of 117; the uncovered lines are binding error paths, which you left for the binding-test PR, one early return, and a constexpr helper that llvm-cov counts as code. radix_select.hpp isn't reported at all, since llvm-cov ignores the headers under libtensor/include/kernels.
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.
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.
dpt.full submits its fill asynchronously, so without depends the extension's clearing fill could run first and leave the initial 7.
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.
The shortcut writing a single median as is now applies to an odd valid count only, so the median of two finfo.max values overflows to inf as in NumPy. The select tests compare floating data types exactly, since the tolerance of float16 let a neighbouring rank pass.
The radix select gathers from src while writing dst, so not even the same array may be passed as both, as the sort and top_k bindings require.
NumPy integers and integer arrays are accepted, while a boolean or an empty kth raises NumPy's error. An empty integer kth gives the indices in order from argpartition, as in NumPy.
The select path only reads the input, so the examples no longer assert that it was modified.
The comparator and its type helpers were used only by the removed dpnp_krnl_sorting.cpp.
…han in an allocation of their own
… of a lower data type as NumPy does
…too, which takes the sort path
…has host USM atomics
ee84109 to
6d970d3
Compare
vlad-perevezentsev
left a comment
vlad-perevezentsev
left a comment
There was a problem hiding this comment.
LGTM
Thank you @ndgrigorian
This PR proposes the addition of updates to radix selection algorithms to generalize for cases of
medianandpartitionrelated functions, which allows the removal ofdpnp_partition_cfrom the DPNP backendThe median and partition functions can pass through either a sorting or a radix selection path, based on input types (and in the case of partition, always when
kthis multiple elements)Supersedes #2423