Caching kernel_bundle after create_program_from_spirv()

Store exec_queue along with kernel_bundle Properly initialize self._kernel_bundle_cache with NullCache() Save kernel_bundle with exec_queue as a key
IntelPython · Feb 13, 2023 · f4d67dc · f4d67dc
1 parent 8a891bb
commit f4d67dc
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 8 deletions.
diff --git a/numba_dpex/core/caching.py b/numba_dpex/core/caching.py
@@ -13,7 +13,9 @@
 from numba_dpex.core.types import USMNdArray
 
 
-def build_key(argtypes, pyfunc, codegen, backend=None, device_type=None):
+def build_key(
+    argtypes, pyfunc, codegen, backend=None, device_type=None, exec_queue=None
+):
     """Constructs a key from python function, context, backend and the device
     type.
 
@@ -64,6 +66,7 @@ def build_key(argtypes, pyfunc, codegen, backend=None, device_type=None):
         codegen.magic_tuple(),
         backend,
         device_type,
+        exec_queue,
         (
             hashlib.sha256(codebytes).hexdigest(),
             hashlib.sha256(cvarbytes).hexdigest(),

diff --git a/numba_dpex/core/kernel_interface/dispatcher.py b/numba_dpex/core/kernel_interface/dispatcher.py
@@ -88,14 +88,21 @@ def __init__(
         # caching related attributes
         if not config.ENABLE_CACHE:
             self._cache = NullCache()
+            self._kernel_bundle_cache = NullCache()
         elif enable_cache:
             self._cache = LRUCache(
                 name="SPIRVKernelCache",
                 capacity=config.CACHE_SIZE,
                 pyfunc=self.pyfunc,
             )
+            self._kernel_bundle_cache = LRUCache(
+                name="KernelBundleCache",
+                capacity=config.CACHE_SIZE,
+                pyfunc=self.pyfunc,
+            )
         else:
             self._cache = NullCache()
+            self._kernel_bundle_cache = NullCache()
         self._cache_hits = 0
 
         if array_access_specifiers:
@@ -627,12 +634,26 @@ def __call__(self, *args):
                     cache=self._cache,
                 )
 
-        # create a sycl::KernelBundle
-        kernel_bundle = dpctl_prog.create_program_from_spirv(
-            exec_queue,
-            device_driver_ir_module,
-            " ".join(self._create_sycl_kernel_bundle_flags),
+        kernel_bundle_key = build_key(
+            tuple(argtypes),
+            self.pyfunc,
+            dpex_kernel_target.target_context.codegen(),
+            exec_queue=exec_queue,
         )
+
+        artifact = self._kernel_bundle_cache.get(kernel_bundle_key)
+
+        if artifact is None:
+            # create a sycl::KernelBundle
+            kernel_bundle = dpctl_prog.create_program_from_spirv(
+                exec_queue,
+                device_driver_ir_module,
+                " ".join(self._create_sycl_kernel_bundle_flags),
+            )
+            self._kernel_bundle_cache.put(kernel_bundle_key, kernel_bundle)
+        else:
+            kernel_bundle = artifact
+
         #  get the sycl::kernel
         sycl_kernel = kernel_bundle.get_sycl_kernel(kernel_module_name)
 

diff --git a/numba_dpex/tests/test_device_array_args.py b/numba_dpex/tests/test_device_array_args.py
@@ -26,7 +26,7 @@ def data_parallel_sum(a, b, c):
 
 
 @skip_no_opencl_cpu
-class TestArrayArgsGPU:
+class TestArrayArgsCPU:
     def test_device_array_args_cpu(self):
         c = np.ones_like(a)
 
@@ -37,7 +37,7 @@ def test_device_array_args_cpu(self):
 
 
 @skip_no_opencl_gpu
-class TestArrayArgsCPU:
+class TestArrayArgsGPU:
     def test_device_array_args_gpu(self):
         c = np.ones_like(a)