11580062618

Committed 29 Oct 2024 06:19PM UTC coverage: 87.705% (+0.04%) from 87.669%

Build # 11580062618

Build Type

Pull #1867

github

Committed by

web-flow

Commit Message

Merge d63dd70d2 into 9b83beff0

Pull Request Pull Request #1867: Feature/radix sort

Run Details

3103 of 3618 branches covered (85.77%)

Branch coverage included in aggregate %.

37 of 39 new or added lines in 1 file covered. (94.87%)

49 existing lines in 3 files now uncovered.

11756 of 13324 relevant lines covered (88.23%)

7098.89 hits per line

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

97.33

/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,
    _sort_ascending,
    _sort_descending,
)
from ._tensor_sorting_radix_impl import (
    _radix_argsort_ascending,
    _radix_argsort_descending,
    _radix_sort_ascending,
    _radix_sort_descending,
    _radix_sort_dtype_supported,
)

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

IntelPython / dpctl / 11580062618

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