[bug] Fix missing old but useful metadata in offline cache

taichi/runtime/llvm/llvm_offline_cache.cpp

@@ -26,22 +26,31 @@ constexpr char kMetadataFilename[] = "metadata";
 std::unique_ptr<LlvmOfflineCacheFileReader> LlvmOfflineCacheFileReader::make(
     const std::string &path,
     LlvmOfflineCache::Format format) {
+  LlvmOfflineCache data;
+  if (!load_meta_data(data, path)) {
+    return nullptr;
+  }
+  return std::unique_ptr<LlvmOfflineCacheFileReader>(
+      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<LlvmOfflineCacheFileReader>(
-      new LlvmOfflineCacheFileReader(path, std::move(data), format));
+  return true;
 }
 
 LlvmOfflineCacheFileReader::LlvmOfflineCacheFileReader(
@@ -111,7 +120,8 @@ std::unique_ptr<llvm::Module> 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,

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,

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);
   }
 }