Implement np.sort in kernel by dpnp

numba_dppy/dpnp_glue/dpnp_array_ops_impl.py

@@ -0,0 +1,73 @@
+# Copyright 2021 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 numba_dppy.dpnp_glue.dpnpimpl as dpnp_ext
+from numba import types
+from numba.core.typing import signature
+from . import stubs
+import numba_dppy.dpnp_glue as dpnp_lowering
+from numba.core.extending import overload
+import numpy as np
+from numba_dppy import dpctl_functions
+import numba_dppy
+
+
+@overload(stubs.dpnp.sort)
+def dpnp_sort_impl(a):
+    name = "sort"
+    dpnp_lowering.ensure_dpnp(name)
+
+    ret_type = types.void
+    """
+    dpnp source:
+    https://github.com/IntelPython/dpnp/blob/0.5.0/dpnp/backend/kernels/dpnp_krnl_sorting.cpp#L90
+
+    Function declaration:
+    void dpnp_sort_c(void* array1_in, void* result1, size_t size)
+
+    """
+    sig = signature(ret_type, types.voidptr, types.voidptr, types.intp)
+    dpnp_func = dpnp_ext.dpnp_func("dpnp_" + name, [a.dtype.name, "NONE"], sig)
+
+    res_dtype = a.dtype
+    PRINT_DEBUG = dpnp_lowering.DEBUG
+
+    def dpnp_impl(a):
+        if a.size == 0:
+            raise ValueError("Passed Empty array")
+
+        sycl_queue = dpctl_functions.get_current_queue()
+
+        a_usm = dpctl_functions.malloc_shared(a.size * a.itemsize, sycl_queue)
+        dpctl_functions.queue_memcpy(sycl_queue, a_usm, a.ctypes, a.size * a.itemsize)
+
+        out = np.arange(a.size, dtype=res_dtype)
+        out_usm = dpctl_functions.malloc_shared(out.size * out.itemsize, sycl_queue)
+
+        dpnp_func(a_usm, out_usm, a.size)
+
+        dpctl_functions.queue_memcpy(
+            sycl_queue, out.ctypes, out_usm, out.size * out.itemsize
+        )
+
+        dpctl_functions.free_with_queue(a_usm, sycl_queue)
+        dpctl_functions.free_with_queue(out_usm, sycl_queue)
+
+        dpnp_ext._dummy_liveness_func([a.size, out.size])
+
+        if PRINT_DEBUG:
+            print("dpnp implementation")
+        return out
+
+    return dpnp_impl

numba_dppy/dpnp_glue/dpnp_fptr_interface.pyx

@@ -201,6 +201,8 @@ cdef DPNPFuncName get_DPNPFuncName_from_str(name):
         return DPNPFuncName.DPNP_FN_DET
     elif name == "dpnp_matrix_rank":
         return DPNPFuncName.DPNP_FN_MATRIX_RANK
+    elif name == "dpnp_sort":
+        return DPNPFuncName.DPNP_FN_SORT
     else:
         raise ValueError("Unknown dpnp function requested: " + name.split("_")[1])
 

numba_dppy/dpnp_glue/stubs.py

@@ -178,3 +178,6 @@ class nansum(Stub):
 
     class nanprod(Stub):
         pass
+
+    class sort(Stub):
+        pass

numba_dppy/rename_numpy_functions_pass.py

@@ -82,6 +82,8 @@
     "nansum": (["numpy"], "nansum"),
     "prod": (["numpy"], "prod"),
     "sum": (["numpy"], "sum"),
+    # array ops
+    "sort": (["numpy"], "sort"),
 }
 
 
@@ -192,6 +194,7 @@ def __init__(self):
         import numba_dppy.dpnp_glue.dpnp_statisticsimpl
         import numba_dppy.dpnp_glue.dpnp_sort_search_countimpl
         import numba_dppy.dpnp_glue.dpnp_randomimpl
+        import numba_dppy.dpnp_glue.dpnp_array_ops_impl
 
     def run_pass(self, state):
         rewrite_function_name_pass = RewriteNumPyOverloadedFunctions(

numba_dppy/tests/test_dpnp_functions.py

@@ -1402,5 +1402,25 @@ def f(a, b):
         self.assertTrue(max_abs_err < 1e-4)
 
 
+@unittest.skipUnless(
+    ensure_dpnp() and dpctl.has_gpu_queues(), "test only when dpNP and GPU is available"
+)
+class Testdpnp_array_ops_functions(unittest.TestCase):
+    tys = [np.int32, np.uint32, np.int64, np.uint64, np.float, np.double]
+
+    def test_sort(self):
+        @njit
+        def f(a):
+            c = np.sort(a)
+            return c
+
+        with assert_dpnp_implementaion():
+            self.assertTrue(
+                check_for_different_datatypes(
+                    f, np.sort, [10], 1, self.tys, np_all=True
+                )
+            )
+
+
 if __name__ == "__main__":
     unittest.main()