[llvm] Add attributes to LLVMCompiledData (#5929)

* [llvm] Add attributes to LLVMCompiledData

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

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

taichi/codegen/llvm/codegen_llvm.cpp

@@ -2625,7 +2625,8 @@ LLVMCompiledData TaskCodeGenLLVM::run_compilation() {
   emit_to_module();
   eliminate_unused_functions();
 
-  return {std::move(this->offloaded_tasks), std::move(this->module)};
+  return {std::move(offloaded_tasks), std::move(module),
+          std::move(used_tree_ids), std::move(struct_for_tls_sizes)};
 }
 
 llvm::Value *TaskCodeGenLLVM::create_xlogue(std::unique_ptr<Block> &block) {
@@ -2701,7 +2702,8 @@ void TaskCodeGenLLVM::visit(FuncCallStmt *stmt) {
 }
 
 LLVMCompiledData LLVMCompiledData::clone() const {
-  return {tasks, llvm::CloneModule(*module)};
+  return {tasks, llvm::CloneModule(*module), used_tree_ids,
+          struct_for_tls_sizes};
 }
 
 TLANG_NAMESPACE_END

taichi/codegen/llvm/codegen_llvm.h

@@ -58,6 +58,8 @@ class TaskCodeGenLLVM : public IRVisitor, public LLVMModuleBuilder {
   llvm::BasicBlock *final_block;
   std::set<std::string> linked_modules;
   bool returned{false};
+  std::unordered_set<int> used_tree_ids;
+  std::unordered_set<int> struct_for_tls_sizes;
 
   std::unordered_map<const Stmt *, std::vector<llvm::Value *>> loop_vars_llvm;
 

taichi/codegen/llvm/llvm_compiled_data.h

@@ -23,11 +23,18 @@ class OffloadedTask {
 struct LLVMCompiledData {
   std::vector<OffloadedTask> tasks;
   std::unique_ptr<llvm::Module> module{nullptr};
+  std::unordered_set<int> used_tree_ids;
+  std::unordered_set<int> struct_for_tls_sizes;
   LLVMCompiledData() = default;
   LLVMCompiledData(LLVMCompiledData &&) = default;
   LLVMCompiledData(std::vector<OffloadedTask> tasks,
-                   std::unique_ptr<llvm::Module> module)
-      : tasks(std::move(tasks)), module(std::move(module)) {
+                   std::unique_ptr<llvm::Module> module,
+                   std::unordered_set<int> used_tree_ids,
+                   std::unordered_set<int> struct_for_tls_sizes)
+      : tasks(std::move(tasks)),
+        module(std::move(module)),
+        used_tree_ids(std::move(used_tree_ids)),
+        struct_for_tls_sizes(std::move(struct_for_tls_sizes)) {
   }
   LLVMCompiledData clone() const;
   TI_IO_DEF(tasks);

taichi/codegen/wasm/codegen_wasm.cpp

@@ -276,7 +276,7 @@ LLVMCompiledData KernelCodeGenWASM::compile_task(
 
   gen->tlctx->jit->global_optimize_module(gen->module.get());
 
-  return {name_list, std::move(gen->module)};
+  return {name_list, std::move(gen->module), {}, {}};
 }
 }  // namespace lang
 }  // namespace taichi

taichi/runtime/llvm/llvm_offline_cache.cpp

@@ -147,8 +147,9 @@ bool LlvmOfflineCacheFileReader::get_kernel_cache(
         return false;  // Must return
       }
     }
-    res.compiled_data_list.emplace_back(data.tasks,
-                                        llvm::CloneModule(*data.module));
+    res.compiled_data_list.emplace_back(
+        data.tasks, llvm::CloneModule(*data.module), data.used_tree_ids,
+        data.struct_for_tls_sizes);
   }
 
   kernel_data.last_used_at = std::time(nullptr);

taichi/runtime/program_impls/llvm/llvm_program.cpp

@@ -145,7 +145,8 @@ void LlvmProgramImpl::cache_kernel(
   kernel_cache.kernel_key = kernel_key;
   for (const auto &data : data_list) {
     kernel_cache.compiled_data_list.emplace_back(
-        data.tasks, llvm::CloneModule(*data.module));
+        data.tasks, llvm::CloneModule(*data.module), data.used_tree_ids,
+        data.struct_for_tls_sizes);
   }
   kernel_cache.args = std::move(args);
   kernel_cache.created_at = std::time(nullptr);

tests/cpp/llvm/llvm_offline_cache_test.cpp

@@ -105,7 +105,10 @@ TEST_P(LlvmOfflineCacheTest, ReadWrite) {
     task.block_dim = kBlockDim;
     task.grid_dim = kGridDim;
     tasks.push_back(task);
-    kcache.compiled_data_list.emplace_back(tasks, make_module(*llvm_ctx));
+    LLVMCompiledData data;
+    data.tasks = tasks;
+    data.module = make_module(*llvm_ctx);
+    kcache.compiled_data_list.push_back(std::move(data));
     kcache.args = arg_infos;
     writer.add_kernel_cache(kKernelName, std::move(kcache));
     writer.set_no_mangle();