Register KernelDispacther with numba's dispatcher_registry.

IntelPython · Nov 6, 2023 · 3f59033 · 3f59033
1 parent 54ba2ea
commit 3f59033
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/numba_dpex/experimental/dpcpp_iface/atomic_ref.py b/numba_dpex/experimental/dpcpp_iface/atomic_ref.py
@@ -9,7 +9,6 @@
 
 from numba_dpex.core import itanium_mangler as ext_itanium_mangler
 from numba_dpex.core.datamodel.models import dpex_data_model_manager
-from numba_dpex.core.exceptions import UnreachableError
 from numba_dpex.core.targets.kernel_target import DPEX_KERNEL_TARGET_NAME
 from numba_dpex.core.types import USMNdArray
 

diff --git a/numba_dpex/experimental/kernel_dispatcher.py b/numba_dpex/experimental/kernel_dispatcher.py
@@ -13,6 +13,7 @@
 from numba.core.compiler import CompileResult
 from numba.core.compiler_lock import global_compiler_lock
 from numba.core.dispatcher import Dispatcher, _FunctionCompiler
+from numba.core.target_extension import dispatcher_registry, target_registry
 from numba.core.typing.typeof import Purpose, typeof
 
 from numba_dpex import config, spirv_generator
@@ -315,3 +316,7 @@ def __call__(self, *args, **kw_args):
         """Functor to launch a kernel."""
 
         raise NotImplementedError
+
+
+_dpex_target = target_registry["dpex_kernel"]
+dispatcher_registry[_dpex_target] = KernelDispatcher