Skip to content

Commit

Permalink
[llvm] Use LlvmProgramImpl::cache_data_ to store compiled kernel info (
Browse files Browse the repository at this point in the history
…#5290)

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

* Rearranged comments
  • Loading branch information
jim19930609 authored Jun 30, 2022
1 parent e9471b6 commit 831b5db
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 6 deletions.
13 changes: 8 additions & 5 deletions taichi/codegen/llvm/codegen_llvm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2328,28 +2328,31 @@ 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()) {
kernel->lower();
}
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);
Expand Down
8 changes: 8 additions & 0 deletions taichi/program/kernel.h
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
43 changes: 42 additions & 1 deletion taichi/runtime/program_impls/llvm/llvm_program.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit 831b5db

Please sign in to comment.