diff --git a/CHANGELOG.md b/CHANGELOG.md index 03f209e9f03..2497c5ad6a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -181,6 +181,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 diff --git a/cpp/src/jit/cache.cpp b/cpp/src/jit/cache.cpp index 7c02d59c283..a0d9a6c3deb 100644 --- a/cpp/src/jit/cache.cpp +++ b/cpp/src/jit/cache.cpp @@ -24,6 +24,8 @@ #include #include +#include + namespace cudf { namespace jit { // Get the directory in home to use for storing the cache @@ -121,6 +123,11 @@ named_prog 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); }); diff --git a/cpp/src/jit/cache.h b/cpp/src/jit/cache.h index 673798841b5..666a8dbbbe9 100644 --- a/cpp/src/jit/cache.h +++ b/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 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 **/ named_prog getProgram( @@ -103,7 +102,8 @@ class cudfJitCache { template using umap_str_shptr = std::unordered_map>; - umap_str_shptr kernel_inst_map; + std::unordered_map> + kernel_inst_context_map; umap_str_shptr program_map; /*