Add experimental kernel decorator to numba jit_registry

Author: Diptorup Deb

numba_dpex/experimental/decorators.py

@@ -8,6 +8,7 @@
 import inspect
 
 from numba.core import sigutils
+from numba.core.target_extension import jit_registry, target_registry
 
 from .kernel_dispatcher import KernelDispatcher
 
@@ -78,3 +79,6 @@ def _specialized_kernel_dispatcher(pyfunc):
             "the return type as void explicitly."
         )
     return _kernel_dispatcher(func)
+
+
+jit_registry[target_registry["dpex_kernel"]] = kernel

numba_dpex/tests/experimental/test_exec_queue_inference.py

@@ -6,6 +6,7 @@
 import dpctl
 import dpnp
 import pytest
+from numba.core import config
 
 import numba_dpex.experimental as exp_dpex
 from numba_dpex import Range
@@ -34,14 +35,16 @@ def test_successful_execution_queue_inference():
     c = dpnp.zeros_like(a, sycl_queue=q)
     r = Range(100)
 
-    # FIXME: This test fails unexpectedly if the NUMBA_CAPTURED_ERRORS is set
-    # to "new_style".
-    # Refer: https://github.com/IntelPython/numba-dpex/issues/1195
+    current_captured_error_style = config.CAPTURED_ERRORS
+    config.CAPTURED_ERRORS = "new_style"
+
     try:
         exp_dpex.call_kernel(add, r, a, b, c)
     except:
         pytest.fail("Unexpected error when calling kernel")
 
+    config.CAPTURED_ERRORS = current_captured_error_style
+
     assert c[0] == b[0] + a[0]
 
 
@@ -59,8 +62,6 @@ def test_execution_queue_inference_error():
     c = dpnp.zeros_like(a, sycl_queue=q1)
     r = Range(100)
 
-    from numba.core import config
-
     current_captured_error_style = config.CAPTURED_ERRORS
     config.CAPTURED_ERRORS = "new_style"