Initial check in of sort/argsort per array API spec

Implemented data-parallel merge-sort algorithm.

Author: oleksandr-pavlyk

dpctl/tensor/CMakeLists.txt

@@ -118,6 +118,10 @@ set(_reduction_sources
 set(_boolean_reduction_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/boolean_reductions.cpp
 )
+set(_sorting_sources
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/sort.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/argsort.cpp
+)
 set(_tensor_impl_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/tensor_ctors.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/simplify_iteration_space.cpp
@@ -148,6 +152,10 @@ set(_tensor_reductions_impl_sources
     ${_boolean_reduction_sources}
     ${_reduction_sources}
 )
+set(_tensor_sorting_impl_sources
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/tensor_sorting.cpp
+    ${_sorting_sources}
+)
 
 set(_py_trgts)
 
@@ -166,6 +174,11 @@ pybind11_add_module(${python_module_name} MODULE ${_tensor_reductions_impl_sourc
 add_sycl_to_target(TARGET ${python_module_name} SOURCES ${_tensor_reductions_impl_sources})
 list(APPEND _py_trgts ${python_module_name})
 
+set(python_module_name _tensor_sorting_impl)
+pybind11_add_module(${python_module_name} MODULE ${_tensor_sorting_impl_sources})
+add_sycl_to_target(TARGET ${python_module_name} SOURCES ${_tensor_sorting_impl_sources})
+list(APPEND _py_trgts ${python_module_name})
+
 set(_clang_prefix "")
 if (WIN32)
   set(_clang_prefix "/clang:")
@@ -179,6 +192,7 @@ set(_no_fast_math_sources
 list(APPEND _no_fast_math_sources
      ${_elementwise_sources}
      ${_reduction_sources}
+     ${_sorting_sources}
 )
 
 foreach(_src_fn ${_no_fast_math_sources})

dpctl/tensor/__init__.py

@@ -179,6 +179,7 @@
     reduce_hypot,
     sum,
 )
+from ._sorting import argsort, sort
 from ._testing import allclose
 
 __all__ = [
@@ -346,4 +347,6 @@
     "__array_namespace_info__",
     "reciprocal",
     "angle",
+    "sort",
+    "argsort",
 ]

dpctl/tensor/_sorting.py

@@ -0,0 +1,107 @@
+from numpy.core.numeric import normalize_axis_index
+
+import dpctl
+import dpctl.tensor as dpt
+import dpctl.tensor._tensor_impl as ti
+
+from ._tensor_sorting_impl import (
+    _argsort_ascending,
+    _argsort_descending,
+    _sort_ascending,
+    _sort_descending,
+)
+
+
+def sort(x, axis=-1, descending=False, stable=False):
+    if not isinstance(x, dpt.usm_ndarray):
+        raise TypeError(
+            f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}"
+        )
+    nd = x.ndim
+    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)
+    exec_q = x.sycl_queue
+    host_tasks_list = []
+    impl_fn = _sort_descending if descending else _sort_ascending
+    if arr.flags.c_contiguous:
+        res = dpt.empty_like(arr, order="C")
+        ht_ev, _ = impl_fn(
+            src=arr, trailing_dims_to_sort=1, dst=res, sycl_queue=exec_q
+        )
+        host_tasks_list.append(ht_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
+        )
+        host_tasks_list.append(ht_ev)
+        res = dpt.empty_like(arr, order="C")
+        ht_ev, _ = impl_fn(
+            src=tmp,
+            trailing_dims_to_sort=1,
+            dst=res,
+            sycl_queue=exec_q,
+            depends=[copy_ev],
+        )
+        host_tasks_list.append(ht_ev)
+    if a1 != nd:
+        inv_perm = sorted(range(nd), key=lambda d: perm[d])
+        res = dpt.permute_dims(res, inv_perm)
+    dpctl.SyclEvent.wait_for(host_tasks_list)
+    return res
+
+
+def argsort(x, axis=-1, descending=False, stable=False):
+    if not isinstance(x, dpt.usm_ndarray):
+        raise TypeError(
+            f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}"
+        )
+    nd = x.ndim
+    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)
+    exec_q = x.sycl_queue
+    host_tasks_list = []
+    impl_fn = _argsort_descending if descending else _argsort_ascending
+    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_fn(
+            src=arr, trailing_dims_to_sort=1, dst=res, sycl_queue=exec_q
+        )
+        host_tasks_list.append(ht_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
+        )
+        host_tasks_list.append(ht_ev)
+        res = dpt.empty_like(arr, dtype=index_dt, order="C")
+        ht_ev, _ = impl_fn(
+            src=tmp,
+            trailing_dims_to_sort=1,
+            dst=res,
+            sycl_queue=exec_q,
+            depends=[copy_ev],
+        )
+        host_tasks_list.append(ht_ev)
+    if a1 != nd:
+        inv_perm = sorted(range(nd), key=lambda d: perm[d])
+        res = dpt.permute_dims(res, inv_perm)
+    dpctl.SyclEvent.wait_for(host_tasks_list)
+    return res