11461930810

Committed 22 Oct 2024 02:05PM UTC coverage: 87.691% (+0.02%) from 87.669%

Build # 11461930810

Build Type

Pull #1867

github

Committed by

web-flow

Commit Message

Merge 3c05c1bda into 9eb8f0303

Pull Request Pull Request #1867: Feature/radix sort

Run Details

3099 of 3614 branches covered (85.75%)

Branch coverage included in aggregate %.

36 of 38 new or added lines in 1 file covered. (94.74%)

6 existing lines in 1 file now uncovered.

11748 of 13317 relevant lines covered (88.22%)

7101.08 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.32

/dpctl/tensor/_sorting.py

#                       Data Parallel Control (dpctl)
#
#  Copyright 2020-2024 Intel Corporation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import dpctl.tensor as dpt
import dpctl.tensor._tensor_impl as ti
import dpctl.utils as du

from ._numpy_helper import normalize_axis_index
from ._tensor_sorting_impl import (
    _argsort_ascending,
    _argsort_descending,
    _radix_argsort_ascending,
    _radix_argsort_descending,
    _radix_sort_ascending,
    _radix_sort_descending,
    _radix_sort_dtype_supported,
    _sort_ascending,
    _sort_descending,
)

__all__ = ["sort", "argsort"]


def _get_mergesort_impl_fn(descending):
    return _sort_descending if descending else _sort_ascending


def _get_radixsort_impl_fn(descending):
    return _radix_sort_descending if descending else _radix_sort_ascending


def sort(x, /, *, axis=-1, descending=False, stable=True, kind=None):
    """sort(x, axis=-1, descending=False, stable=True)

    Returns a sorted copy of an input array `x`.

    Args:
        x (usm_ndarray):
            input array.
        axis (Optional[int]):
            axis along which to sort. If set to `-1`, the function
            must sort along the last axis. Default: `-1`.
        descending (Optional[bool]):
            sort order. If `True`, the array must be sorted in descending
            order (by value). If `False`, the array must be sorted in
            ascending order (by value). Default: `False`.
        stable (Optional[bool]):
            sort stability. If `True`, the returned array must maintain the
            relative order of `x` values which compare as equal. If `False`,
            the returned array may or may not maintain the relative order of
            `x` values which compare as equal. Default: `True`.
        kind (Optional[Literal["stable", "mergesort", "radixsort"]]):
            Sorting algorithm. The default is `"stable"`, which uses parallel
            merge-sort or parallel radix-sort algorithms depending on the
            array data type.
    Returns:
        usm_ndarray:
            a sorted array. The returned array has the same data type and
            the same shape as the input array `x`.
    """
    if not isinstance(x, dpt.usm_ndarray):
        raise TypeError(
            f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}"
        )
    nd = x.ndim
    if nd == 0:
        axis = normalize_axis_index(axis, ndim=1, msg_prefix="axis")
        return dpt.copy(x, order="C")
    else:
        axis = normalize_axis_index(axis, ndim=nd, msg_prefix="axis")
    a1 = axis + 1
    if a1 == nd:
        perm = list(range(nd))
        arr = x
    else:
        perm = [i for i in range(nd) if i != axis] + [
            axis,
        ]
        arr = dpt.permute_dims(x, perm)
    if kind is None:
        kind = "stable"
    if not isinstance(kind, str) or kind not in [
        "stable",
        "radixsort",
        "mergesort",
    ]:
        raise ValueError(
            "Unsupported kind value. Expected 'stable', 'mergesort', "
            f"or 'radixsort', but got '{kind}'"
        )
    if kind == "mergesort":
        impl_fn = _get_mergesort_impl_fn(descending)
    elif kind == "radixsort":
        if _radix_sort_dtype_supported(x.dtype.num):
            impl_fn = _get_radixsort_impl_fn(descending)
        else:
            raise ValueError(f"Radix sort is not supported for {x.dtype}")
    else:
        dt = x.dtype
        if dt in [dpt.bool, dpt.uint8, dpt.int8, dpt.int16, dpt.uint16]:
            impl_fn = _get_radixsort_impl_fn(descending)
        else:
            impl_fn = _get_mergesort_impl_fn(descending)
    exec_q = x.sycl_queue
    _manager = du.SequentialOrderManager[exec_q]
    dep_evs = _manager.submitted_events
    if arr.flags.c_contiguous:
        res = dpt.empty_like(arr, order="C")
        ht_ev, impl_ev = impl_fn(
            src=arr,
            trailing_dims_to_sort=1,
            dst=res,
            sycl_queue=exec_q,
            depends=dep_evs,
        )
        _manager.add_event_pair(ht_ev, impl_ev)
    else:
        tmp = dpt.empty_like(arr, order="C")
        ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
            src=arr, dst=tmp, sycl_queue=exec_q, depends=dep_evs
        )
        _manager.add_event_pair(ht_ev, copy_ev)
        res = dpt.empty_like(arr, order="C")
        ht_ev, impl_ev = impl_fn(
            src=tmp,
            trailing_dims_to_sort=1,
            dst=res,
            sycl_queue=exec_q,
            depends=[copy_ev],
        )
        _manager.add_event_pair(ht_ev, impl_ev)
    if a1 != nd:
        inv_perm = sorted(range(nd), key=lambda d: perm[d])
        res = dpt.permute_dims(res, inv_perm)
    return res


