Merge pull request #5219 from devavret/jit-context-cache

[REVIEW] Jit context cache

CHANGELOG.md

@@ -184,6 +184,7 @@
 - PR #5214 Move docs build script into repository
 - PR #5155 Fix cudf documentation misspellings
 - PR #5214 Move docs build script into repository
+- PR #5219 Add per context cache for JIT kernels
 
 ## Bug Fixes
 

cpp/src/jit/cache.cpp

@@ -24,6 +24,8 @@
 #include <unistd.h>
 #include <boost/filesystem.hpp>
 
+#include <cuda.h>
+
 namespace cudf {
 namespace jit {
 // Get the directory in home to use for storing the cache
@@ -121,6 +123,11 @@ named_prog<jitify::experimental::KernelInstantiation> cudfJitCache::getKernelIns
   std::string kern_inst_name = prog_name + '.' + kern_name;
   for (auto&& arg : arguments) kern_inst_name += '_' + arg;
 
+  CUcontext c;
+  cuCtxGetCurrent(&c);
+
+  auto& kernel_inst_map = kernel_inst_context_map[c];
+
   return getCached(kern_inst_name, kernel_inst_map, [&]() {
     return program.kernel(kern_name).instantiate(arguments);
   });

cpp/src/jit/cache.h

@@ -67,9 +67,9 @@ class cudfJitCache {
    * Searches an internal in-memory cache and file based cache for the kernel
    * and if not found, JIT compiles and returns the kernel
    *
-   * @param kern_name [in] name of kernel to return
-   * @param program   [in] Jitify preprocessed program to get the kernel from
-   * @param arguments [in] template arguments for kernel in vector of strings
+   * @param kern_name  name of kernel to return
+   * @param program    Jitify preprocessed program to get the kernel from
+   * @param arguments  template arguments for kernel in vector of strings
    * @return  Pair of string kernel identifier and compiled kernel object
    **/
   named_prog<jitify::experimental::KernelInstantiation> getKernelInstantiation(
@@ -83,13 +83,12 @@ class cudfJitCache {
    * Searches an internal in-memory cache and file based cache for the Jitify
    * pre-processed program and if not found, JIT processes and returns it
    *
-   * @param prog_file_name [in] name of program to return
-   * @param cuda_source    [in] string source code of program to compile
-   * @param given_headers  [in] vector of strings representing source or names of
-   *  each header included in cuda_source
-   * @param given_options  [in] vector of strings options to pass to NVRTC
-   * @param file_callback  [in] pointer to callback function to call whenever a
-   *  header needs to be loaded
+   * @param prog_file_name name of program to return
+   * @param cuda_source    string source code of program to compile
+   * @param given_headers  vector of strings representing source or names of each header included in
+   *                       cuda_source
+   * @param given_options  vector of strings options to pass to NVRTC
+   * @param file_callback  pointer to callback function to call whenever a header needs to be loaded
    * @return named_prog<jitify::experimental::Program>
    **/
   named_prog<jitify::experimental::Program> getProgram(
@@ -103,7 +102,8 @@ class cudfJitCache {
   template <typename Tv>
   using umap_str_shptr = std::unordered_map<std::string, std::shared_ptr<Tv>>;
 
-  umap_str_shptr<jitify::experimental::KernelInstantiation> kernel_inst_map;
+  std::unordered_map<CUcontext, umap_str_shptr<jitify::experimental::KernelInstantiation>>
+    kernel_inst_context_map;
   umap_str_shptr<jitify::experimental::Program> program_map;
 
   /*