Reorganize new features that belong to experimental module.

numba_dpex/core/datamodel/models.py

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from llvmlite import ir as llvmir
 from numba.core import datamodel, types
 from numba.core.datamodel.models import PrimitiveModel, StructModel
 from numba.core.extending import register_model
@@ -12,11 +11,9 @@
 
 from ..types import (
     Array,
-    AtomicRefType,
     DpctlSyclEvent,
     DpctlSyclQueue,
     DpnpNdArray,
-    IntEnumLiteral,
     NdRangeType,
     RangeType,
     USMNdArray,
@@ -227,23 +224,6 @@ def __init__(self, dmm, fe_type):
         ]
         super(NdRangeModel, self).__init__(dmm, fe_type, members)
 
-
-class AtomicRefModel(StructModel):
-    def __init__(self, dmm, fe_type):
-        members = [
-            (
-                "ref",
-                types.CPointer(fe_type.dtype, addrspace=fe_type.address_space),
-            ),
-        ]
-        super(AtomicRefModel, self).__init__(dmm, fe_type, members)
-
-
-class LiteralIntEnumModel(PrimitiveModel):
-    def __init__(self, dmm, fe_type):
-        be_type = llvmir.IntType(fe_type.bitwidth)
-        super(LiteralIntEnumModel, self).__init__(dmm, fe_type, be_type)
-
     @property
     def flattened_field_count(self):
         """Return the number of fields in an instance of a NdRangeModel."""
@@ -285,10 +265,6 @@ def _init_data_model_manager() -> datamodel.DataModelManager:
     # model manager. The dpex_data_model_manager is used by the DpexKernelTarget
     dmm.register(DpctlSyclQueue, SyclQueueModel)
 
-    dmm.register(AtomicRefType, AtomicRefModel)
-
-    dmm.register(IntEnumLiteral, LiteralIntEnumModel)
-
     return dmm
 
 
@@ -314,7 +290,3 @@ def _init_data_model_manager() -> datamodel.DataModelManager:
 
 # Register the NdRangeType type
 register_model(NdRangeType)(NdRangeModel)
-# XXX remove registration of AtomicRef and IntEnumLiteral from default registry
-register_model(AtomicRefType)(AtomicRefModel)
-
-register_model(IntEnumLiteral)(LiteralIntEnumModel)

numba_dpex/core/types/__init__.py

@@ -3,10 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .array_type import Array
-from .dpcpp_types import AtomicRefType
 from .dpctl_types import DpctlSyclEvent, DpctlSyclQueue
 from .dpnp_ndarray_type import DpnpNdArray
