From f1137835193e52fa4c9cb232f5b9ff6c9e63c6b3 Mon Sep 17 00:00:00 2001
From: Devavret Makkar <dmakkar@nvidia.com>
Date: Tue, 12 May 2020 18:57:30 +0530
Subject: [PATCH 1/2] Add context level cache for kernels

---
 cpp/src/jit/cache.cpp |  7 +++++++
 cpp/src/jit/cache.h   | 22 +++++++++++-----------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/cpp/src/jit/cache.cpp b/cpp/src/jit/cache.cpp
index a0a0750f03d..1e0f8ce9217 100644
--- a/cpp/src/jit/cache.cpp
+++ b/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);
   });
diff --git a/cpp/src/jit/cache.h b/cpp/src/jit/cache.h
index d102ffa08c5..389e42208f9 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<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;
 
   /*

From 5d83e2e939120b1075e87b9b816dfdf19bdab833 Mon Sep 17 00:00:00 2001
From: Devavret Makkar <dmakkar@nvidia.com>
Date: Tue, 19 May 2020 02:39:57 +0530
Subject: [PATCH 2/2] changelog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

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