def _get_mergeargsort_impl_fn(descending):
    return _argsort_descending if descending else _argsort_ascending


def _get_radixargsort_impl_fn(descending):
    return _radix_argsort_descending if descending else _radix_argsort_ascending


def argsort(x, axis=-1, descending=False, stable=True, kind=None):
    """argsort(x, axis=-1, descending=False, stable=True)

    Returns the indices that sort an array `x` along a specified axis.

    Args:
        x (usm_ndarray):
            input array.
        axis (Optional[int]):
            axis along which to sort. If set to `-1`, the function
            must sort along the last axis. Default: `-1`.
        descending (Optional[bool]):
            sort order. If `True`, the array must be sorted in descending
            order (by value). If `False`, the array must be sorted in
            ascending order (by value). Default: `False`.
        stable (Optional[bool]):
            sort stability. If `True`, the returned array must maintain the
            relative order of `x` values which compare as equal. If `False`,
            the returned array may or may not maintain the relative order of
            `x` values which compare as equal. Default: `True`.
        kind (Optional[Literal["stable", "mergesort", "radixsort"]]):
            Sorting algorithm. The default is `"stable"`, which uses parallel
            merge-sort or parallel radix-sort algorithms depending on the
            array data type.

    Returns:
        usm_ndarray:
            an array of indices. The returned array has the  same shape as
            the input array `x`. The return array has default array index
            data type.
    """
    if not isinstance(x, dpt.usm_ndarray):
        raise TypeError(
            f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}"
        )
    nd = x.ndim
    if nd == 0:
        axis = normalize_axis_index(axis, ndim=1, msg_prefix="axis")
        return dpt.zeros_like(
            x, dtype=ti.default_device_index_type(x.sycl_queue), order="C"
        )
    else:
        axis = normalize_axis_index(axis, ndim=nd, msg_prefix="axis")
    a1 = axis + 1
    if a1 == nd:
        perm = list(range(nd))
        arr = x
    else:
        perm = [i for i in range(nd) if i != axis] + [
            axis,
        ]
        arr = dpt.permute_dims(x, perm)
    if kind is None:
        kind = "stable"
    if not isinstance(kind, str) or kind not in [
        "stable",
        "radixsort",
        "mergesort",
    ]:
        raise ValueError(
            "Unsupported kind value. Expected 'stable', 'mergesort', "
            f"or 'radixsort', but got '{kind}'"
        )
    if kind == "mergesort":
        impl_fn = _get_mergeargsort_impl_fn(descending)
    elif kind == "radixsort":
        if _radix_sort_dtype_supported(x.dtype.num):
            impl_fn = _get_radixargsort_impl_fn(descending)
        else:
            raise ValueError(f"Radix sort is not supported for {x.dtype}")
    else:
        dt = x.dtype
        if dt in [dpt.bool, dpt.uint8, dpt.int8, dpt.int16, dpt.uint16]:
            impl_fn = _get_radixargsort_impl_fn(descending)
        else:
            impl_fn = _get_mergeargsort_impl_fn(descending)
    exec_q = x.sycl_queue
    _manager = du.SequentialOrderManager[exec_q]
    dep_evs = _manager.submitted_events
    index_dt = ti.default_device_index_type(exec_q)
    if arr.flags.c_contiguous:
        res = dpt.empty_like(arr, dtype=index_dt, order="C")
        ht_ev, impl_ev = impl_fn(
            src=arr,
            trailing_dims_to_sort=1,
            dst=res,
            sycl_queue=exec_q,
            depends=dep_evs,
        )
        _manager.add_event_pair(ht_ev, impl_ev)
    else:
        tmp = dpt.empty_like(arr, order="C")
        ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
            src=arr, dst=tmp, sycl_queue=exec_q, depends=dep_evs
        )
        _manager.add_event_pair(ht_ev, copy_ev)
        res = dpt.empty_like(arr, dtype=index_dt, order="C")
        ht_ev, impl_ev = impl_fn(
            src=tmp,
            trailing_dims_to_sort=1,
            dst=res,
            sycl_queue=exec_q,
            depends=[copy_ev],
        )
        _manager.add_event_pair(ht_ev, impl_ev)
    if a1 != nd:
        inv_perm = sorted(range(nd), key=lambda d: perm[d])
        res = dpt.permute_dims(res, inv_perm)
    return res

