diff --git a/taichi/runtime/llvm/llvm_offline_cache.cpp b/taichi/runtime/llvm/llvm_offline_cache.cpp index 13c727ebc177b..799ca20484867 100644 --- a/taichi/runtime/llvm/llvm_offline_cache.cpp +++ b/taichi/runtime/llvm/llvm_offline_cache.cpp @@ -26,22 +26,31 @@ constexpr char kMetadataFilename[] = "metadata"; std::unique_ptr LlvmOfflineCacheFileReader::make( const std::string &path, LlvmOfflineCache::Format format) { + LlvmOfflineCache data; + if (!load_meta_data(data, path)) { + return nullptr; + } + return std::unique_ptr( + new LlvmOfflineCacheFileReader(path, std::move(data), format)); +} + +bool LlvmOfflineCacheFileReader::load_meta_data( + LlvmOfflineCache &data, + const std::string &cache_file_path) { std::stringstream tcb_ss; - tcb_ss << path << "/" << kMetadataFilename << ".tcb"; + tcb_ss << cache_file_path << "/" << kMetadataFilename << ".tcb"; const auto tcb_path = tcb_ss.str(); { // No the best way to check for filepath existence, but whatever... See // https://stackoverflow.com/questions/12774207/fastest-way-to-check-if-a-file-exists-using-standard-c-c11-14-17-c std::ifstream fs(tcb_path, std::ios::in | std::ios::binary); if (!fs.good()) { - TI_DEBUG("LLVM cache {} does not exist", path); - return nullptr; + TI_DEBUG("LLVM cache {} does not exist", cache_file_path); + return false; } } - LlvmOfflineCache data; read_from_binary_file(data, tcb_path); - return std::unique_ptr( - new LlvmOfflineCacheFileReader(path, std::move(data), format)); + return true; } LlvmOfflineCacheFileReader::LlvmOfflineCacheFileReader( @@ -111,7 +120,8 @@ std::unique_ptr LlvmOfflineCacheFileReader::load_module( } void LlvmOfflineCacheFileWriter::dump(const std::string &path, - LlvmOfflineCache::Format format) { + LlvmOfflineCache::Format format, + bool merge_with_old) { taichi::create_directories(path); for (auto &[k, v] : data_.kernels) { std::stringstream filename_ss; @@ -149,6 +159,14 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path, } } { + // Merge with old metadata + if (merge_with_old) { + LlvmOfflineCache old_data; + if (LlvmOfflineCacheFileReader::load_meta_data(old_data, path)) { + add_data(std::move(old_data)); + } + } + // Dump metadata std::stringstream prefix_ss; prefix_ss << path << "/" << kMetadataFilename; const std::string file_prefix = prefix_ss.str(); @@ -160,6 +178,23 @@ void LlvmOfflineCacheFileWriter::dump(const std::string &path, } } +void LlvmOfflineCacheFileWriter::add_data(LlvmOfflineCache &&data) { + // Note: merge this->data_ with data, new cover old + auto &new_kernels = data_.kernels; + auto &new_fields = data_.fields; + auto &old_kernels = data.kernels; + auto &old_fields = data.fields; + + for (auto &[k, v] : new_fields) { + old_fields[k] = std::move(v); + } + for (auto &[k, v] : new_kernels) { + old_kernels[k] = std::move(v); + } + + data_ = std::move(data); +} + void LlvmOfflineCacheFileWriter::mangle_offloaded_task_name( const std::string &kernel_key, llvm::Module *module, diff --git a/taichi/runtime/llvm/llvm_offline_cache.h b/taichi/runtime/llvm/llvm_offline_cache.h index c8d9b6090b1ba..89b23628c62df 100644 --- a/taichi/runtime/llvm/llvm_offline_cache.h +++ b/taichi/runtime/llvm/llvm_offline_cache.h @@ -111,6 +111,9 @@ class LlvmOfflineCacheFileReader { const std::string &path, LlvmOfflineCache::Format format = LlvmOfflineCache::Format::LL); + static bool load_meta_data(LlvmOfflineCache &data, + const std::string &cache_file_path); + private: LlvmOfflineCacheFileReader(const std::string &path, LlvmOfflineCache &&data, @@ -138,13 +141,16 @@ class LlvmOfflineCacheFileWriter { } void dump(const std::string &path, - LlvmOfflineCache::Format format = LlvmOfflineCache::Format::LL); + LlvmOfflineCache::Format format = LlvmOfflineCache::Format::LL, + bool merge_with_old = false); void set_no_mangle() { mangled_ = true; } private: + void add_data(LlvmOfflineCache &&data); + void mangle_offloaded_task_name( const std::string &kernel_key, llvm::Module *module, diff --git a/taichi/runtime/program_impls/llvm/llvm_program.cpp b/taichi/runtime/program_impls/llvm/llvm_program.cpp index 078f437fc7a67..0176971541912 100644 --- a/taichi/runtime/program_impls/llvm/llvm_program.cpp +++ b/taichi/runtime/program_impls/llvm/llvm_program.cpp @@ -142,7 +142,8 @@ void LlvmProgramImpl::dump_cache_data_to_disk() { if (config->offline_cache && !cache_data_.kernels.empty()) { LlvmOfflineCacheFileWriter writer{}; writer.set_data(std::move(cache_data_)); - writer.dump(config->offline_cache_file_path); + // Note: For offline-cache, new-metadata should be merged with old-metadata + writer.dump(config->offline_cache_file_path, LlvmOfflineCache::LL, true); } }