-
Notifications
You must be signed in to change notification settings - Fork 916
/
cache.h
206 lines (183 loc) · 8.04 KB
/
cache.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
/*
* Copyright (c) 2019, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef CUDF_JIT_CACHE_H_
#define CUDF_JIT_CACHE_H_
#include <cudf/utilities/error.hpp>
#include <jitify.hpp>
#include <unordered_map>
#include <string>
#include <memory>
#include <mutex>
#include <boost/filesystem.hpp>
namespace cudf {
namespace jit {
template <typename Tv>
using named_prog = std::pair<std::string, std::shared_ptr<Tv>>;
/**
* @brief Get the string path to the JITIFY kernel cache directory.
*
* This path can be overridden at runtime by defining an environment variable
* named `LIBCUDF_KERNEL_CACHE_PATH`. The value of this variable must be a path
* under which the process' user has read/write priveleges.
*
* This function returns a path to the cache directory, creating it if it
* doesn't exist.
*
* The default cache directory `$TEMPDIR/cudf_$CUDF_VERSION`.
**/
boost::filesystem::path getCacheDir();
class cudfJitCache
{
public:
/**---------------------------------------------------------------------------*
* @brief Get a process wide singleton cache object
*
*---------------------------------------------------------------------------**/
static cudfJitCache& Instance() {
// Meyers' singleton is thread safe in C++11
// Link: https://stackoverflow.com/a/1661564
static cudfJitCache cache;
return cache;
}
cudfJitCache();
~cudfJitCache();
/**---------------------------------------------------------------------------*
* @brief Get the Kernel Instantiation object
*
* 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
* @return Pair of string kernel identifier and compiled kernel object
*---------------------------------------------------------------------------**/
named_prog<jitify::experimental::KernelInstantiation> getKernelInstantiation(
std::string const& kern_name,
named_prog<jitify::experimental::Program> const& program,
std::vector<std::string> const& arguments);
/**---------------------------------------------------------------------------*
* @brief Get the Jitify preprocessed Program object
*
* 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
* @return named_prog<jitify::experimental::Program>
*---------------------------------------------------------------------------**/
named_prog<jitify::experimental::Program> getProgram(
std::string const& prog_file_name,
std::string const& cuda_source = "",
std::vector<std::string> const& given_headers = {},
std::vector<std::string> const& given_options = {},
jitify::experimental::file_callback_type file_callback = nullptr);
private:
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;
umap_str_shptr<jitify::experimental::Program> program_map;
/*
Even though this class can be used as a non-singleton, the file cache
access should remain limited to one thread per process. The lockf locks can
prevent multiple processes from accessing the file but are ineffective in
preventing multiple threads from doing so as the lock is shared by the
entire process.
Therefore the mutexes are static.
*/
static std::mutex _kernel_cache_mutex;
static std::mutex _program_cache_mutex;
private:
/**---------------------------------------------------------------------------*
* @brief Class to allow process wise exclusive access to cache files
*
*---------------------------------------------------------------------------**/
class cacheFile
{
private:
std::string _file_name;
bool successful_read = false;
bool successful_write = false;
public:
cacheFile(std::string file_name);
~cacheFile();
/**---------------------------------------------------------------------------*
* @brief Read this file and return the contents as a std::string
*
*---------------------------------------------------------------------------**/
std::string read();
/**---------------------------------------------------------------------------*
* @brief Write the passed string to this file
*
*---------------------------------------------------------------------------**/
void write(std::string);
/**---------------------------------------------------------------------------*
* @brief Check whether the read() operation on the file completed successfully
*
* @return true Read was successful. String returned by `read()` is valid
* @return false Read was unsuccessful. String returned by `read()` is empty
*---------------------------------------------------------------------------**/
bool is_read_successful() { return successful_read; }
/**---------------------------------------------------------------------------*
* @brief Check whether the write() operation on the file completed successfully
*
* @return true Write was successful.
* @return false Write was unsuccessful. File state is undefined
*---------------------------------------------------------------------------**/
bool is_write_successful() { return successful_write; }
};
private:
template <typename T, typename FallbackFunc>
named_prog<T> getCached(
std::string const& name,
umap_str_shptr<T>& map,
FallbackFunc func) {
// Find memory cached T object
auto it = map.find(name);
if ( it != map.end()) {
return std::make_pair(name, it->second);
}
else { // Find file cached T object
bool successful_read = false;
std::string serialized;
#if defined(JITIFY_USE_CACHE)
boost::filesystem::path file_name = getCacheDir() / name;
cacheFile file{file_name.string()};
serialized = file.read();
successful_read = file.is_read_successful();
#endif
if (not successful_read) {
// JIT compile and write to file if possible
serialized = func().serialize();
#if defined(JITIFY_USE_CACHE)
file.write(serialized);
#endif
}
// Add deserialized T to cache and return
auto program = std::make_shared<T>(T::deserialize(serialized));
map[name] = program;
return std::make_pair(name, program);
}
}
};
} // namespace jit
} // namespace cudf
#endif // CUDF_JIT_CACHE_H_