From 52b356ade18c9dab5e33e3055963470951821bac Mon Sep 17 00:00:00 2001
From: Diptorup Deb <diptorup.deb@intel.com>
Date: Tue, 24 Oct 2023 11:38:52 -0500
Subject: [PATCH] Add unit tests

---
 numba_dpex/tests/experimental/__init__.py     |   3 +
 .../experimental/test_kernel_dispatcher.py    | 105 ++++++++++++++++++
 2 files changed, 108 insertions(+)
 create mode 100644 numba_dpex/tests/experimental/__init__.py
 create mode 100644 numba_dpex/tests/experimental/test_kernel_dispatcher.py

diff --git a/numba_dpex/tests/experimental/__init__.py b/numba_dpex/tests/experimental/__init__.py
new file mode 100644
index 0000000000..3a217e6325
--- /dev/null
+++ b/numba_dpex/tests/experimental/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
diff --git a/numba_dpex/tests/experimental/test_kernel_dispatcher.py b/numba_dpex/tests/experimental/test_kernel_dispatcher.py
new file mode 100644
index 0000000000..2199ec13af
--- /dev/null
+++ b/numba_dpex/tests/experimental/test_kernel_dispatcher.py
@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import dpctl
+import dpnp
+
+import numba_dpex.experimental as exp_dpex
+from numba_dpex import NdRange, Range, dpjit
+
+
+@exp_dpex.kernel(
+    release_gil=False,
+    no_compile=True,
+    no_cpython_wrapper=True,
+    no_cfunc_wrapper=True,
+)
+def add(a, b, c):
+    c[0] = b[0] + a[0]
+
+
+@exp_dpex.kernel(
+    release_gil=False,
+    no_compile=True,
+    no_cpython_wrapper=True,
+    no_cfunc_wrapper=True,
+)
+def sq(a, b):
+    a[0] = b[0] * b[0]
+
+
+def test_call_kernel_from_cpython():
+    """
+    Tests if we can call a kernel function from CPython using the call_kernel
+    dpjit function.
+    """
+
+    q = dpctl.SyclQueue("gpu")
+    a = dpnp.ones(100, sycl_queue=q)
+    b = dpnp.ones_like(a, sycl_queue=q)
+    c = dpnp.zeros_like(a, sycl_queue=q)
+    r = Range(100)
+    ndr = NdRange(global_size=(100,), local_size=(1,))
+
+    exp_dpex.call_kernel(add, r, a, b, c)
+
+    assert c[0] == b[0] + a[0]
+
+    exp_dpex.call_kernel(add, ndr, a, b, c)
+
+    assert c[0] == b[0] + a[0]
+
+
+def test_call_kernel_from_dpjit():
+    """
+    Tests if we can call a kernel function from a dpjit function using the
+    call_kernel dpjit function.
+    """
+
+    @dpjit
+    def range_kernel_caller(q, a, b, c):
+        r = Range(100)
+        exp_dpex.call_kernel(add, r, a, b, c)
+        return c
+
+    @dpjit
+    def ndrange_kernel_caller(q, a, b, c):
+        gr = Range(100)
+        lr = Range(1)
+        ndr = NdRange(gr, lr)
+        exp_dpex.call_kernel(add, ndr, a, b, c)
+        return c
+
+    q = dpctl.SyclQueue("gpu")
+    a = dpnp.ones(100, sycl_queue=q)
+    b = dpnp.ones_like(a, sycl_queue=q)
+    c = dpnp.zeros_like(a, sycl_queue=q)
+
+    range_kernel_caller(q, a, b, c)
+
+    assert c[0] == b[0] + a[0]
+
+    ndrange_kernel_caller(q, a, b, c)
+
+    assert c[0] == b[0] + a[0]
+
+
+def test_call_multiple_kernels():
+    """
+    Tests if the call_kernel dpjit function supports calling different types of
+    kernel with different number of arguments.
+    """
+    q = dpctl.SyclQueue("gpu")
+    a = dpnp.ones(100, sycl_queue=q)
+    b = dpnp.ones_like(a, sycl_queue=q)
+    c = dpnp.zeros_like(a, sycl_queue=q)
+    r = Range(100)
+
+    exp_dpex.call_kernel(add, r, a, b, c)
+
+    assert c[0] == b[0] + a[0]
+
+    exp_dpex.call_kernel(sq, r, a, c)
+
+    assert a[0] == c[0] * c[0]