From b3b19d8345c508c2e4e119b0c4fe699c27929435 Mon Sep 17 00:00:00 2001
From: Diptorup Deb <diptorup.deb@intel.com>
Date: Tue, 26 Dec 2023 22:14:14 -0600
Subject: [PATCH] Unit tests to evaluate AtomicRef and fetch_add

---
 .../kernel_iface/spv_overloads/__init__.py    |  3 ++
 .../spv_overloads/test_atomic_fetch_phi.py    | 40 +++++++++++++++++
 .../spv_overloads/test_atomic_ref.py          | 45 +++++++++++++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 numba_dpex/tests/experimental/kernel_iface/spv_overloads/__init__.py
 create mode 100644 numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_fetch_phi.py
 create mode 100644 numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_ref.py

diff --git a/numba_dpex/tests/experimental/kernel_iface/spv_overloads/__init__.py b/numba_dpex/tests/experimental/kernel_iface/spv_overloads/__init__.py
new file mode 100644
index 0000000000..3a217e6325
--- /dev/null
+++ b/numba_dpex/tests/experimental/kernel_iface/spv_overloads/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
diff --git a/numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_fetch_phi.py b/numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_fetch_phi.py
new file mode 100644
index 0000000000..b52e8763e8
--- /dev/null
+++ b/numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_fetch_phi.py
@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import dpnp
+import pytest
+
+import numba_dpex as dpex
+import numba_dpex.experimental as dpex_exp
+from numba_dpex.experimental.kernel_iface import AtomicRef
+from numba_dpex.tests._helper import get_all_dtypes
+
+list_of_supported_dtypes = get_all_dtypes(
+    no_bool=True, no_float16=True, no_none=True, no_complex=True
+)
+
+
+@pytest.fixture(params=list_of_supported_dtypes)
+def input_arrays(request):
+    # The size of input and out arrays to be used
+    N = 10
+    a = dpnp.ones(N, dtype=request.param)
+    b = dpnp.zeros(N, dtype=request.param)
+    return a, b
+
+
+def test_fetch_add(input_arrays):
+    @dpex_exp.kernel
+    def atomic_ref_kernel(a, b, ref_index):
+        i = dpex.get_global_id(0)
+        v = AtomicRef(b, index=ref_index)
+        v.fetch_add(a[i])
+
+    a, b = input_arrays
+    ref_index = 0
+
+    dpex_exp.call_kernel(atomic_ref_kernel, dpex.Range(10), a, b, ref_index)
+
+    # Verify that `a` was accumulated at b[ref_index]
+    assert b[ref_index] == 10
diff --git a/numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_ref.py b/numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_ref.py
new file mode 100644
index 0000000000..fce6ad3477
--- /dev/null
+++ b/numba_dpex/tests/experimental/kernel_iface/spv_overloads/test_atomic_ref.py
@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import dpnp
+import pytest
+from numba.core.errors import TypingError
+
+import numba_dpex as dpex
+import numba_dpex.experimental as dpex_exp
+from numba_dpex.experimental.kernel_iface import AddressSpace, AtomicRef
+
+
+def test_atomic_ref_compilation():
+    @dpex_exp.kernel
+    def atomic_ref_kernel(a, b):
+        i = dpex.get_global_id(0)
+        v = AtomicRef(b, index=0, address_space=AddressSpace.GLOBAL)
+        v.fetch_add(a[i])
+
+    a = dpnp.ones(10)
+    b = dpnp.zeros(10)
+    try:
+        dpex_exp.call_kernel(atomic_ref_kernel, dpex.Range(10), a, b)
+    except Exception:
+        pytest.fail("Unexpected execution failure")
+
+
+def test_atomic_ref_compilation_failure():
+    """A negative test that verifies that a TypingError is raised if we try to
+    create an AtomicRef in the local address space from a global address space
+    ref.
+    """
+
+    @dpex_exp.kernel
+    def atomic_ref_kernel(a, b):
+        i = dpex.get_global_id(0)
+        v = AtomicRef(b, index=0, address_space=AddressSpace.LOCAL)
+        v.fetch_add(a[i])
+
+    a = dpnp.ones(10)
+    b = dpnp.zeros(10)
+
+    with pytest.raises(TypingError):
+        dpex_exp.call_kernel(atomic_ref_kernel, dpex.Range(10), a, b)