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