[llvm] Avoid creating new LLVM contexts when updating struct module

taichi/codegen/llvm/codegen_llvm.cpp

@@ -925,8 +925,9 @@ void CodeGenLLVM::emit_gc(OffloadedStmt *stmt) {
 }
 
 llvm::Value *CodeGenLLVM::create_call(llvm::Value *func,
-                                      llvm::ArrayRef<llvm::Value *> args) {
-  check_func_call_signature(func, args);
+                                      llvm::ArrayRef<llvm::Value *> args_arr) {
+  std::vector<llvm::Value *> args = args_arr;
+  check_func_call_signature(func, args, builder.get());
 #ifdef TI_LLVM_15
   llvm::FunctionType *func_ty = nullptr;
   if (auto *fn = llvm::dyn_cast<llvm::Function>(func)) {

taichi/codegen/llvm/llvm_codegen_utils.cpp

@@ -6,12 +6,59 @@ namespace lang {
 std::string type_name(llvm::Type *type) {
   std::string type_name_str;
   llvm::raw_string_ostream rso(type_name_str);
-  type->print(rso);
+  type->print(rso, /*IsForDebug=*/false, /*NoDetails=*/true);
   return type_name_str;
 }
 
+/*
+ * Determine whether two types are the same
+ * (a type is a renamed version of the other one) based on the
+ * type name. Check recursively if the types are function types.
+ *
+ * The name of a type imported multiple times is added a suffix starting with a
+ * "." following by a number. For example, "RuntimeContext" may be renamed to
+ * names like "RuntimeContext.0" and "RuntimeContext.8".
+ */
+bool is_same_type(llvm::Type *required, llvm::Type *provided) {
+  if (required == provided) {
+    return true;
+  }
+  if (required->isPointerTy() != provided->isPointerTy()) {
+    return false;
+  }
+  if (required->isPointerTy()) {
+    required = required->getPointerElementType();
+    provided = provided->getPointerElementType();
+  }
+  if (required->isFunctionTy() != provided->isFunctionTy()) {
+    return false;
+  }
+  if (required->isFunctionTy()) {
+    auto req_func = llvm::dyn_cast<llvm::FunctionType>(required);
+    auto prov_func = llvm::dyn_cast<llvm::FunctionType>(provided);
+    if (!is_same_type(req_func->getReturnType(), prov_func->getReturnType())) {
+      return false;
+    }
+    if (req_func->getNumParams() != prov_func->getNumParams()) {
+      return false;
+    }
+    for (int j = 0; j < req_func->getNumParams(); j++) {
+      if (!is_same_type(req_func->getParamType(j),
+                        prov_func->getParamType(j))) {
+        return false;
+      }
+    }
+    return true;
+  }
+  auto req_name = type_name(required);
+  auto prov_name = type_name(provided);
+  int min_len = std::min(req_name.size(), prov_name.size());
+  return req_name.substr(0, min_len) == prov_name.substr(0, min_len);
+}
+
 void check_func_call_signature(llvm::Value *func,
-                               std::vector<llvm::Value *> arglist) {
+                               std::vector<llvm::Value *> &arglist,
+                               llvm::IRBuilder<> *builder) {
   llvm::FunctionType *func_type = nullptr;
   if (llvm::Function *fn = llvm::dyn_cast<llvm::Function>(func)) {
     func_type = fn->getFunctionType();
@@ -31,7 +78,19 @@ void check_func_call_signature(llvm::Value *func,
   for (int i = 0; i < num_params; i++) {
     auto required = func_type->getFunctionParamType(i);
     auto provided = arglist[i]->getType();
+    /*
+     * When importing a module from file, the imported `llvm::Type`s can get
+     * conflict with the same type in the llvm::Context. In such scenario,
+     * the imported types will be renamed from "original_type" to
+     * "original_type.xxx", making them separate types in essence.
+     * To make the types of the argument and parameter the same,
+     * a pointer cast must be performed.
+     */
     if (required != provided) {
+      if (is_same_type(required, provided)) {
+        arglist[i] = builder->CreatePointerCast(arglist[i], required);
+        continue;
+      }
       TI_INFO("Function : {}", std::string(func->getName()));
       TI_INFO("    Type : {}", type_name(func->getType()));
       if (&required->getContext() != &provided->getContext()) {

taichi/codegen/llvm/llvm_codegen_utils.h

@@ -44,7 +44,8 @@ inline constexpr char kLLVMPhysicalCoordinatesName[] = "PhysicalCoordinates";
 std::string type_name(llvm::Type *type);
 
 void check_func_call_signature(llvm::Value *func,
-                               std::vector<llvm::Value *> arglist);
+                               std::vector<llvm::Value *> &arglist,
+                               llvm::IRBuilder<> *builder);
 
 class LLVMModuleBuilder {
  public:
@@ -121,8 +122,9 @@ class LLVMModuleBuilder {
                     const std::string &func_name,
                     const std::vector<llvm::Value *> &arglist) {
     auto func = get_runtime_function(func_name);
-    check_func_call_signature(func, arglist);
-    return builder->CreateCall(func, arglist);
+    std::vector<llvm::Value *> args = arglist;
+    check_func_call_signature(func, args, builder);
+    return builder->CreateCall(func, args);
   }
 
   template <typename... Args>
@@ -131,7 +133,7 @@ class LLVMModuleBuilder {
                     Args &&...args) {
     auto func = get_runtime_function(func_name);
     auto arglist = std::vector<llvm::Value *>({args...});
-    check_func_call_signature(func, arglist);
+    check_func_call_signature(func, arglist, builder);
     return builder->CreateCall(func, arglist);
   }
 
@@ -186,7 +188,7 @@ class RuntimeObject {
   llvm::Value *call(const std::string &func_name, Args &&...args) {
     auto func = get_func(func_name);
     auto arglist = std::vector<llvm::Value *>({ptr, args...});
-    check_func_call_signature(func, arglist);
+    check_func_call_signature(func, arglist, builder);
     return builder->CreateCall(func, arglist);
   }
 

taichi/runtime/llvm/llvm_context.cpp

@@ -529,12 +529,6 @@ void TaichiLLVMContext::set_struct_module(
       continue;
     }
     TI_ASSERT(!data->runtime_module);
-    data->struct_module.reset();
-    old_contexts_.push_back(std::move(data->thread_safe_llvm_context));
-    data->thread_safe_llvm_context =
-        std::make_unique<llvm::orc::ThreadSafeContext>(
-            std::make_unique<llvm::LLVMContext>());
-    data->llvm_context = data->thread_safe_llvm_context->getContext();
     data->struct_module = clone_module_to_context(
         this_thread_data->struct_module.get(), data->llvm_context);
   }

taichi/runtime/llvm/llvm_context.h

@@ -163,9 +163,6 @@ class TaichiLLVMContext {
   std::mutex thread_map_mut_;
 
   std::unordered_map<int, std::vector<std::string>> snode_tree_funcs_;
-  std::vector<std::unique_ptr<llvm::orc::ThreadSafeContext>>
-      old_contexts_;  // Contains old contexts that have modules stored inside
-                      // the offline cache.
 };
 
 class LlvmModuleBitcodeLoader {