[misc] Lazy load spirv code from disk during offline cache

taichi/cache/gfx/cache_manager.cpp

@@ -134,6 +134,7 @@ CacheManager::CacheManager(Params &&init_params)
       gfx::AotModuleParams params;
       params.module_path = path_;
       params.runtime = runtime_;
+      params.enable_lazy_loading = true;
       cached_module_ = gfx::make_aot_module(params, init_params.arch);
     }
   }

taichi/runtime/gfx/aot_module_loader_impl.cpp

@@ -24,28 +24,31 @@ class FieldImpl : public aot::Field {
 class AotModuleImpl : public aot::Module {
  public:
   explicit AotModuleImpl(const AotModuleParams &params, Arch device_api_backend)
-      : runtime_(params.runtime), device_api_backend_(device_api_backend) {
+      : module_path_(params.module_path),
+        runtime_(params.runtime),
+        device_api_backend_(device_api_backend) {
     const std::string bin_path =
         fmt::format("{}/metadata.tcb", params.module_path);
     if (!read_from_binary_file(ti_aot_data_, bin_path)) {
       mark_corrupted();
       return;
     }
 
-    for (int i = 0; i < ti_aot_data_.kernels.size(); ++i) {
-      auto k = ti_aot_data_.kernels[i];
-
-      std::vector<std::vector<uint32_t>> spirv_sources_codes;
-      for (int j = 0; j < k.tasks_attribs.size(); ++j) {
-        std::vector<uint32_t> res =
-            read_spv_file(params.module_path, k.tasks_attribs[j]);
-        if (res.size() == 0) {
-          mark_corrupted();
-          return;
+    if (!params.enable_lazy_loading) {
+      for (int i = 0; i < ti_aot_data_.kernels.size(); ++i) {
+        auto k = ti_aot_data_.kernels[i];
+        std::vector<std::vector<uint32_t>> spirv_sources_codes;
+        for (int j = 0; j < k.tasks_attribs.size(); ++j) {
+          std::vector<uint32_t> res =
+              read_spv_file(params.module_path, k.tasks_attribs[j]);
+          if (res.size() == 0) {
+            mark_corrupted();
+            return;
+          }
+          spirv_sources_codes.push_back(res);
         }
-        spirv_sources_codes.push_back(res);
+        ti_aot_data_.spirv_codes.push_back(spirv_sources_codes);
       }
-      ti_aot_data_.spirv_codes.push_back(spirv_sources_codes);
     }
 
     const std::string graph_path =
@@ -104,6 +107,9 @@ class AotModuleImpl : public aot::Module {
       // AOT, only use the name of the function which should be the first part
       // of the struct
       if (ti_aot_data_.kernels[i].name.rfind(name, 0) == 0) {
+        if (!try_load_spv_kernel(i)) {
+          return false;
+        }
         kernel.kernel_attribs = ti_aot_data_.kernels[i];
         kernel.task_spirv_source_codes = ti_aot_data_.spirv_codes[i];
         // We don't have to store the number of SNodeTree in |ti_aot_data_| yet,
@@ -141,19 +147,40 @@ class AotModuleImpl : public aot::Module {
     return std::make_unique<FieldImpl>(runtime_, field);
   }
 
-  std::vector<uint32_t> read_spv_file(const std::string &output_dir,
-                                      const TaskAttributes &k) {
+  bool try_load_spv_kernel(std::size_t index) {
+    if (index >= ti_aot_data_.spirv_codes.size() ||
+        ti_aot_data_.spirv_codes[index].empty()) {
+      ti_aot_data_.spirv_codes.resize(index + 1);
+      auto &codes = ti_aot_data_.spirv_codes[index];
+      const auto &k = ti_aot_data_.kernels[index];
+      for (const auto &t : k.tasks_attribs) {
+        auto spv = read_spv_file(module_path_, t);
+        if (spv.empty()) {
+          mark_corrupted();
+          return false;
+        }
+        codes.push_back(spv);
+      }
+    }
+    return true;
+  }
+
+  static std::vector<uint32_t> read_spv_file(const std::string &output_dir,
+                                             const TaskAttributes &k) {
     const std::string spv_path = fmt::format("{}/{}.spv", output_dir, k.name);
     std::vector<uint32_t> source_code;
     std::ifstream fs(spv_path, std::ios_base::binary | std::ios::ate);
-    size_t size = fs.tellg();
-    fs.seekg(0, std::ios::beg);
-    source_code.resize(size / sizeof(uint32_t));
-    fs.read((char *)source_code.data(), size);
-    fs.close();
+    if (fs.is_open()) {
+      size_t size = fs.tellg();
+      fs.seekg(0, std::ios::beg);
+      source_code.resize(size / sizeof(uint32_t));
+      fs.read((char *)source_code.data(), size);
+      fs.close();
+    }
     return source_code;
   }
 
+  std::string module_path_;
   TaichiAotData ti_aot_data_;
   GfxRuntime *runtime_{nullptr};
   Arch device_api_backend_;

taichi/runtime/gfx/aot_module_loader_impl.h

@@ -19,6 +19,13 @@ namespace gfx {
 struct TI_DLL_EXPORT AotModuleParams {
   std::string module_path;
   GfxRuntime *runtime{nullptr};
+  bool enable_lazy_loading{false};
+
+  AotModuleParams() = default;
+
+  AotModuleParams(const std::string &path, GfxRuntime *rt)
+      : module_path(path), runtime(rt) {
+  }
 };
 
 TI_DLL_EXPORT std::unique_ptr<aot::Module> make_aot_module(