Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[REVIEW] Jit context cache #5219

Merged
merged 3 commits into from
May 19, 2020
Merged

[REVIEW] Jit context cache #5219

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f113783
Add context level cache for kernels
devavret May 12, 2020
90c39b9
Merge branch 'branch-0.14' into jit-context-cache
devavret May 18, 2020
5d83e2e
changelog
devavret May 18, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
7 changes: 7 additions & 0 deletions cpp/src/jit/cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
});
Expand Down
22 changes: 11 additions & 11 deletions cpp/src/jit/cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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;

/*
Expand Down