1	# Data Parallel Control (dpctl)
2	#
3	# Copyright 2020-2024 Intel Corporation
4	#
5	# Licensed under the Apache License, Version 2.0 (the "License");
6	# you may not use this file except in compliance with the License.
7	# You may obtain a copy of the License at
8	#
9	# http://www.apache.org/licenses/LICENSE-2.0
10	#
11	# Unless required by applicable law or agreed to in writing, software
12	# distributed under the License is distributed on an "AS IS" BASIS,
13	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14	# See the License for the specific language governing permissions and
15	# limitations under the License.
16
17	import dpctl.tensor as dpt	1✔
18	import dpctl.tensor._tensor_impl as ti	1✔
19	import dpctl.utils as du	1✔
20
21	from ._numpy_helper import normalize_axis_index	1✔
22	from ._tensor_sorting_impl import (	1✔
23	_argsort_ascending,
24	_argsort_descending,
25	_radix_argsort_ascending,
26	_radix_argsort_descending,
27	_radix_sort_ascending,
28	_radix_sort_descending,
29	_radix_sort_dtype_supported,
30	_sort_ascending,
31	_sort_descending,
32	)
33
34	__all__ = ["sort", "argsort"]	1✔
35
36
37	def _get_mergesort_impl_fn(descending):	1✔
38	return _sort_descending if descending else _sort_ascending	1✔
39
40
41	def _get_radixsort_impl_fn(descending):	1✔
42	return _radix_sort_descending if descending else _radix_sort_ascending	1✔
43
44
45	def sort(x, /, *, axis=-1, descending=False, stable=True, kind=None):	1✔
46	"""sort(x, axis=-1, descending=False, stable=True)
47
48	Returns a sorted copy of an input array `x`.
49
50	Args:
51	x (usm_ndarray):
52	input array.
53	axis (Optional[int]):
54	axis along which to sort. If set to `-1`, the function
55	must sort along the last axis. Default: `-1`.
56	descending (Optional[bool]):
57	sort order. If `True`, the array must be sorted in descending
58	order (by value). If `False`, the array must be sorted in
59	ascending order (by value). Default: `False`.
60	stable (Optional[bool]):
61	sort stability. If `True`, the returned array must maintain the
62	relative order of `x` values which compare as equal. If `False`,
63	the returned array may or may not maintain the relative order of
64	`x` values which compare as equal. Default: `True`.
65	kind (Optional[Literal["stable", "mergesort", "radixsort"]]):
66	Sorting algorithm. The default is `"stable"`, which uses parallel
67	merge-sort or parallel radix-sort algorithms depending on the
68	array data type.
69	Returns:
70	usm_ndarray:
71	a sorted array. The returned array has the same data type and
72	the same shape as the input array `x`.
73	"""
74	if not isinstance(x, dpt.usm_ndarray):	1✔
75	raise TypeError(	1✔
76	f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}"
77	)
78	nd = x.ndim	1✔
79	if nd == 0:	1✔
80	axis = normalize_axis_index(axis, ndim=1, msg_prefix="axis")	1✔
81	return dpt.copy(x, order="C")	1✔
82	else:
83	axis = normalize_axis_index(axis, ndim=nd, msg_prefix="axis")	1✔
84	a1 = axis + 1	1✔
85	if a1 == nd:	1✔
86	perm = list(range(nd))	1✔
87	arr = x	1✔
88	else:
89	perm = [i for i in range(nd) if i != axis] + [	1✔
90	axis,
91	]
92	arr = dpt.permute_dims(x, perm)	1✔
93	if kind is None:	1✔
94	kind = "stable"	1✔
95	if not isinstance(kind, str) or kind not in [	1✔
96	"stable",
97	"radixsort",
98	"mergesort",
99	]:
100	raise ValueError(	1✔
101	"Unsupported kind value. Expected 'stable', 'mergesort', "
102	f"or 'radixsort', but got '{kind}'"
103	)
104	if kind == "mergesort":	1✔
105	impl_fn = _get_mergesort_impl_fn(descending)	1✔
106	elif kind == "radixsort":	1✔
107	if _radix_sort_dtype_supported(x.dtype.num):	1!
108	impl_fn = _get_radixsort_impl_fn(descending)	1✔
109	else:
NEW 110	raise ValueError(f"Radix sort is not supported for {x.dtype}")	×
111	else:
112	dt = x.dtype	1✔
113	if dt in [dpt.bool, dpt.uint8, dpt.int8, dpt.int16, dpt.uint16]:	1✔
114	impl_fn = _get_radixsort_impl_fn(descending)	1✔
115	else:
116	impl_fn = _get_mergesort_impl_fn(descending)	1✔
117	exec_q = x.sycl_queue	1✔
118	_manager = du.SequentialOrderManager[exec_q]	1✔
119	dep_evs = _manager.submitted_events	1✔
120	if arr.flags.c_contiguous:	1✔
121	res = dpt.empty_like(arr, order="C")	1✔
122	ht_ev, impl_ev = impl_fn(	1✔
123	src=arr,
124	trailing_dims_to_sort=1,
125	dst=res,
126	sycl_queue=exec_q,
127	depends=dep_evs,
128	)
129	_manager.add_event_pair(ht_ev, impl_ev)	1✔
130	else:
131	tmp = dpt.empty_like(arr, order="C")	1✔
132	ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(	1✔
133	src=arr, dst=tmp, sycl_queue=exec_q, depends=dep_evs
134	)
135	_manager.add_event_pair(ht_ev, copy_ev)	1✔
136	res = dpt.empty_like(arr, order="C")	1✔
137	ht_ev, impl_ev = impl_fn(	1✔
138	src=tmp,
139	trailing_dims_to_sort=1,
140	dst=res,
141	sycl_queue=exec_q,
142	depends=[copy_ev],
143	)
144	_manager.add_event_pair(ht_ev, impl_ev)	1✔
145	if a1 != nd:	1✔
146	inv_perm = sorted(range(nd), key=lambda d: perm[d])	1✔
147	res = dpt.permute_dims(res, inv_perm)	1✔
148	return res	1✔
149
150
151	def _get_mergeargsort_impl_fn(descending):	1✔
152	return _argsort_descending if descending else _argsort_ascending	1✔
153
154
155	def _get_radixargsort_impl_fn(descending):	1✔
156	return _radix_argsort_descending if descending else _radix_argsort_ascending	1✔
157
158
159	def argsort(x, axis=-1, descending=False, stable=True, kind=None):	1✔
160	"""argsort(x, axis=-1, descending=False, stable=True)
161
162	Returns the indices that sort an array `x` along a specified axis.
163
164	Args:
165	x (usm_ndarray):
166	input array.
167	axis (Optional[int]):
168	axis along which to sort. If set to `-1`, the function
169	must sort along the last axis. Default: `-1`.
170	descending (Optional[bool]):
171	sort order. If `True`, the array must be sorted in descending
172	order (by value). If `False`, the array must be sorted in
173	ascending order (by value). Default: `False`.
174	stable (Optional[bool]):
175	sort stability. If `True`, the returned array must maintain the
176	relative order of `x` values which compare as equal. If `False`,
177	the returned array may or may not maintain the relative order of
178	`x` values which compare as equal. Default: `True`.
179	kind (Optional[Literal["stable", "mergesort", "radixsort"]]):
180	Sorting algorithm. The default is `"stable"`, which uses parallel
181	merge-sort or parallel radix-sort algorithms depending on the
182	array data type.
183
184	Returns:
185	usm_ndarray:
186	an array of indices. The returned array has the same shape as
187	the input array `x`. The return array has default array index
188	data type.
189	"""
190	if not isinstance(x, dpt.usm_ndarray):	1✔
191	raise TypeError(	1✔
192	f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}"
193	)
194	nd = x.ndim	1✔
195	if nd == 0:	1✔
196	axis = normalize_axis_index(axis, ndim=1, msg_prefix="axis")	1✔
197	return dpt.zeros_like(	1✔
198	x, dtype=ti.default_device_index_type(x.sycl_queue), order="C"
199	)
200	else:
201	axis = normalize_axis_index(axis, ndim=nd, msg_prefix="axis")	1✔
202	a1 = axis + 1	1✔
203	if a1 == nd:	1✔
204	perm = list(range(nd))	1✔
205	arr = x	1✔
206	else:
207	perm = [i for i in range(nd) if i != axis] + [	1✔
208	axis,
209	]
210	arr = dpt.permute_dims(x, perm)	1✔
211	if kind is None:	1✔
212	kind = "stable"	1✔
213	if not isinstance(kind, str) or kind not in [	1✔
214	"stable",
215	"radixsort",
216	"mergesort",
217	]:
218	raise ValueError(	1✔
219	"Unsupported kind value. Expected 'stable', 'mergesort', "
220	f"or 'radixsort', but got '{kind}'"
221	)
222	if kind == "mergesort":	1✔
223	impl_fn = _get_mergeargsort_impl_fn(descending)	1✔
224	elif kind == "radixsort":	1✔
225	if _radix_sort_dtype_supported(x.dtype.num):	1!
226	impl_fn = _get_radixargsort_impl_fn(descending)	1✔
227	else:
NEW 228	raise ValueError(f"Radix sort is not supported for {x.dtype}")	×
229	else:
230	dt = x.dtype	1✔
231	if dt in [dpt.bool, dpt.uint8, dpt.int8, dpt.int16, dpt.uint16]:	1✔
232	impl_fn = _get_radixargsort_impl_fn(descending)	1✔
233	else:
234	impl_fn = _get_mergeargsort_impl_fn(descending)	1✔
235	exec_q = x.sycl_queue	1✔
236	_manager = du.SequentialOrderManager[exec_q]	1✔
237	dep_evs = _manager.submitted_events	1✔
238	index_dt = ti.default_device_index_type(exec_q)	1✔
239	if arr.flags.c_contiguous:	1✔
240	res = dpt.empty_like(arr, dtype=index_dt, order="C")	1✔
241	ht_ev, impl_ev = impl_fn(	1✔
242	src=arr,
243	trailing_dims_to_sort=1,
244	dst=res,
245	sycl_queue=exec_q,
246	depends=dep_evs,
247	)
248	_manager.add_event_pair(ht_ev, impl_ev)	1✔
249	else:
250	tmp = dpt.empty_like(arr, order="C")	1✔
251	ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(	1✔
252	src=arr, dst=tmp, sycl_queue=exec_q, depends=dep_evs
253	)
254	_manager.add_event_pair(ht_ev, copy_ev)	1✔
255	res = dpt.empty_like(arr, dtype=index_dt, order="C")	1✔
256	ht_ev, impl_ev = impl_fn(	1✔
257	src=tmp,
258	trailing_dims_to_sort=1,
259	dst=res,
260	sycl_queue=exec_q,
261	depends=[copy_ev],
262	)
263	_manager.add_event_pair(ht_ev, impl_ev)	1✔
264	if a1 != nd:	1✔
265	inv_perm = sorted(range(nd), key=lambda d: perm[d])	1✔
266	res = dpt.permute_dims(res, inv_perm)	1✔
267	return res	1✔

IntelPython / dpctl / 11461930810

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous