Implements `top_k` functions in dpctl.tensor #1921

ndgrigorian · 2024-12-04T10:38:49Z

This PR implements the functions top_k, top_k_indices, and top_k_values as per proposal in array API spec.

Radix and merge sorting are used, and modified merge-sort kernels are introduced which sort the array in chunks and write out to a temporary the k largest or smallest values.

Have you provided a meaningful PR description?
Have you added a test, reproducer or referred to an issue with a reproducer?
Have you tested your changes locally for CPU and GPU devices?
Have you made sure that new changes do not introduce compiler warnings?
Have you checked performance impact of proposed changes?
Have you added documentation for your changes, if necessary?
Have you added your changes to the changelog?
If this PR is a work in progress, are you opening the PR as a draft?

github-actions · 2024-12-04T11:16:50Z

View rendered docs @ https://intelpython.github.io/dpctl/pulls/1921/index.html

github-actions · 2024-12-04T11:26:19Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310hdf72452_296 ran successfully.
Passed: 895
Failed: 1
Skipped: 118

coveralls · 2024-12-04T11:27:46Z

coverage: 87.716% (+0.06%) from 87.659%
when pulling 2865f27 on feature/topk
into 39a19c1 on master.

oleksandr-pavlyk · 2024-12-04T21:10:49Z

@ndgrigorian Please add top_k to the docs/doc_sources/api_reference/dpctl/tensor.sorting_functions.rst:

diff --git a/docs/doc_sources/api_reference/dpctl/tensor.sorting_functions.rst b/docs/doc_sources/api_reference/dpctl/tensor.sorting_functions.rst
index ae1605d988..ef20f4654c 100644
--- a/docs/doc_sources/api_reference/dpctl/tensor.sorting_functions.rst
+++ b/docs/doc_sources/api_reference/dpctl/tensor.sorting_functions.rst
@@ -10,3 +10,4 @@ Sorting functions

    argsort
    sort
+   top_k

github-actions · 2024-12-05T01:00:55Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310hdf72452_295 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

github-actions · 2024-12-06T00:08:39Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310hdf72452_297 ran successfully.
Passed: 895
Failed: 1
Skipped: 118

github-actions · 2024-12-11T19:31:47Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_326 ran successfully.
Passed: 895
Failed: 1
Skipped: 118

github-actions · 2024-12-11T22:19:48Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_327 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

github-actions · 2024-12-12T10:09:12Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_331 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

github-actions · 2024-12-13T01:42:26Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_331 ran successfully.
Passed: 895
Failed: 1
Skipped: 118

github-actions · 2024-12-18T22:28:52Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_331 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

github-actions · 2024-12-27T19:40:03Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_388 ran successfully.
Passed: 893
Failed: 3
Skipped: 118

The implementation leverages existing merge-sort code, and partially sorts the array in cases where a parial sort reduces the size of temporary memory allocation

Reduces amount of casting. `k` will need to fit in `py::ssize_t` regardless.

Instead of using an overload to handle the `axis=None` case, use std::optional and check for trailing_dims_to_search in validation logic

Factored out map_back_impl projects indexing from flat index to a row-wise index. Removed dead code excluded by preprocessor conditional.

Replaced it with hand-written implementation of ceil_log2(n), such that n <= (dectype(n){1} << ceil_log2(n)) is true for all positive values of `n` in the range.

Add check of computed against expected indices

One asserts that at least one unique pointer is specified. Another that specified arguments are unique pointers with USMDeleter.

github-actions · 2024-12-28T01:29:31Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_391 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

github-actions · 2024-12-28T01:43:20Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_386 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

github-actions · 2024-12-28T01:56:53Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_385 ran successfully.
Passed: 895
Failed: 1
Skipped: 118

github-actions · 2024-12-28T01:59:08Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_386 ran successfully.
Passed: 895
Failed: 1
Skipped: 118

github-actions · 2024-12-28T02:52:06Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_387 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

ndgrigorian · 2024-12-28T05:27:43Z

@oleksandr-pavlyk
I've added a skip for both the smallest and largest tests with int8

