[bug] Fix crashing when loading old offline cache files (#6089)

Issue: #4401, fixes #6081

In future, if necessary, we should maintain a version number for offline
cache instead of using the Taichi version directly.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PGZXB

cmake/TaichiTests.cmake

@@ -20,7 +20,8 @@ file(GLOB_RECURSE TAICHI_TESTS_SOURCE
         "tests/cpp/llvm/*.cpp"
         "tests/cpp/program/*.cpp"
         "tests/cpp/struct/*.cpp"
-        "tests/cpp/transforms/*.cpp")
+        "tests/cpp/transforms/*.cpp"
+        "tests/cpp/offline_cache/*.cpp")
 
 if (TI_WITH_OPENGL OR TI_WITH_VULKAN)
     file(GLOB TAICHI_TESTS_GFX_UTILS_SOURCE

taichi/cache/gfx/cache_manager.cpp

@@ -47,13 +47,6 @@ struct CacheCleanerUtils<gfx::CacheManager::Metadata> {
   using MetadataType = gfx::CacheManager::Metadata;
   using KernelMetaData = MetadataType::KernelMetadata;
 
-  // To load metadata from file
-  static bool load_metadata(const CacheCleanerConfig &config,
-                            MetadataType &result) {
-    return read_from_binary_file(
-        result, taichi::join_path(config.path, config.metadata_filename));
-  }
-
   // To save metadata as file
   static bool save_metadata(const CacheCleanerConfig &config,
                             const MetadataType &data) {
@@ -81,13 +74,6 @@ struct CacheCleanerUtils<gfx::CacheManager::Metadata> {
     return true;
   }
 
-  // To check version
-  static bool check_version(const CacheCleanerConfig &config,
-                            const Version &version) {
-    return version[0] == TI_VERSION_MAJOR && version[1] == TI_VERSION_MINOR &&
-           version[2] == TI_VERSION_PATCH;
-  }
-
   // To get cache files name
   static std::vector<std::string> get_cache_files(
       const CacheCleanerConfig &config,
@@ -106,6 +92,12 @@ struct CacheCleanerUtils<gfx::CacheManager::Metadata> {
         taichi::join_path(config.path, kDebuggingAotMetadataFilename));
     taichi::remove(taichi::join_path(config.path, kGraphMetadataFilename));
   }
+
+  // To check if a file is cache file
+  static bool is_valid_cache_file(const CacheCleanerConfig &config,
+                                  const std::string &name) {
+    return filename_extension(name) == "spv";
+  }
 };
 
 }  // namespace offline_cache
@@ -184,10 +176,12 @@ void CacheManager::dump_with_merging() const {
       cache_builder->dump(path_, "");
 
       // Update offline_cache_metadata.tcb
+      using offline_cache::load_metadata_with_checking;
+      using Error = offline_cache::LoadMetadataError;
       Metadata old_data;
       const auto filename =
           taichi::join_path(path_, kOfflineCacheMetadataFilename);
-      if (read_from_binary_file(old_data, filename)) {
+      if (load_metadata_with_checking(old_data, filename) == Error::kNoError) {
         for (auto &[k, v] : offline_cache_metadata_.kernels) {
           auto iter = old_data.kernels.find(k);
           if (iter != old_data.kernels.end()) {  // Update

taichi/common/serialization.h

@@ -345,15 +345,20 @@ class BinarySerializer : public Serializer {
     preserved = 0;
   }
 
+  template <bool writing_ = writing>
+  typename std::enable_if<!writing_, std::size_t>::type retrieve_length() {
+    return *reinterpret_cast<std::size_t *>(c_data);
+  }
+
   void finalize() {
-    if (writing) {
+    if constexpr (writing) {
       if (c_data) {
         *reinterpret_cast<std::size_t *>(&c_data[0]) = head;
       } else {
         *reinterpret_cast<std::size_t *>(&data[0]) = head;
       }
     } else {
-      assert(head == *reinterpret_cast<std::size_t *>(c_data));
+      assert(head == retrieve_length());
     }
   }
 
@@ -880,6 +885,21 @@ operator<<(std::ostream &os, const T &t) {
 }
 
 // Returns true if deserialization succeeded.
+template <typename T>
+bool read_from_binary(T &t,
+                      const void *bin,
+                      std::size_t len,
+                      bool match_all = true) {
+  BinaryInputSerializer reader;
+  reader.initialize(const_cast<void *>(bin));
+  if (len != reader.retrieve_length()) {
+    return false;
+  }
+  reader(t);
+  auto head = reader.head;
+  return match_all ? head == len : head <= len;
+}
+
 template <typename T>
 bool read_from_binary_file(T &t, const std::string &file_name) {
   BinaryInputSerializer reader;

taichi/runtime/llvm/llvm_offline_cache.cpp

@@ -27,13 +27,6 @@ using Format = LlvmOfflineCache::Format;
 constexpr char kMetadataFilename[] = "metadata";
 constexpr char kMetadataFileLockName[] = "metadata.lock";
 
-static bool is_current_llvm_cache_version(
-    const LlvmOfflineCache::Version &ver) {
-  // TODO(PGZXB): Do more detailed checking
-  return ver[0] == TI_VERSION_MAJOR && ver[1] == TI_VERSION_MINOR &&
-         ver[2] == TI_VERSION_PATCH;
-}
-
 static std::string get_llvm_cache_metadata_file_path(const std::string &dir) {
   return taichi::join_path(dir, std::string(kMetadataFilename) + ".tcb");
 }
@@ -60,13 +53,6 @@ struct CacheCleanerUtils<LlvmOfflineCache> {
   using MetadataType = LlvmOfflineCache;
   using KernelMetaData = typename MetadataType::KernelMetadata;
 
-  // To load metadata from file
-  static bool load_metadata(const CacheCleanerConfig &config,
-                            MetadataType &result) {
-    return read_from_binary_file(
-        result, taichi::join_path(config.path, config.metadata_filename));
-  }
-
   // To save metadata as file
   static bool save_metadata(const CacheCleanerConfig &config,
                             const MetadataType &data) {
@@ -84,12 +70,6 @@ struct CacheCleanerUtils<LlvmOfflineCache> {
     return true;
   }
 
-  // To check version
-  static bool check_version(const CacheCleanerConfig &config,
-                            const Version &version) {
-    return is_current_llvm_cache_version(version);
-  }
-
   // To get cache files name
   static std::vector<std::string> get_cache_files(
       const CacheCleanerConfig &config,
@@ -106,6 +86,13 @@ struct CacheCleanerUtils<LlvmOfflineCache> {
   static void remove_other_files(const CacheCleanerConfig &config) {
     // Do nothing
   }
+
+  // To check if a file is cache file
+  static bool is_valid_cache_file(const CacheCleanerConfig &config,
+                                  const std::string &name) {
+    std::string ext = filename_extension(name);
+    return ext == "ll" || ext == "bc";
+  }
 };
 
 }  // namespace offline_cache
@@ -126,20 +113,12 @@ bool LlvmOfflineCacheFileReader::load_meta_data(
     LlvmOfflineCache &data,
     const std::string &cache_file_path,
     bool with_lock) {
+  using offline_cache::load_metadata_with_checking;
+  using Error = offline_cache::LoadMetadataError;
   const auto tcb_path = get_llvm_cache_metadata_file_path(cache_file_path);
-  {
-    // 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", cache_file_path);
-      return false;
-    }
-  }
 
   if (!with_lock) {
-    read_from_binary_file(data, tcb_path);
-    return true;
+    return Error::kNoError == load_metadata_with_checking(data, tcb_path);
   }
 
   std::string lock_path =
@@ -150,8 +129,7 @@ bool LlvmOfflineCacheFileReader::load_meta_data(
         TI_WARN("Unlock {} failed", lock_path);
       }
     });
-    read_from_binary_file(data, tcb_path);
-    return true;
+    return Error::kNoError == load_metadata_with_checking(data, tcb_path);
   }
   TI_WARN("Lock {} failed", lock_path);
   return false;
@@ -389,10 +367,6 @@ void LlvmOfflineCacheFileWriter::mangle_offloaded_task_name(
     for (auto &offload : compiled_data.tasks) {
       std::string mangled_name =
           offline_cache::mangle_name(offload.name, kernel_key);
-      TI_DEBUG(
-          "Mangle offloaded-task from internal name '{}' to offline cache "
-          "key '{}'",
-          offload.name, mangled_name);
       auto func = compiled_data.module->getFunction(offload.name);
       TI_ASSERT(func != nullptr);
       func->setName(mangled_name);

taichi/runtime/llvm/llvm_offline_cache.h

@@ -100,6 +100,7 @@ struct LlvmOfflineCache {
   std::unordered_map<std::string, KernelCacheData>
       kernels;  // key = kernel_name
 
+  // NOTE: The "version" must be the first field to be serialized
   TI_IO_DEF(version, size, fields, kernels);
 };
 

taichi/util/io.h

@@ -14,6 +14,11 @@
 
 #if defined(TI_PLATFORM_WINDOWS)
 #include <filesystem>
+#else  // POSIX
+#include <unistd.h>
+#include <dirent.h>
+#include <sys/types.h>
+#include <sys/stat.h>
 #endif
 
 TI_NAMESPACE_BEGIN
@@ -51,6 +56,47 @@ inline bool remove(const std::string &path) {
   return std::remove(path.c_str()) == 0;
 }
 
+template <typename Visitor>  // void(const std::string &name, bool is_dir)
+inline bool traverse_directory(const std::string &dir, Visitor v) {
+#if defined(TI_PLATFORM_WINDOWS)
+  namespace fs = std::filesystem;
+  std::error_code ec{};
+  auto iter = fs::directory_iterator(dir, ec);
+  if (ec) {
+    return false;
+  }
+  for (auto &f : iter) {
+    v(f.path().filename().string(), f.is_directory());
+  }
+  return true;
+#else  // POSIX
+  struct dirent *f = nullptr;
+  DIR *directory = ::opendir(dir.c_str());
+  if (!directory) {
+    return false;
+  }
+  while ((f = ::readdir(directory))) {
+    struct stat *stat_buf = nullptr;
+    auto fullpath = join_path(dir, f->d_name);
+    auto ret = ::stat(fullpath.c_str(), stat_buf);
+    TI_ASSERT(ret == 0);
+    v(f->d_name, S_ISDIR(stat_buf->st_mode));
+  }
+  auto ret = ::closedir(directory);
+  TI_ASSERT(ret == 0);
+  return true;
+#endif
+}
+
+inline std::string filename_extension(const std::string &filename) {
+  std::string postfix;
+  auto pos = filename.find_last_of('.');
+  if (pos != std::string::npos) {
+    postfix = filename.substr(pos + 1);
+  }
+  return postfix;
+}
+
 template <typename T>
 void write_to_disk(const T &dat, std::string fn) {
   FILE *f = fopen(fn.c_str(), "wb");