-from .literal_intenum_type import IntEnumLiteral
 from .numba_types_short_names import (
     b1,
     bool_,
@@ -35,13 +33,11 @@
 
 __all__ = [
     "Array",
-    "AtomicRefType",
     "DpctlSyclQueue",
     "DpctlSyclEvent",
     "DpnpNdArray",
     "RangeType",
     "NdRangeType",
-    "IntEnumLiteral",
     "USMNdArray",
     "none",
     "boolean",

numba_dpex/core/typing/typeof.py

@@ -8,11 +8,9 @@
 from numba.extending import typeof_impl
 from numba.np import numpy_support
 
-from numba_dpex.core.kernel_interface.dpcpp_iface import AtomicRef
 from numba_dpex.utils import address_space
 
 from ..kernel_interface.indexers import NdRange, Range
-from ..types.dpcpp_types import AtomicRefType
 from ..types.dpctl_types import DpctlSyclEvent, DpctlSyclQueue
 from ..types.dpnp_ndarray_type import DpnpNdArray
 from ..types.range_types import NdRangeType, RangeType
@@ -152,8 +150,3 @@ def typeof_ndrange(val, c):
     Returns: A numba_dpex.core.types.range_types.RangeType instance.
     """
     return NdRangeType(val.global_range.ndim)
-
-
-@typeof_impl.register(AtomicRef)
-def typeof_atomic_ref(val, c):
-    return AtomicRefType(val)

numba_dpex/decorators.py

@@ -16,7 +16,11 @@
 from numba_dpex.core.pipelines.dpjit_compiler import DpjitCompiler
 
 
-def kernel(func_or_sig=None, debug=False, enable_cache=True, *args, **kws):
+def kernel(
+    func_or_sig=None,
+    debug=False,
+    enable_cache=True,
+):
     """A decorator to define a kernel function.
 
     A kernel function is conceptually equivalent to a SYCL kernel function, and
@@ -163,4 +167,3 @@ def dpjit(*args, **kws):
 # add it to the decorator registry, this is so e.g. @overload can look up a
 # JIT function to do the compilation work.
 jit_registry[target_registry["dpex"]] = dpjit
-# jit_registry[target_registry["dpex_kernel"]] = kernel

numba_dpex/experimental/__init__.py

@@ -9,6 +9,9 @@
 from numba.core.imputils import Registry
 
 from .decorators import kernel
+from .dpcpp_iface import AddressSpace, AtomicRef, MemoryOrder, MemoryScope
+from .dpcpp_types import AtomicRefType
+from .literal_intenum_type import IntEnumLiteral
 from .kernel_dispatcher import KernelDispatcher
 from .launcher import call_kernel
 from .models import *
@@ -26,4 +29,14 @@ def dpex_dispatcher_const(context):
     return context.get_dummy_value()
 
 
-__all__ = ["kernel", "KernelDispatcher", "call_kernel"]
+__all__ = [
+    "kernel",
+    "call_kernel",
+    "AddressSpace",
+    "AtomicRef",
+    "AtomicRefType",
+    "IntEnumLiteral",
+    "KernelDispatcher",
+    "MemoryOrder",
+    "MemoryScope",
+]

numba_dpex/core/kernel_interface/dpcpp_iface/__init__.py → numba_dpex/experimental/dpcpp_iface/__init__.py

@@ -1,7 +1,10 @@
-# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
+# SPDX-FileCopyrightText: 2023 Intel Corporation
 #
 # SPDX-License-Identifier: Apache-2.0
 
+"""Python classes that are analogous to dpcpp's SYCL API used to write kernels.
+"""
+
 from .atomic_ref import AtomicRef
 from .memory_enums import AddressSpace, MemoryOrder, MemoryScope
 

numba_dpex/core/kernel_interface/dpcpp_iface/memory_enums.py → numba_dpex/experimental/dpcpp_iface/memory_enums.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from numba_dpex.core.flag_enum import FlagEnum
+from numba_dpex.experimental.flag_enum import FlagEnum
 
 
 class MemoryOrder(FlagEnum):

numba_dpex/core/flag_enum.py → numba_dpex/experimental/flag_enum.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
 from enum import IntEnum
 
 

numba_dpex/core/types/literal_intenum_type.py → numba_dpex/experimental/literal_intenum_type.py

@@ -5,7 +5,7 @@
 from numba.core.types import Integer, Literal
 from numba.core.typing.typeof import typeof
 
-from numba_dpex.core.flag_enum import FlagEnum
+from numba_dpex.experimental.flag_enum import FlagEnum
 
 
 class IntEnumLiteral(Literal, Integer):

numba_dpex/experimental/models.py

@@ -6,15 +6,40 @@
 numba_dpex.experimental module.
 """
 
+from dpcpp_types import AtomicRefType
+from llvmlite import ir as llvmir
+from numba.core import types
 from numba.core.datamodel import models
+from numba.core.datamodel.models import PrimitiveModel, StructModel
 from numba.core.extending import register_model
 
 from numba_dpex.core.datamodel.models import dpex_data_model_manager as dmm
 
+from .literal_intenum_type import IntEnumLiteral
 from .types import KernelDispatcherType
 
+
+class AtomicRefModel(StructModel):
+    def __init__(self, dmm, fe_type):
+        members = [
+            (
+                "ref",
+                types.CPointer(fe_type.dtype, addrspace=fe_type.address_space),
+            ),
+        ]
+        super(AtomicRefModel, self).__init__(dmm, fe_type, members)
+
+
+class LiteralIntEnumModel(PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        be_type = llvmir.IntType(fe_type.bitwidth)
+        super(LiteralIntEnumModel, self).__init__(dmm, fe_type, be_type)
+
+
 # Register the types and datamodel in the DpexKernelTargetContext
 dmm.register(KernelDispatcherType, models.OpaqueModel)
+dmm.register(AtomicRefType, AtomicRefModel)
+dmm.register(IntEnumLiteral, LiteralIntEnumModel)
 
 # Register the types and datamodel in the DpexTargetContext
 register_model(KernelDispatcherType)(models.OpaqueModel)

numba_dpex/experimental/typeof.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from numba.extending import typeof_impl
+
+from .dpcpp_iface import AtomicRef
+
+
+from .dpcpp_types import AtomicRefType
+
+
+@typeof_impl.register(AtomicRef)
+def typeof_atomic_ref(val, c):
+    return AtomicRefType(val)