atomic_fence implementation with test case

Author: adarshyoga

numba_dpex/experimental/__init__.py

@@ -13,6 +13,7 @@
 from numba_dpex.kernel_api_impl.spirv.dispatcher import SPIRVKernelDispatcher
 
 from ._kernel_dpcpp_spirv_overloads import (
+    _atomic_fence_overloads,
     _atomic_ref_overloads,
     _group_barrier_overloads,
     _index_space_id_overloads,

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_atomic_fence_overloads.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Provides overloads for functions included in kernel_api.atomic_fence
+that generate dpcpp SPIR-V LLVM IR intrinsic function calls.
+"""
+from llvmlite import ir as llvmir
+from numba.core import types
+from numba.extending import intrinsic, overload
+
+from numba_dpex.kernel_api import atomic_fence
+
+from ..target import DPEX_KERNEL_EXP_TARGET_NAME
+from ._spv_atomic_inst_helper import get_memory_semantics_mask, get_scope
+from .spv_fn_declarations import (
+    _SUPPORT_CONVERGENT,
+    get_or_insert_spv_atomic_fence_fn,
+)
+
+
+@intrinsic(target=DPEX_KERNEL_EXP_TARGET_NAME)
+def _intrinsic_atomic_fence(
+    ty_context, ty_spirv_mem_sem_mask, ty_spirv_scope
+):  # pylint: disable=unused-argument
+
+    # Signature of `__spirv_MemoryBarrier` call that is
+    # generated for atomic_fence. It takes two arguments -
+    # scope and memory_semantics_mask.
+    # All arguments have to be of type unsigned int32.
+    sig = types.void(types.uint32, types.uint32)
+
+    def _intrinsic_atomic_fence_gen(
+        context, builder, sig, args
+    ):  # pylint: disable=unused-argument
+        callinst = builder.call(
+            get_or_insert_spv_atomic_fence_fn(builder.module),
+            [
+                builder.trunc(args[1], llvmir.IntType(32)),  # scope
+                builder.trunc(args[0], llvmir.IntType(32)),  # semantics mask
+            ],
+        )
+
+        if _SUPPORT_CONVERGENT:  # pylint: disable=duplicate-code
+            callinst.attributes.add("convergent")
+        callinst.attributes.add("nounwind")
+
+    return (
+        sig,
+        _intrinsic_atomic_fence_gen,
+    )
+
+
+@overload(
+    atomic_fence,
+    prefer_literal=True,
+    target=DPEX_KERNEL_EXP_TARGET_NAME,
+)
+def ol_atomic_fence(memory_order, memory_scope):
+    """SPIR-V overload for
+    :meth:`numba_dpex.kernel_api.atomic_fence`.
+
+    Generates the same LLVM IR instruction as DPC++ for the SYCL
+    `atomic_fence` function.
+    """
+    spirv_memory_semantics_mask = get_memory_semantics_mask(
+        memory_order.literal_value
+    )
+    spirv_scope = get_scope(memory_scope.literal_value)
+
+    def ol_atomic_fence_impl(
+        memory_order, memory_scope
+    ):  # pylint: disable=unused-argument
+        # pylint: disable=no-value-for-parameter
+        return _intrinsic_atomic_fence(spirv_memory_semantics_mask, spirv_scope)
+
+    return ol_atomic_fence_impl

…_overloads/spv_atomic_fn_declarations.py …p_spirv_overloads/spv_fn_declarations.pynumba_dpex/experimental/_kernel_dpcpp_spirv_overloads/spv_atomic_fn_declarations.py renamed to numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/spv_fn_declarations.py numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/spv_atomic_fn_declarations.py renamed to numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/spv_fn_declarations.py

@@ -226,3 +226,56 @@ def get_or_insert_spv_atomic_compare_exchange_fn(
     fn.attributes.add("nounwind")
 
     return fn
+
+
+def get_or_insert_spv_group_barrier_fn(module):
+    """
+    Gets or inserts a declaration for a __spirv_ControlBarrier call into the
+    specified LLVM IR module.
+    """
+    mangled_fn_name = ext_itanium_mangler.mangle_ext(
+        "__spirv_ControlBarrier", [types.uint32, types.uint32, types.uint32]
+    )
+
+    spirv_fn_arg_types = [
+        llvmir.IntType(32),
+        llvmir.IntType(32),
+        llvmir.IntType(32),
+    ]
+
+    fn = cgutils.get_or_insert_function(
+        module,
+        llvmir.FunctionType(llvmir.VoidType(), spirv_fn_arg_types),
+        mangled_fn_name,
+    )
+    fn.calling_convention = CC_SPIR_FUNC
+
+    if _SUPPORT_CONVERGENT:
+        fn.attributes.add("convergent")
+    fn.attributes.add("nounwind")
+
+    return fn
+
+
+def get_or_insert_spv_atomic_fence_fn(module):
+    """
+    Gets or inserts a declaration for a __spirv_MemoryBarrier call into the
+    specified LLVM IR module.
+    """
+    mangled_fn_name = ext_itanium_mangler.mangle_ext(
+        "__spirv_MemoryBarrier", [types.uint32, types.uint32]
+    )
+
+    fn = cgutils.get_or_insert_function(
+        module,
+        llvmir.FunctionType(
+            llvmir.VoidType(), [llvmir.IntType(32), llvmir.IntType(32)]
+        ),
+        mangled_fn_name,
+    )
+    fn.calling_convention = CC_SPIR_FUNC
+    if _SUPPORT_CONVERGENT:
+        fn.attributes.add("convergent")
+    fn.attributes.add("nounwind")
+
+    return fn

numba_dpex/kernel_api/__init__.py

@@ -9,6 +9,7 @@
 numba_dpex.
 """
 
+from .atomic_fence import atomic_fence
 from .atomic_ref import AtomicRef
 from .barrier import group_barrier
 from .index_space_ids import Group, Item, NdItem
@@ -18,6 +19,7 @@
 
 __all__ = [
     "AddressSpace",
+    "atomic_fence",
     "AtomicRef",
     "MemoryOrder",
     "MemoryScope",

numba_dpex/kernel_api/atomic_fence.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: 2023 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Python functions that simulate SYCL's atomic_fence primitives.
+"""
+
+
+def atomic_fence(memory_order, memory_scope):  # pylint: disable=unused-argument
+    """The function for performing memory fence across all work-items.
+    Modeled after ``sycl::atomic_fence`` function.
+    It provides control over re-ordering of memory load
+    and store operations. The ``atomic_fence`` function acts as a
+    fence across all work-items and devices specified by a
+    memory_scope argument.
+
+    Args:
+    memory_order: The memory synchronization order.
+
+    memory_scope: The set of work-items and devices to which
+    the memory ordering constraints apply.
+
+    """

numba_dpex/tests/experimental/kernel_api_overloads/spv_overloads/test_atomic_fence.py

@@ -0,0 +1,43 @@
+import dpnp
+
+import numba_dpex as dpex
+import numba_dpex.experimental as dpex_exp
+from numba_dpex.kernel_api import (
+    AtomicRef,
+    Item,
+    MemoryOrder,
+    MemoryScope,
+    atomic_fence,
+)
+from numba_dpex.tests._helper import skip_windows
+
+
+# TODO: https://github.com/IntelPython/numba-dpex/issues/1308
+@skip_windows
+def test_atomic_fence():
+    """A test for atomic_fence function."""
+
+    @dpex_exp.kernel
+    def _kernel(item: Item, a, b):
+        i = item.get_id(0)
+
+        bref = AtomicRef(b)
+
+        if i == 1:
+            a[i] += 1
+            atomic_fence(MemoryOrder.RELEASE, MemoryScope.DEVICE)
+            bref.store(1)
+        elif i == 0:
+            while not bref.load():
+                continue
+            atomic_fence(MemoryOrder.ACQUIRE, MemoryScope.DEVICE)
+            for idx in range(1, a.size):
+                a[0] += a[idx]
+
+    N = 2
+    a = dpnp.ones(N, dtype=dpnp.int64)
+    b = dpnp.zeros(1, dtype=dpnp.int64)
+
+    dpex_exp.call_kernel(_kernel, dpex.Range(N), a, b)
+
+    assert a[0] == N + 1