oleksandr-pavlyk · 2024-12-28T14:30:33Z

Good to see the CI green again!

I was suggesting to only skip it in the test command used in the workflow so that we can still provide a reproducer to the CPU team.

I was thinking we could add a file of tests to skip, and pass it as argument to pytest.

I think this is the approach taken by dpnp.

…tly SYCL bundle DPC++ compiler

gid-lane_id is already a multiple of sg_size.

Change kernel to process few data elements in the work-item.

…t conversion

Counters can not exceed uint16_t max, because the kernel assumes that the number of elements to sort fits into uint16_t. The change reduces the kernel SLM footprint. Also, remove use of std::move, uint16_t->std::uint16_t, etc Replace size_t->std::size_t, uint32_t->std::uint32_t Use `if constexpr` in order-preservign-cast for better readability.

The team developing OpenCL:CPU device runtime and compiler was notified. See CMPLRLLVM-64592 Once fixed, the work-around should be removed.

was applied in C++. Add tests for 2d input arrays, for axis=0 and axis=1 Add a test for non-contiguous input, 0d input, validation 100% coverage of top_k function implementation achieved

oleksandr-pavlyk

I think this PR is ready to go in.

Thank you for working on this feature, @ndgrigorian

Any additional changes can be deferred to future PRs.

github-actions · 2025-01-04T19:42:18Z

Array API standard conformance tests for dpctl=0.19.0dev0=py310h93fe807_399 ran successfully.
Passed: 894
Failed: 2
Skipped: 118

oleksandr-pavlyk · 2025-01-04T19:45:40Z

I'd suggest rebasing the branch on top of the targeted base branch to remove two cherry-picked commits fixing the workflow for building with nightly DPC++ bundle

oleksandr-pavlyk · 2025-01-04T20:56:45Z

Also, to make GH rule checker happy, add a line to changelog, so that your commit is the last one, otherwise my approve does not count

ndgrigorian requested a review from AlexanderKalistratov December 4, 2024 10:38

ndgrigorian force-pushed the feature/topk branch from 79b97d9 to 882c70d Compare December 5, 2024 00:16

ndgrigorian force-pushed the feature/topk branch from a56e21c to 26718f3 Compare December 11, 2024 18:46

ndgrigorian marked this pull request as ready for review December 12, 2024 09:24

ndgrigorian requested a review from oleksandr-pavlyk as a code owner December 12, 2024 09:24

ndgrigorian force-pushed the feature/topk branch 2 times, most recently from 8bcb100 to 8f38b80 Compare December 13, 2024 01:00

ndgrigorian force-pushed the feature/topk branch 3 times, most recently from 775b423 to 84d1388 Compare December 28, 2024 01:00

ndgrigorian added 8 commits December 27, 2024 17:12

Implements top_k

b53e1e2

The implementation leverages existing merge-sort code, and partially sorts the array in cases where a parial sort reduces the size of temporary memory allocation

Add top_k to doc sources

dac1daa

Use std::size_t for k instead of py::ssize_t

3fa4167

Reduces amount of casting. `k` will need to fit in `py::ssize_t` regardless.

Add implementation of top_k using radix sort

db7b352

Reduce code duplication by using std::optional in py_topk directly

678a165

Instead of using an overload to handle the `axis=None` case, use std::optional and check for trailing_dims_to_search in validation logic

Clean up top_k docstring

3cfc943

Move top_k into _sorting.py

f285582

top_k raises when k > 1 and input is a scalar

594c9eb

oleksandr-pavlyk added 7 commits December 27, 2024 17:12

Typo: smart_malloc_jost -> smart_malloc_host

624141a

Reimplemented map_back_impl to process few elements per work-item

fe99957

Factored out map_back_impl projects indexing from flat index to a row-wise index. Removed dead code excluded by preprocessor conditional.

Replace use of std::log2(size_t_value)

cb5d213

Replaced it with hand-written implementation of ceil_log2(n), such that n <= (dectype(n){1} << ceil_log2(n)) is true for all positive values of `n` in the range.

