Merge pull request #1215 from IntelPython/atomic_fence

Implementing atomic_fence in dpex.kernel

Author: Diptorup Deb

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

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_atomic_ref_overloads.py

@@ -32,7 +32,7 @@
     get_memory_semantics_mask,
     get_scope,
 )
-from .spv_atomic_fn_declarations import (
+from .spv_fn_declarations import (
     _SUPPORT_CONVERGENT,
     get_or_insert_atomic_load_fn,
     get_or_insert_spv_atomic_compare_exchange_fn,

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_group_barrier_overloads.py

@@ -3,23 +3,25 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Provides overloads for functions included in kernel_iface.barrier that
+Provides overloads for functions included in kernel_api.barrier that
 generate dpcpp SPIR-V LLVM IR intrinsic function calls.
 """
 
 from llvmlite import ir as llvmir
-from numba.core import cgutils, types
+from numba.core import types
 from numba.core.errors import TypingError
 from numba.extending import intrinsic, overload
 
-from numba_dpex.core import itanium_mangler as ext_itanium_mangler
 from numba_dpex.core.types.kernel_api.index_space_ids import GroupType
 from numba_dpex.experimental.target import DPEX_KERNEL_EXP_TARGET_NAME
 from numba_dpex.kernel_api import group_barrier
 from numba_dpex.kernel_api.memory_enums import MemoryOrder, MemoryScope
 
 from ._spv_atomic_inst_helper import get_memory_semantics_mask, get_scope
-from .spv_atomic_fn_declarations import _SUPPORT_CONVERGENT
+from .spv_fn_declarations import (
+    _SUPPORT_CONVERGENT,
+    get_or_insert_spv_group_barrier_fn,
+)
 
 
 def _get_memory_scope(fence_scope):
@@ -35,7 +37,7 @@ def _intrinsic_barrier(
     ty_mem_scope,  # pylint: disable=unused-argument
     ty_spirv_mem_sem_mask,  # pylint: disable=unused-argument
 ):
-    # Signature of `__spirv_control_barrier` call that is
+    # Signature of `__spirv_ControlBarrier` call that is
     # generated for group_barrier. It takes three arguments -
     # exec_scope, memory_scope and memory_semantics_mask.
     # All arguments have to be of type unsigned int32.
@@ -44,42 +46,17 @@ def _intrinsic_barrier(
     def _intrinsic_barrier_codegen(
         context, builder, sig, args  # pylint: disable=unused-argument
     ):
-        exec_scope_arg = builder.trunc(args[0], llvmir.IntType(32))
-        mem_scope_arg = builder.trunc(args[1], llvmir.IntType(32))
-        spirv_memory_semantics_mask_arg = builder.trunc(
-            args[2], llvmir.IntType(32)
-        )
-
         fn_args = [
-            exec_scope_arg,
-            mem_scope_arg,
-            spirv_memory_semantics_mask_arg,
+            builder.trunc(args[0], llvmir.IntType(32)),
+            builder.trunc(args[1], llvmir.IntType(32)),
+            builder.trunc(args[2], llvmir.IntType(32)),
         ]
 
-        mangled_fn_name = ext_itanium_mangler.mangle_ext(
-            "__spirv_ControlBarrier", [types.uint32, types.uint32, types.uint32]
+        callinst = builder.call(
+            get_or_insert_spv_group_barrier_fn(builder.module), fn_args
         )
 
-        spirv_fn_arg_types = [
-            llvmir.IntType(32),
-            llvmir.IntType(32),
-            llvmir.IntType(32),
-        ]
-
-        fn = cgutils.get_or_insert_function(
-            builder.module,
-            llvmir.FunctionType(llvmir.VoidType(), spirv_fn_arg_types),
-            mangled_fn_name,
-        )
-
-        if _SUPPORT_CONVERGENT:
-            fn.attributes.add("convergent")
-        fn.attributes.add("nounwind")
-        fn.calling_convention = "spir_func"
-
-        callinst = builder.call(fn, fn_args)
-
-        if _SUPPORT_CONVERGENT:
+        if _SUPPORT_CONVERGENT:  # pylint: disable=duplicate-code
             callinst.attributes.add("convergent")
         callinst.attributes.add("nounwind")
 

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