From 831b5db056f923cd8ef5305ced7843b776a90abe Mon Sep 17 00:00:00 2001
From: Zhanlue Yang <jim19930609@gmail.com>
Date: Thu, 30 Jun 2022 14:49:02 +0800
Subject: [PATCH] [llvm] Use LlvmProgramImpl::cache_data_ to store compiled
 kernel info (#5290)

* [llvm] Use LlvmProgramImpl::cache_data_ to store compiled kernel info

* Rearranged comments
---
 taichi/codegen/llvm/codegen_llvm.cpp          | 13 +++---
 taichi/program/kernel.h                       |  8 ++++
 .../runtime/program_impls/llvm/llvm_program.h | 43 ++++++++++++++++++-
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/taichi/codegen/llvm/codegen_llvm.cpp b/taichi/codegen/llvm/codegen_llvm.cpp
index 8fc36feaec3b6..e9907ac2b90eb 100644
--- a/taichi/codegen/llvm/codegen_llvm.cpp
+++ b/taichi/codegen/llvm/codegen_llvm.cpp
@@ -2328,18 +2328,17 @@ void CodeGenLLVM::emit_to_module() {
 }
 
 CodeGenLLVM::CompiledData CodeGenLLVM::run_compilation() {
-  bool needs_cache = false;
   const auto &config = prog->config;
-  std::string kernel_key;
+  std::string kernel_key =
+      get_hashed_offline_cache_key(&kernel->program->config, kernel);
+  kernel->set_kernel_key_for_cache(kernel_key);
   if (config.offline_cache && !config.async_mode &&
       this->supports_offline_cache() && !kernel->is_evaluator) {
-    kernel_key = get_hashed_offline_cache_key(&kernel->program->config, kernel);
     CompiledData res;
     const bool ok = maybe_read_compilation_from_cache(kernel_key, &res);
     if (ok) {
       return res;
     }
-    needs_cache = true;
   }
 
   if (!kernel->lowered()) {
@@ -2347,9 +2346,13 @@ CodeGenLLVM::CompiledData CodeGenLLVM::run_compilation() {
   }
   emit_to_module();
   eliminate_unused_functions();
-  if (needs_cache) {
+
+  // Updates LlvmProgramImpl->cache_data_ to save the compiled kernel
+  // information for successive uses in AOT or CGraph.
+  if (!kernel->is_evaluator) {
     cache_module(kernel_key);
   }
+
   CompiledData res;
   res.offloaded_tasks = std::move(this->offloaded_tasks);
   res.llvm_module = std::move(this->module);
diff --git a/taichi/program/kernel.h b/taichi/program/kernel.h
index d5a0ecfc920b9..7d36a8f621d2e 100644
--- a/taichi/program/kernel.h
+++ b/taichi/program/kernel.h
@@ -140,6 +140,14 @@ class TI_DLL_EXPORT Kernel : public Callable {
    */
   static bool supports_lowering(Arch arch);
 
+  void set_kernel_key_for_cache(const std::string &kernel_key) {
+    kernel_key_ = kernel_key;
+  }
+
+  const std::string &get_cached_kernel_key() {
+    return kernel_key_;
+  }
+
  private:
   void init(Program &program,
             const std::function<void()> &func,
diff --git a/taichi/runtime/program_impls/llvm/llvm_program.h b/taichi/runtime/program_impls/llvm/llvm_program.h
index dc75efcccfa36..cc7b3c380287c 100644
--- a/taichi/runtime/program_impls/llvm/llvm_program.h
+++ b/taichi/runtime/program_impls/llvm/llvm_program.h
@@ -223,10 +223,51 @@ class LlvmProgramImpl : public ProgramImpl {
     return runtime_exec_->llvm_device();
   }
 
+  // TODO(zhanlue): Rearrange llvm::Context's ownership
+  //
+  // In LLVM backend, most of the compiled information are stored in
+  // llvm::Module:
+  // 1. Runtime functions are compiled into runtime_module,
+  // 2. Fields are compiled into struct_module,
+  // 3. Each kernel is compiled into individual kernel_module
+  //
+  // However, all the llvm::Modules are owned by llvm::Context, which belongs to
+  // TaichiLLVMContext. Upon destruction, there's an implicit requirement that
+  // TaichiLLVMContext has to stay alive until all the llvm::Modules are
+  // destructed, otherwise there will be risks of dangling references.
+  //
+  // To guarantee the life cycle of llvm::Module stay aligned with
+  // llvm::Context, we better make llvm::Context a more global-scoped variable,
+  // instead of owned by TaichiLLVMContext.
+  //
+  // Objects owning llvm::Module so far (from direct to indirect):
+  // 1. LlvmOfflineCache::CachedKernelData(direct owner)
+  // 2. LlvmOfflineCache
+  //   3.1 LlvmProgramImpl
+  //   3.2 LlvmAotModuleBuilder
+  //   3.3 llvm_aot::KernelImpl (for use in CGraph)
+  //
+  // Objects owning llvm::Context (from direct to indirect)
+  // 1. TaichiLLVMContext
+  // 2. LlvmProgramImpl
+  //
+  // Make sure the above mentioned objects are destructed in order.
+  ~LlvmProgramImpl() {
+    // Explicitly enforce "LlvmOfflineCache::CachedKernelData::owned_module"
+    // destructs before
+    // "LlvmRuntimeExecutor::TaichiLLVMContext::ThreadSafeContext"
+
+    // 1. Destructs cahce_data_
+    cache_data_ = LlvmOfflineCache();
+
+    // 2. Destructs runtime_exec_
+    runtime_exec_.reset();
+  }
+
  private:
   std::size_t num_snode_trees_processed_{0};
-  LlvmOfflineCache cache_data_;
   std::unique_ptr<LlvmRuntimeExecutor> runtime_exec_;
+  LlvmOfflineCache cache_data_;
 };
 
 LlvmProgramImpl *get_llvm_program(Program *prog);