Remove dead code disable via preprocessor conditional

5001200

Add information displayed on failure, renamed variables

723c3f4

Add check of computed against expected indices

Add comment to explain ceil_log2 algo

b71b128

Add static assers to async_smart_free

809cb70

One asserts that at least one unique pointer is specified. Another that specified arguments are unique pointers with USMDeleter.

ndgrigorian force-pushed the feature/topk branch from 84d1388 to 809cb70 Compare December 28, 2024 01:13

Merge branch 'master' into feature/topk

7ad7b73

Skip top_k test known to fail for some CPU architectures

33f1ed4

oleksandr-pavlyk added 9 commits January 4, 2025 07:20

Use Ubuntu 24.04 to run the OS-LLVM workflow to build dpctl with nigh…

7436de6

…tly SYCL bundle DPC++ compiler

No need to install libtinfo5 on Ubuntu 24.04

bff71fc

Fixed blunder in work-item id to data_id computation

addb341

gid-lane_id is already a multiple of sg_size.

Replace map_back_impl in sort_utils

399cdd1

Change kernel to process few data elements in the work-item.

Use get_global_linear_id instead of get_global_id and rely on implici…

1a64014

…t conversion

Apply work-around for failing tests with CPU device and short sub-groups

7c7d8f9

The team developing OpenCL:CPU device runtime and compiler was notified. See CMPLRLLVM-64592 Once fixed, the work-around should be removed.

Remove skipping of tests for i1/i2 dtypes since work-around

e1b7540

was applied in C++. Add tests for 2d input arrays, for axis=0 and axis=1 Add a test for non-contiguous input, 0d input, validation 100% coverage of top_k function implementation achieved

Merge pull request #1951 from IntelPython/improve-top_k-coverage

2865f27

oleksandr-pavlyk approved these changes Jan 4, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implements `top_k` functions in dpctl.tensor #1921

Implements `top_k` functions in dpctl.tensor #1921

ndgrigorian commented Dec 4, 2024 •

edited by oleksandr-pavlyk

Loading

github-actions bot commented Dec 4, 2024

github-actions bot commented Dec 4, 2024

coveralls commented Dec 4, 2024 •

edited

Loading

oleksandr-pavlyk commented Dec 4, 2024

github-actions bot commented Dec 5, 2024

github-actions bot commented Dec 6, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 13, 2024

github-actions bot commented Dec 18, 2024

github-actions bot commented Dec 27, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

ndgrigorian commented Dec 28, 2024

oleksandr-pavlyk commented Dec 28, 2024

oleksandr-pavlyk left a comment

github-actions bot commented Jan 4, 2025

oleksandr-pavlyk commented Jan 4, 2025

oleksandr-pavlyk commented Jan 4, 2025

Implements top_k functions in dpctl.tensor #1921

Are you sure you want to change the base?

Implements top_k functions in dpctl.tensor #1921

Conversation

ndgrigorian commented Dec 4, 2024 • edited by oleksandr-pavlyk Loading

github-actions bot commented Dec 4, 2024

github-actions bot commented Dec 4, 2024

coveralls commented Dec 4, 2024 • edited Loading

oleksandr-pavlyk commented Dec 4, 2024

github-actions bot commented Dec 5, 2024

github-actions bot commented Dec 6, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 13, 2024

github-actions bot commented Dec 18, 2024

github-actions bot commented Dec 27, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

github-actions bot commented Dec 28, 2024

ndgrigorian commented Dec 28, 2024

oleksandr-pavlyk commented Dec 28, 2024

oleksandr-pavlyk left a comment

Choose a reason for hiding this comment

github-actions bot commented Jan 4, 2025

oleksandr-pavlyk commented Jan 4, 2025

oleksandr-pavlyk commented Jan 4, 2025

Implements `top_k` functions in dpctl.tensor #1921

Implements `top_k` functions in dpctl.tensor #1921

ndgrigorian commented Dec 4, 2024 •

edited by oleksandr-pavlyk

Loading

coveralls commented Dec 4, 2024 •

edited

Loading