Skip to content

Commit 2b37ebb

Browse files
author
Diptorup Deb
committed
Register KernelDispacther with numba's dispatcher_registry.
1 parent 1ab07a8 commit 2b37ebb

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

numba_dpex/experimental/kernel_dispatcher.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from numba.core.compiler import CompileResult
1414
from numba.core.compiler_lock import global_compiler_lock
1515
from numba.core.dispatcher import Dispatcher, _FunctionCompiler
16+
from numba.core.target_extension import dispatcher_registry, target_registry
1617
from numba.core.typing.typeof import Purpose, typeof
1718

1819
from numba_dpex import config, spirv_generator
@@ -315,3 +316,7 @@ def __call__(self, *args, **kw_args):
315316
"""Functor to launch a kernel."""
316317

317318
raise NotImplementedError
319+
320+
321+
_dpex_target = target_registry["dpex_kernel"]
322+
dispatcher_registry[_dpex_target] = KernelDispatcher

0 commit comments

Comments
 (0)