Skip to content

Commit

Permalink
Reenable serialisation
Browse files Browse the repository at this point in the history 
Triggering compile of the preprocessed source after deserialisation is still fast.
  • Loading branch information
@Robadob
Robadob committed Nov 17, 2023
1 parent febcb51 commit d2e8f46
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 56 deletions.
7 changes: 4 additions & 3 deletions include/flamegpu/detail/JitifyCache.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

namespace jitify2 {
class KernelData;
class PreprocessedProgramData;
} // namespace jitify2

namespace flamegpu {
Expand All @@ -30,7 +31,7 @@ class JitifyCache {
// dynamic header concatenated to kernel
// We check this is an exact match before loading from cache
std::string long_reference;
std::string serialised_kernel;
std::string serialised_program;
};

public:
Expand Down Expand Up @@ -89,9 +90,9 @@ class JitifyCache {
* In the case of FLAME GPU 2, these args are likely to be the user defined function_impl and the message i/o types.
* @param kernel_src Source code for the user defined agent function/condition
* @param dynamic_header Dynamic header source generated by curve rtc
* @return A jitify RTC kernel instance of the provided kernel sources
* @return A jitify preprocessed program instance of the provided kernel sources
*/
static std::unique_ptr<jitify2::KernelData> compileKernel(
static std::unique_ptr<jitify2::PreprocessedProgramData> preprocessKernel(
const std::string &func_name,
const std::vector<std::string> &template_args,
const std::string &kernel_src,
Expand Down
2 changes: 1 addition & 1 deletion include/flamegpu/detail/compute_capability.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ std::vector<int> getNVRTCSupportedComputeCapabilties();
* Get the best matching compute capability from a vector of compute capabililties in ascending order
* I.e. get the maximum CC value which is less than or equal to the target CC
*
* This method has been separated from JitifyCache::compileKernel so that it can be tested generically, without having to write tests which are relative to the linked nvrtc and/or the current device.
* This method has been separated from JitifyCache::preprocessKernel so that it can be tested generically, without having to write tests which are relative to the linked nvrtc and/or the current device.
*
* @param target compute capability to find the best match for
* @param archictectures a vector of architectures in ascending order
Expand Down
103 changes: 51 additions & 52 deletions src/flamegpu/detail/JitifyCache.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ std::vector<std::filesystem::path> getIncludeDirs() {
}
return rtn;
}
/* Temp comment to suppress warning in CUDA 11.0 CI

std::string loadFile(const std::filesystem::path &filepath) {
std::ifstream ifs;
ifs.open(filepath, std::ifstream::binary);
Expand All @@ -97,7 +97,7 @@ std::string loadFile(const std::filesystem::path &filepath) {
ifs.close();
return rtn;
}
*/

/**
* Find the cuda include directory.
* Throws exceptions if it can not be found.
Expand Down Expand Up @@ -308,8 +308,8 @@ bool confirmFLAMEGPUHeaderVersion(const std::string &flamegpuIncludeDir, const s
} // namespace

std::mutex JitifyCache::instance_mutex;
std::unique_ptr<jitify2::KernelData> JitifyCache::compileKernel(const std::string &func_name, const std::vector<std::string> &template_args, const std::string &kernel_src, const std::string &dynamic_header) {
flamegpu::util::nvtx::Range range{"JitifyCache::compileKernel"};
std::unique_ptr<jitify2::PreprocessedProgramData> JitifyCache::preprocessKernel(const std::string &func_name, const std::vector<std::string> &template_args, const std::string &kernel_src, const std::string &dynamic_header) {
flamegpu::util::nvtx::Range range{"JitifyCache::preprocessKernel"};
// find and validate the cuda include directory via CUDA_PATH or CUDA_HOME.
static const std::string cuda_include_dir = getCUDAIncludeDir();
// find and validate the the flamegpu include directory
Expand Down Expand Up @@ -423,39 +423,7 @@ std::unique_ptr<jitify2::KernelData> JitifyCache::compileKernel(const std::strin
"in JitifyCache::buildProgram().",
func_name.c_str());
}
// Build the name of the template configuration to be instantiated
std::stringstream name_expression;
if (template_args.size() == 1) {
name_expression << "flamegpu::agent_function_condition_wrapper<";
name_expression << template_args[0];
name_expression << ">";
} else if (template_args.size() == 3) {
name_expression << "flamegpu::agent_function_wrapper<";
name_expression << template_args[0] << "," << template_args[1] << "," << template_args[2];
name_expression << ">";
} else {
THROW exception::UnknownInternalError("Unexpected AgentFunction template arg count!");
}
jitify2::LoadedProgram loaded_program = program->load({ name_expression.str() });
if (!loaded_program.ok()) {
const jitify2::ErrorMsg &compile_error = loaded_program.error();
fprintf(stderr, "Failed to load program for agent function (condition) '%s', log:\n%s",
func_name.c_str(), compile_error.c_str());
THROW exception::InvalidAgentFunc("Error loading agent function (or function condition) ('%s'): function had compilation errors (see std::cout), "
"in JitifyCache::buildProgram().",
func_name.c_str());
}
jitify2::Kernel loaded_kernel = loaded_program->get_kernel(name_expression.str());
if (loaded_kernel.ok()) {
return std::make_unique<jitify2::KernelData>(loaded_kernel.value());
} else {
const jitify2::ErrorMsg &compile_error = loaded_kernel.error();
fprintf(stderr, "Failed to compile and link agent function (condition) '%s', log:\n%s",
func_name.c_str(), compile_error.c_str());
THROW exception::InvalidAgentFunc("Error compiling runtime agent function (or function condition) ('%s'): function had compilation errors (see std::cout), "
"in JitifyCache::buildProgram().",
func_name.c_str());
}
return std::make_unique<jitify2::PreprocessedProgramData>(program.value());
}
void JitifyCache::getKnownHeaders(std::vector<std::string>& headers) {
// Add known headers from hierarchy
Expand Down Expand Up @@ -561,16 +529,17 @@ std::unique_ptr<jitify2::KernelData> JitifyCache::loadKernel(const std::string &
#endif
// Use jitify hash methods for consistent hashing between OSs
jitify2::detail::sha256(kernel_src + dynamic_header);
std::unique_ptr<jitify2::PreprocessedProgramData> program;
// Does a copy with the right reference exist in memory?
/*if (use_memory_cache) {
if (use_memory_cache) {
const auto it = cache.find(short_reference);
if (it != cache.end()) {
// Check long reference
if (it->second.long_reference == long_reference) {
// Deserialize and return program
jitify2::Kernel prog = jitify2::Kernel::deserialize(it->second.serialised_kernel);
jitify2::PreprocessedProgram prog = jitify2::PreprocessedProgram::deserialize(it->second.serialised_program);
if (prog.ok()) {
return std::make_unique<jitify2::KernelData>(prog.value());
program = std::make_unique<jitify2::PreprocessedProgramData>(prog.value());
}
// Fail silently and try to build code
}
Expand All @@ -579,7 +548,7 @@ std::unique_ptr<jitify2::KernelData> JitifyCache::loadKernel(const std::string &
// Does a copy with the right reference exist on disk?
const std::filesystem::path cache_file = getTMP() / short_reference;
const std::filesystem::path reference_file = cache_file.parent_path() / std::filesystem::path(cache_file.filename().string() + ".ref");
if (use_disk_cache && std::filesystem::exists(cache_file)) {
if (!program && use_disk_cache && std::filesystem::exists(cache_file)) {
// Load the long reference for the cache file
const std::string file_long_reference = loadFile(reference_file);
if (file_long_reference == long_reference) {
Expand All @@ -589,29 +558,28 @@ std::unique_ptr<jitify2::KernelData> JitifyCache::loadKernel(const std::string &
// Add it to cache for later loads
cache.emplace(short_reference, CachedProgram{long_reference, serialised_kernelinst});
// Deserialize and return program
jitify2::Kernel prog = jitify2::Kernel::deserialize(serialised_kernelinst);
jitify2::PreprocessedProgram prog = jitify2::PreprocessedProgram::deserialize(serialised_kernelinst);
if (prog.ok()) {
return std::make_unique<jitify2::KernelData>(prog.value());
program = std::make_unique<jitify2::PreprocessedProgramData>(prog.value());
}
// Fail silently and try to build code
}
}
}*/
}
// Kernel has not yet been cached
{
if (!program) {
// Build kernel
std::unique_ptr<jitify2::KernelData> kernel = compileKernel(func_name, template_args, kernel_src, dynamic_header);
/*
program = preprocessKernel(func_name, template_args, kernel_src, dynamic_header);
// Add it to cache for later loads
const std::string serialised_kernel = use_memory_cache || use_disk_cache ? kernel->serialize() : "";
const std::string serialised_program = use_memory_cache || use_disk_cache ? program->serialize() : "";
if (use_memory_cache) {
cache.emplace(short_reference, CachedProgram{long_reference, serialised_kernel });
cache.emplace(short_reference, CachedProgram{long_reference, serialised_program });
}
// Save it to disk
if (use_disk_cache) {
std::ofstream ofs(cache_file, std::ofstream::out | std::ofstream::binary | std::ofstream::trunc);
if (ofs) {
ofs << serialised_kernel;
ofs << serialised_program;
ofs.close();
}
ofs = std::ofstream(reference_file, std::ofstream::out | std::ofstream::binary | std::ofstream::trunc);
Expand All @@ -620,9 +588,40 @@ std::unique_ptr<jitify2::KernelData> JitifyCache::loadKernel(const std::string &
ofs.close();
}
}
*/
return kernel;
}
// Build the name of the template configuration to be instantiated
std::stringstream name_expression;
if (template_args.size() == 1) {
name_expression << "flamegpu::agent_function_condition_wrapper<";
name_expression << template_args[0];
name_expression << ">";
} else if (template_args.size() == 3) {
name_expression << "flamegpu::agent_function_wrapper<";
name_expression << template_args[0] << "," << template_args[1] << "," << template_args[2];
name_expression << ">";
} else {
THROW exception::UnknownInternalError("Unexpected AgentFunction template arg count!");
}
// Actually trigger compilation
jitify2::LoadedProgram loaded_program = program->load({ name_expression.str() });
if (!loaded_program.ok()) {
const jitify2::ErrorMsg& compile_error = loaded_program.error();
fprintf(stderr, "Failed to load program for agent function (condition) '%s', log:\n%s",
func_name.c_str(), compile_error.c_str());
THROW exception::InvalidAgentFunc("Error loading agent function (or function condition) ('%s'): function had compilation errors (see std::cout), "
"in JitifyCache::loadKernel().",
func_name.c_str());
}
jitify2::Kernel loaded_kernel = loaded_program->get_kernel(name_expression.str());
if (loaded_kernel.ok()) {
return std::make_unique<jitify2::KernelData>(loaded_kernel.value());
}
const jitify2::ErrorMsg& compile_error = loaded_kernel.error();
fprintf(stderr, "Failed to compile and link agent function (condition) '%s', log:\n%s",
func_name.c_str(), compile_error.c_str());
THROW exception::InvalidAgentFunc("Error compiling runtime agent function (or function condition) ('%s'): function had compilation errors (see std::cout), "
"in JitifyCache::loadKernel().",
func_name.c_str());
}
void JitifyCache::useMemoryCache(bool yesno) {
std::lock_guard<std::mutex> lock(cache_mutex);
Expand Down

0 comments on commit d2e8f46

Please sign in to comment.