Skip to content

Implements top_k in dpctl.tensor #1921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ndgrigorian merged 30 commits into from
Jan 5, 2025
Merged

Implements top_k in dpctl.tensor #1921

Changes from 1 commit
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
0b353bf
Implements `top_k`
ndgrigorian Dec 4, 2024
f809573
Add `top_k` to doc sources
ndgrigorian Dec 5, 2024
b59e59d
Use `std::size_t` for `k` instead of `py::ssize_t`
ndgrigorian Dec 5, 2024
b4d7ba4
Add implementation of `top_k` using radix sort
ndgrigorian Dec 5, 2024
c37f1de
Reduce code duplication by using std::optional in py_topk directly
ndgrigorian Dec 11, 2024
1f0e9fd
Clean up top_k docstring
ndgrigorian Dec 11, 2024
fa74e5e
Move top_k into _sorting.py
ndgrigorian Dec 12, 2024
cd83881
top_k raises when k > 1 and input is a scalar
ndgrigorian Dec 13, 2024
70c2ef0
Fix bug in top_k partial merge sort implementation
ndgrigorian Dec 13, 2024
65f14be
Add test file for top_k functionality
oleksandr-pavlyk Dec 14, 2024
275827a
Add sort_utils.hpp with iota_impl
oleksandr-pavlyk Dec 20, 2024
cd1243f
Simplify write-out kernels in topk implementation (avoid recomputing …
oleksandr-pavlyk Dec 22, 2024
823c201
Add explicit read into registers
oleksandr-pavlyk Dec 22, 2024
3a4d8ab
Use unique_ptr as temporary owner of USM allocation
oleksandr-pavlyk Dec 20, 2024
b15e302
Factor out write-out kernel into separate detail function
oleksandr-pavlyk Dec 25, 2024
4ef615e
Typo: smart_malloc_jost -> smart_malloc_host
oleksandr-pavlyk Dec 25, 2024
2608aea
Reimplemented map_back_impl to process few elements per work-item
oleksandr-pavlyk Dec 25, 2024
15da303
Replace use of std::log2(size_t_value)
oleksandr-pavlyk Dec 25, 2024
fb9bc43
Remove dead code disable via preprocessor conditional
oleksandr-pavlyk Dec 27, 2024
6cf36b1
Add information displayed on failure, renamed variables
oleksandr-pavlyk Dec 25, 2024
4ec252e
Add comment to explain ceil_log2 algo
oleksandr-pavlyk Dec 27, 2024
1436522
Add static assers to async_smart_free
oleksandr-pavlyk Dec 27, 2024
f0e3245
Skip top_k test known to fail for some CPU architectures
ndgrigorian Dec 28, 2024
8f4dcdd
Fixed blunder in work-item id to data_id computation
oleksandr-pavlyk Dec 31, 2024
b6806ba
Replace map_back_impl in sort_utils
oleksandr-pavlyk Dec 31, 2024
b47d9f3
Use get_global_linear_id instead of get_global_id and rely on implici…
oleksandr-pavlyk Jan 2, 2025
0493485
Counters in one-workgroup kernel to use uint16_t from uint32_t
oleksandr-pavlyk Jan 2, 2025
5125e11
Apply work-around for failing tests with CPU device and short sub-groups
oleksandr-pavlyk Jan 4, 2025
505b64c
Remove skipping of tests for i1/i2 dtypes since work-around
oleksandr-pavlyk Dec 28, 2024
8c6abf5
Add entry to changelog for top_k
ndgrigorian Jan 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Add information displayed on failure, renamed variables 
Add check of computed against expected indices
  • Loading branch information
@oleksandr-pavlyk @ndgrigorian
oleksandr-pavlyk authored and ndgrigorian committed Jan 5, 2025
commit 6cf36b1760132b2f2e1d200808cfaa79858c11f2
97 changes: 83 additions & 14 deletions dpctl/tests/test_usm_ndarray_top_k.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,38 @@
from dpctl.tests.helper import get_queue_or_skip, skip_if_dtype_not_supported


def _expected_largest_inds(inp, n, shift, k):
"Computed expected top_k indices for mode='largest'"
assert k < n
ones_start_id = shift % (2 * n)

alloc_dev = inp.device

if ones_start_id < n:
expected_inds = dpt.arange(
ones_start_id, ones_start_id + k, dtype="i8", device=alloc_dev
)
else:
# wrap-around
ones_end_id = (ones_start_id + n) % (2 * n)
if ones_end_id >= k:
expected_inds = dpt.arange(k, dtype="i8", device=alloc_dev)
else:
expected_inds = dpt.concat(
(
dpt.arange(ones_end_id, dtype="i8", device=alloc_dev),
dpt.arange(
ones_start_id,
ones_start_id + k - ones_end_id,
dtype="i8",
device=alloc_dev,
),
)
)

return expected_inds


@pytest.mark.parametrize(
"dtype",
[
Expand All @@ -38,23 +70,57 @@
"c16",
],
)
@pytest.mark.parametrize("n", [33, 255, 511, 1021, 8193])
def test_topk_1d_largest(dtype, n):
@pytest.mark.parametrize("n", [33, 43, 255, 511, 1021, 8193])
def test_top_k_1d_largest(dtype, n):
q = get_queue_or_skip()
skip_if_dtype_not_supported(dtype, q)

shift, k = 734, 5
o = dpt.ones(n, dtype=dtype)
z = dpt.zeros(n, dtype=dtype)
zo = dpt.concat((o, z))
inp = dpt.roll(zo, 734)
k = 5
oz = dpt.concat((o, z))
inp = dpt.roll(oz, shift)

expected_inds = _expected_largest_inds(oz, n, shift, k)

s = dpt.top_k(inp, k, mode="largest")
assert s.values.shape == (k,)
assert s.values.dtype == inp.dtype
assert s.indices.shape == (k,)
assert dpt.all(s.values == dpt.ones(k, dtype=dtype))
assert dpt.all(s.values == inp[s.indices])
assert dpt.all(s.indices == expected_inds)
assert dpt.all(s.values == dpt.ones(k, dtype=dtype)), s.values
assert dpt.all(s.values == inp[s.indices]), s.indices


def _expected_smallest_inds(inp, n, shift, k):
"Computed expected top_k indices for mode='smallest'"
assert k < n
zeros_start_id = (n + shift) % (2 * n)
zeros_end_id = (shift) % (2 * n)

alloc_dev = inp.device

if zeros_start_id < zeros_end_id:
expected_inds = dpt.arange(
zeros_start_id, zeros_start_id + k, dtype="i8", device=alloc_dev
)
else:
if zeros_end_id >= k:
expected_inds = dpt.arange(k, dtype="i8", device=alloc_dev)
else:
expected_inds = dpt.concat(
(
dpt.arange(zeros_end_id, dtype="i8", device=alloc_dev),
dpt.arange(
zeros_start_id,
zeros_start_id + k - zeros_end_id,
dtype="i8",
device=alloc_dev,
),
)
)

return expected_inds


@pytest.mark.parametrize(
Expand All @@ -75,20 +141,23 @@ def test_topk_1d_largest(dtype, n):
"c16",
],
)
@pytest.mark.parametrize("n", [33, 255, 257, 513, 1021, 8193])
def test_topk_1d_smallest(dtype, n):
@pytest.mark.parametrize("n", [37, 39, 61, 255, 257, 513, 1021, 8193])
def test_top_k_1d_smallest(dtype, n):
q = get_queue_or_skip()
skip_if_dtype_not_supported(dtype, q)

shift, k = 734, 5
o = dpt.ones(n, dtype=dtype)
z = dpt.zeros(n, dtype=dtype)
zo = dpt.concat((o, z))
inp = dpt.roll(zo, 734)
k = 5
oz = dpt.concat((o, z))
inp = dpt.roll(oz, shift)

expected_inds = _expected_smallest_inds(oz, n, shift, k)

s = dpt.top_k(inp, k, mode="smallest")
assert s.values.shape == (k,)
assert s.values.dtype == inp.dtype
assert s.indices.shape == (k,)
assert dpt.all(s.values == dpt.zeros(k, dtype=dtype))
assert dpt.all(s.values == inp[s.indices])
assert dpt.all(s.indices == expected_inds)
assert dpt.all(s.values == dpt.zeros(k, dtype=dtype)), s.values
assert dpt.all(s.values == inp[s.indices]), s.indices