Implement np.take in kernel by dpnp

numba_dppy/dpnp_glue/dpnp_array_ops_impl.py

@@ -149,3 +149,57 @@ def dpnp_impl(a):
         return out
 
     return dpnp_impl
+
+
+@overload(stubs.dpnp.take)
+def dpnp_take_impl(a, ind):
+    name = "take"
+    dpnp_lowering.ensure_dpnp(name)
+
+    ret_type = types.void
+    """
+    dpnp source:
+    https://github.com/IntelPython/dpnp/blob/0.5.1/dpnp/backend/kernels/dpnp_krnl_indexing.cpp#L34
+    Function declaration:
+    void dpnp_take_c(void* array1_in, void* indices1, void* result1, size_t size)
+    """
+    sig = signature(ret_type, types.voidptr, 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, ind):
+        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)
+
+        ind_usm = dpctl_functions.malloc_shared(ind.size * ind.itemsize, sycl_queue)
+        dpctl_functions.queue_memcpy(
+            sycl_queue, ind_usm, ind.ctypes, ind.size * ind.itemsize
+        )
+
+        out = np.arange(ind.size, dtype=res_dtype).reshape(ind.shape)
+        out_usm = dpctl_functions.malloc_shared(out.size * out.itemsize, sycl_queue)
+
+        dpnp_func(a_usm, ind_usm, out_usm, ind.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(ind_usm, sycl_queue)
+        dpctl_functions.free_with_queue(out_usm, sycl_queue)
+
+        dpnp_ext._dummy_liveness_func([a.size, ind.size, out.size])
+
+        if PRINT_DEBUG:
+            print("dpnp implementation")
+        return out
+
+    return dpnp_impl

numba_dppy/dpnp_glue/dpnp_fptr_interface.pyx

@@ -103,6 +103,7 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_STD
         DPNP_FN_SUBTRACT
         DPNP_FN_SUM
+        DPNP_FN_TAKE
         DPNP_FN_TAN
         DPNP_FN_TANH
         DPNP_FN_TRANSPOSE
@@ -215,6 +216,8 @@ cdef DPNPFuncName get_DPNPFuncName_from_str(name):
         return DPNPFuncName.DPNP_FN_CUMPROD
     elif name == "dpnp_sort":
         return DPNPFuncName.DPNP_FN_SORT
+    elif name == "dpnp_take":
+        return DPNPFuncName.DPNP_FN_TAKE
     else:
         raise ValueError("Unknown dpnp function requested: " + name.split("_")[1])
 

numba_dppy/dpnp_glue/stubs.py

@@ -199,3 +199,6 @@ class cumprod(Stub):
 
     class sort(Stub):
         pass
+
+    class take(Stub):
+        pass

numba_dppy/rename_numpy_functions_pass.py

@@ -91,6 +91,7 @@
     "cumsum": (["numpy"], "cumsum"),
     "cumprod": (["numpy"], "cumprod"),
     "sort": (["numpy"], "sort"),
+    "take": (["numpy"], "take"),
 }
 
 

numba_dppy/tests/test_dpnp_functions.py

@@ -1488,6 +1488,54 @@ def f(a):
                 )
             )
 
+    def check_take_for_different_datatypes(
+        self, fn, test_fn, ind, dims, tys, matrix=False
+    ):
+        for ty in tys:
+            if matrix:
+                a = np.arange(np.prod(dims), dtype=ty).reshape(dims[0], dims[1])
+            else:
+                a = np.arange(dims[0], dtype=ty)
+
+            c = fn(a, ind)
+
+            d = test_fn(a, ind)
+            if c.shape == d.shape:
+                max_abs_err = np.all(c - d)
+            if not (max_abs_err < 1e-4) and c.dtype != d.dtype:
+                return False
+
+        return True
+
+    def test_take(self):
+        @njit
+        def f(a, ind):
+            c = np.take(a, ind)
+            return c
+
+        test_indices = []
+        test_indices.append(np.array([[1, 5, 1], [11, 3, 0]]))
+        test_indices.append(np.array([[[1, 5, 1], [11, 3, 0]]]))
+        test_indices.append(np.array([[[[1, 5]], [[11, 0]], [[1, 2]]]]))
+
+        self.assertTrue(
+            self.check_take_for_different_datatypes(
+                f, np.take, np.array([1, 5, 1, 11, 3]), [12], self.tys
+            )
+        )
+
+        for ind in test_indices:
+            self.assertTrue(
+                self.check_take_for_different_datatypes(
+                    f,
+                    np.take,
+                    ind,
+                    [3, 4],
+                    [np.float],
+                    matrix=True,
+                )
+            )
+
 
 @unittest.skipUnless(
     ensure_dpnp() and dpctl.has_gpu_queues(), "test only when dpNP and GPU is available"