Skip to content

Add iota kernel #1946

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
oleksandr-pavlyk merged 15 commits into from
Dec 27, 2024
Merged

Add iota kernel #1946

Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
be9f5ce
Add test file for top_k functionality
oleksandr-pavlyk Dec 14, 2024
ff81365
Adds test to triage AMD CPU top_k failure
ndgrigorian Dec 17, 2024
58da416
radix sort triage test now parametrizes for identical sizes to topk_1…
ndgrigorian Dec 18, 2024
fd4bab3
Add sort_utils.hpp with iota_impl
oleksandr-pavlyk Dec 20, 2024
d9cc919
Simplify write-out kernels in topk implementation (avoid recomputing …
oleksandr-pavlyk Dec 22, 2024
eb28d1a
Add explicit read into registers
oleksandr-pavlyk Dec 22, 2024
da3fbcc
Use unique_ptr as temporary owner of USM allocation
oleksandr-pavlyk Dec 20, 2024
79feecd
Factor out write-out kernel into separate detail function
oleksandr-pavlyk Dec 25, 2024
83b1771
Typo: smart_malloc_jost -> smart_malloc_host
oleksandr-pavlyk Dec 25, 2024
1c25d0b
Reimplemented map_back_impl to process few elements per work-item
oleksandr-pavlyk Dec 25, 2024
a74c788
Replace use of std::log2(size_t_value)
oleksandr-pavlyk Dec 25, 2024
9294095
Remove dead code disable via preprocessor conditional
oleksandr-pavlyk Dec 27, 2024
1bb83bc
Add information displayed on failure, renamed variables
oleksandr-pavlyk Dec 25, 2024
fde7408
Add comment to explain ceil_log2 algo
oleksandr-pavlyk Dec 27, 2024
4c8d7d7
Add static assers to async_smart_free
oleksandr-pavlyk Dec 27, 2024
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
oleksandr-pavlyk committed Dec 27, 2024
commit 1bb83bc33ad73b55605536d4d4f7f1e185575bf0
157 changes: 131 additions & 26 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,41 +141,80 @@ 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


# triage failing top k radix implementation on CPU
# replicates from Python behavior of radix sort topk implementation
@pytest.mark.parametrize("n", [33, 255, 511, 1021, 8193])
def test_topk_largest_1d_radix_i1_255(n):
@pytest.mark.parametrize(
"n",
[
33,
34,
35,
36,
37,
38,
39,
40,
41,
42,
43,
44,
45,
46,
47,
48,
49,
50,
61,
137,
255,
511,
1021,
8193,
],
)
def test_top_k_largest_1d_radix_i1(n):
get_queue_or_skip()
dt = "i1"

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

sorted = dpt.copy(dpt.sort(inp, descending=True, kind="radixsort")[:k])
argsorted = dpt.copy(
dpt.argsort(inp, descending=True, kind="radixsort")[:k]
)
assert dpt.all(sorted == dpt.ones(k, dtype=dt))
assert dpt.all(sorted == inp[argsorted])
oz = dpt.concat((o, z))
inp = dpt.roll(oz, shift)

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

sorted_v = dpt.sort(inp, descending=True, kind="radixsort")
argsorted = dpt.argsort(inp, descending=True, kind="radixsort")

assert dpt.all(sorted_v == inp[argsorted])

topk_vals = dpt.copy(sorted_v[:k])
topk_inds = dpt.copy(argsorted[:k])

assert dpt.all(topk_vals == dpt.ones(k, dtype=dt))
assert dpt.all(topk_inds == expected_inds)

assert dpt.all(topk_vals == inp[topk_inds]), topk_inds