This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c4ab46c7a6 [FFI] Use thread-safe packed function initialization
(#19939)
c4ab46c7a6 is described below
commit c4ab46c7a6f18c34503b86d5f19bd7ce4421ecef
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Jul 5 04:33:31 2026 +0800
[FFI] Use thread-safe packed function initialization (#19939)
Generated LLVM modules currently cache imported packed functions through
an unsynchronized null check and plain store. Concurrent first calls can
therefore race the lookup and cache update.
This change:
- routes generated C and LLVM lookups directly through
TVMFFIEnvModLookupFromImports and removes the superseded
TVMBackendGetFuncFromEnv wrapper;
- gives each LLVM cached function a readable internal hidden initializer
and publishes its module-owned result through TVMFFIHandleInitOnce;
- removes the unused TVMBackendRunOnce export and dead LLVM static-init
producer while preserving the live static-handle path.
Validated with the existing LLVM, C-host, common codegen, static-init,
and C++ test suites.
---
include/tvm/runtime/c_backend_api.h | 26 ------
src/runtime/device_api.cc | 12 ---
src/runtime/module.cc | 3 +-
src/target/llvm/codegen_cpu.cc | 165 ++++++++++++++++--------------------
src/target/llvm/codegen_cpu.h | 20 +++--
src/target/source/codegen_c_host.cc | 5 +-
6 files changed, 91 insertions(+), 140 deletions(-)
diff --git a/include/tvm/runtime/c_backend_api.h
b/include/tvm/runtime/c_backend_api.h
index 5a3bd4c216..07d1d305de 100644
--- a/include/tvm/runtime/c_backend_api.h
+++ b/include/tvm/runtime/c_backend_api.h
@@ -34,19 +34,6 @@
extern "C" {
#endif
-/*!
- * \brief Backend function for modules to get function
- * from its environment mod_node (its imports and global function).
- * The user do should not call TVMFuncFree on func.
- *
- * \param mod_node The module handle.
- * \param func_name The name of the function.
- * \param out The result function.
- * \return 0 when no error is thrown, -1 when failure happens
- */
-TVM_RUNTIME_DLL int TVMBackendGetFuncFromEnv(void* mod_node, const char*
func_name,
- TVMFFIObjectHandle* out);
-
/*!
* \brief Backend function to allocate temporal workspace.
*
@@ -116,19 +103,6 @@ TVM_RUNTIME_DLL int
TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* c
*/
TVM_RUNTIME_DLL int TVMBackendParallelBarrier(int task_id,
TVMParallelGroupEnv* penv);
-/*!
- * \brief Simple static initialization function.
- * Run f once and set handle to be not null.
- * This function is mainly used for test purpose.
- *
- * \param handle A global address to indicate f
- * \param f The function to be run
- * \param cdata The closure data to pass to the function.
- * \param nbytes Number of bytes in the closure data.
- * \return 0 when no error is thrown, -1 when failure happens
- */
-TVM_RUNTIME_DLL int TVMBackendRunOnce(void** handle, int (*f)(void*), void*
cdata, int nbytes);
-
#ifdef __cplusplus
} // TVM_EXTERN_C
#endif
diff --git a/src/runtime/device_api.cc b/src/runtime/device_api.cc
index 959cd619ab..e573b8fb4e 100644
--- a/src/runtime/device_api.cc
+++ b/src/runtime/device_api.cc
@@ -241,10 +241,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
using namespace tvm::runtime;
-int TVMBackendGetFuncFromEnv(void* mod_node, const char* func_name,
TVMFFIObjectHandle* func) {
- return TVMFFIEnvModLookupFromImports(mod_node, func_name, func);
-}
-
void* TVMBackendAllocWorkspace(int device_type, int device_id, uint64_t size,
int dtype_code_hint,
int dtype_bits_hint) {
DLDevice dev;
@@ -266,11 +262,3 @@ int TVMBackendFreeWorkspace(int device_type, int
device_id, void* ptr) {
DeviceAPIManager::Get(dev)->FreeWorkspace(dev, ptr);
return 0;
}
-
-int TVMBackendRunOnce(void** handle, int (*f)(void*), void* cdata, int nbytes)
{
- if (*handle == nullptr) {
- *handle = reinterpret_cast<void*>(1);
- return (*f)(cdata);
- }
- return 0;
-}
diff --git a/src/runtime/module.cc b/src/runtime/module.cc
index b72607134d..7892f3eb4e 100644
--- a/src/runtime/module.cc
+++ b/src/runtime/module.cc
@@ -78,7 +78,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// Initialize the functions
TVM_INIT_CONTEXT_FUNC(TVMFFIFunctionCall);
TVM_INIT_CONTEXT_FUNC(TVMFFIErrorSetRaisedFromCStr);
- TVM_INIT_CONTEXT_FUNC(TVMBackendGetFuncFromEnv);
+ TVM_INIT_CONTEXT_FUNC(TVMFFIEnvModLookupFromImports);
+ TVM_INIT_CONTEXT_FUNC(TVMFFIHandleInitOnce);
TVM_INIT_CONTEXT_FUNC(TVMBackendAllocWorkspace);
TVM_INIT_CONTEXT_FUNC(TVMBackendFreeWorkspace);
TVM_INIT_CONTEXT_FUNC(TVMBackendParallelLaunch);
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index a3e0f9030b..35689e4f9a 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -76,6 +76,7 @@ void CodeGenCPU::Init(const std::string& module_name,
LLVMTarget* llvm_target,
system_lib_prefix_ = system_lib_prefix;
dbg_info_ = CreateDebugInfo(module_.get());
func_handle_map_.clear();
+ func_handle_init_map_.clear();
export_system_symbols_.clear();
// Runtime types.
@@ -128,11 +129,20 @@ void CodeGenCPU::Init(const std::string& module_name,
LLVMTarget* llvm_target,
t_void_,
{llvmGetPointerTo(t_char_, 0),
llvmGetPointerTo(llvmGetPointerTo(t_char_, 0), 0), t_int_},
false);
- // Defined in include/tvm/runtime/c_backend_api.h:
- // int TVMBackendGetFuncFromEnv(void* mod_node, const char* func_name,
TVMFunctionHandle* out);
- ftype_tvm_get_func_from_env_ = llvm::FunctionType::get(
+ // Defined in include/tvm/ffi/extra/c_env_api.h:
+ // int TVMFFIEnvModLookupFromImports(TVMFFIObjectHandle library_ctx, const
char* func_name,
+ // TVMFFIObjectHandle* out);
+ ftype_tvm_ffi_env_mod_lookup_from_imports_ = llvm::FunctionType::get(
t_int_, {t_void_p_, llvmGetPointerTo(t_char_, 0),
llvmGetPointerTo(t_tvm_func_handle_, 0)},
false);
+ // Defined in include/tvm/ffi/c_api.h:
+ // int TVMFFIHandleInitOnce(void** handle_addr, int (*init_func)(void**
result));
+ ftype_tvm_ffi_handle_init_callback_ =
+ llvm::FunctionType::get(t_int_, {llvmGetPointerTo(t_void_p_, 0)}, false);
+ ftype_tvm_ffi_handle_init_once_ = llvm::FunctionType::get(
+ t_int_,
+ {llvmGetPointerTo(t_void_p_, 0),
llvmGetPointerTo(ftype_tvm_ffi_handle_init_callback_, 0)},
+ false);
// Defined in include/tvm/runtime/c_backend_api.h:
// int TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* cdata, int
num_task);
ftype_tvm_parallel_launch_ = llvm::FunctionType::get(
@@ -141,12 +151,6 @@ void CodeGenCPU::Init(const std::string& module_name,
LLVMTarget* llvm_target,
// int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv);
ftype_tvm_parallel_barrier_ = llvm::FunctionType::get(
t_int_, {t_int_, llvmGetPointerTo(t_tvm_parallel_group_env_, 0)}, false);
- ftype_tvm_static_init_callback_ = llvm::FunctionType::get(t_int_,
{t_void_p_}, false);
- ftype_tvm_static_init_ = llvm::FunctionType::get(
- t_int_,
- {llvmGetPointerTo(t_void_p_, 0),
llvmGetPointerTo(ftype_tvm_static_init_callback_, 0),
- t_void_p_, t_int_},
- false);
// initialize TVM runtime API
if (system_lib_prefix_.has_value() && !target_c_runtime) {
// We will need this in environment for backward registration.
@@ -168,9 +172,12 @@ void CodeGenCPU::Init(const std::string& module_name,
LLVMTarget* llvm_target,
f_tvm_ffi_set_raised_from_c_str_parts_ = llvm::Function::Create(
ftype_tvm_ffi_error_set_raised_from_c_str_parts_,
llvm::Function::ExternalLinkage,
"TVMFFIErrorSetRaisedFromCStrParts", module_.get());
- f_tvm_get_func_from_env_ =
- llvm::Function::Create(ftype_tvm_get_func_from_env_,
llvm::Function::ExternalLinkage,
- "TVMBackendGetFuncFromEnv", module_.get());
+ f_tvm_ffi_env_mod_lookup_from_imports_ = llvm::Function::Create(
+ ftype_tvm_ffi_env_mod_lookup_from_imports_,
llvm::Function::ExternalLinkage,
+ "TVMFFIEnvModLookupFromImports", module_.get());
+ f_tvm_ffi_handle_init_once_ =
+ llvm::Function::Create(ftype_tvm_ffi_handle_init_once_,
llvm::Function::ExternalLinkage,
+ "TVMFFIHandleInitOnce", module_.get());
f_tvm_parallel_launch_ =
llvm::Function::Create(ftype_tvm_parallel_launch_,
llvm::Function::ExternalLinkage,
"TVMBackendParallelLaunch", module_.get());
@@ -462,8 +469,11 @@ void CodeGenCPU::InitGlobalContext(bool dynamic_lookup) {
if (!dynamic_lookup) {
gv_tvm_ffi_func_call_ =
InitContextPtr(llvmGetPointerTo(ftype_tvm_ffi_func_call_, 0),
"__TVMFFIFunctionCall");
- gv_tvm_get_func_from_env_ =
InitContextPtr(llvmGetPointerTo(ftype_tvm_get_func_from_env_, 0),
- "__TVMBackendGetFuncFromEnv");
+ gv_tvm_ffi_env_mod_lookup_from_imports_ =
+
InitContextPtr(llvmGetPointerTo(ftype_tvm_ffi_env_mod_lookup_from_imports_, 0),
+ "__TVMFFIEnvModLookupFromImports");
+ gv_tvm_ffi_handle_init_once_ = InitContextPtr(
+ llvmGetPointerTo(ftype_tvm_ffi_handle_init_once_, 0),
"__TVMFFIHandleInitOnce");
gv_tvm_ffi_set_last_error_c_str_ =
InitContextPtr(llvmGetPointerTo(ftype_tvm_ffi_error_set_raised_by_c_str_, 0),
"__TVMFFIErrorSetRaisedFromCStr");
@@ -689,49 +699,40 @@ llvm::Value* CodeGenCPU::CreateStaticHandle() {
return gv;
}
-void CodeGenCPU::CreateStaticInit(const std::string& init_fname, const Stmt&
body) {
- // closure data
- llvm::Function* f =
- llvm::Function::Create(ftype_tvm_static_init_callback_,
llvm::Function::PrivateLinkage,
- "__tvm_static_init_lambda", module_.get());
- SetTargetAttributes(f);
- llvm::Value* gv = CreateStaticHandle();
- llvm::Function* finit = module_->getFunction(init_fname);
- if (finit == nullptr) {
- finit = llvm::Function::Create(ftype_tvm_static_init_,
llvm::Function::ExternalLinkage,
- init_fname, module_.get());
- }
- // allocate and setup the closure, call the closure.
- uint64_t nbytes;
- ffi::Array<Var> vfields = tirx::UndefinedVars(body, {});
- TypedPointer cdata = PackClosureData(vfields, &nbytes);
- llvm::BasicBlock* init_end = CheckCallSuccess(builder_->CreateCall(
- finit, {gv, f, builder_->CreatePointerCast(cdata.addr, t_void_p_),
ConstInt32(nbytes)}));
- // Setup the closure function.
- auto* lambda_entry = llvm::BasicBlock::Create(*llvm_target_->GetContext(),
"entry", f);
- builder_->SetInsertPoint(lambda_entry);
- auto it = f->arg_begin();
- cdata.addr = builder_->CreatePointerCast(&(*it++), cdata.addr->getType());
- // setup new variable map, swap it with current var context.
- std::unordered_map<const VarNode*, llvm::Value*> new_vmap;
- UnpackClosureData(cdata, vfields, &new_vmap);
- TVM_FFI_ICHECK(parallel_env_.penv == nullptr);
- auto new_analyzer = arith::Analyzer();
- std::swap(function_, f);
- std::swap(analyzer_, new_analyzer);
- std::swap(var_map_, new_vmap);
- this->VisitStmt(body);
- builder_->CreateRet(ConstInt32(0));
- // swap the var map back, now we are back on track.
- std::swap(var_map_, new_vmap);
- std::swap(analyzer_, new_analyzer);
- std::swap(function_, f);
- builder_->SetInsertPoint(init_end);
+llvm::Function* CodeGenCPU::CreatePackedFuncInit(const std::string& fname) {
+ // Generate the equivalent of:
+ //
+ // static int __tvm_func_handle_init.<fname>(void** result) {
+ // return TVMFFIEnvModLookupFromImports(__tvm_ffi__library_ctx,
"<fname>", result);
+ // }
+ //
+ // The lookup result is owned by the module and remains valid for the
generated module's lifetime.
+ llvm::IRBuilderBase::InsertPoint saved_ip = builder_->saveIP();
+ llvm::Function* saved_function = function_;
+ llvm::LLVMContext* ctx = llvm_target_->GetContext();
+ llvm::Function* init_func =
+ llvm::Function::Create(ftype_tvm_ffi_handle_init_callback_,
llvm::Function::InternalLinkage,
+ "__tvm_func_handle_init." + fname, module_.get());
+ init_func->setVisibility(llvm::GlobalValue::HiddenVisibility);
+ SetTargetAttributes(init_func);
+ function_ = init_func;
+ builder_->SetInsertPoint(llvm::BasicBlock::Create(*ctx, "entry", init_func));
+
+ llvm::Value* result = &*init_func->arg_begin();
+ llvm::Value* module_ctx =
builder_->CreateAlignedLoad(gv_mod_ctx_->getValueType(), gv_mod_ctx_,
+
llvm::Align(gv_mod_ctx_->getAlignment()));
+ auto lookup_callee =
llvm::FunctionCallee(ftype_tvm_ffi_env_mod_lookup_from_imports_,
+
RuntimeTVMFFIEnvModLookupFromImports());
+ builder_->CreateRet(
+ builder_->CreateCall(lookup_callee, {module_ctx, GetConstString(fname),
result}));
+
+ function_ = saved_function;
+ builder_->restoreIP(saved_ip);
+ return init_func;
}
llvm::Value* CodeGenCPU::GetPackedFuncHandle(const std::string& fname) {
- // We will store the packed function handle in global space.
- // Initialize it during the first call.
+ // Cache the module-owned packed-function handle in global space with
thread-safe initialization.
#if TVM_LLVM_VERSION >= 200
llvm::DataLayout layout(module_.get()->getDataLayout());
#else
@@ -741,49 +742,25 @@ llvm::Value* CodeGenCPU::GetPackedFuncHandle(const
std::string& fname) {
auto it = func_handle_map_.find(fname);
llvm::GlobalVariable* hptr;
+ llvm::Function* init_func;
if (it == func_handle_map_.end()) {
- // create global location for the handle
- // create the function handle
hptr =
new llvm::GlobalVariable(*module_, t_tvm_func_handle_, false,
llvm::GlobalValue::InternalLinkage, nullptr,
".tvm_func." + fname);
hptr->setAlignment(llvm::Align(align));
hptr->setInitializer(llvm::Constant::getNullValue(t_tvm_func_handle_));
func_handle_map_[fname] = hptr;
+ init_func = CreatePackedFuncInit(fname);
+ func_handle_init_map_[fname] = init_func;
} else {
hptr = it->second;
+ init_func = func_handle_init_map_.at(fname);
}
- // create emit codes that checks and load the function.
- llvm::LLVMContext* ctx = llvm_target_->GetContext();
- llvm::BasicBlock* pre_block = builder_->GetInsertBlock();
- auto* init_block = llvm::BasicBlock::Create(*ctx, "handle_init", function_);
- auto* end_block = llvm::BasicBlock::Create(*ctx, "handle_init_end",
function_);
- llvm::Value* handle = builder_->CreateAlignedLoad(hptr->getValueType(),
hptr, llvm::Align(align));
- llvm::Value* handle_not_null =
- builder_->CreateICmpNE(handle,
llvm::Constant::getNullValue(t_tvm_func_handle_));
- builder_->CreateCondBr(handle_not_null, end_block, init_block,
md_very_likely_branch_);
- // Initialize the handle if needed.
- builder_->SetInsertPoint(init_block);
- llvm::Value* out =
- WithFunctionEntry([&]() { return
builder_->CreateAlloca(t_tvm_func_handle_); });
- llvm::LoadInst* ctx_load =
builder_->CreateAlignedLoad(gv_mod_ctx_->getValueType(), gv_mod_ctx_,
-
llvm::Align(gv_mod_ctx_->getAlignment()));
- ctx_load->setMetadata(
- "tbaa", md_builder_->createTBAAStructTagNode(md_tbaa_ctx_ptr_,
md_tbaa_ctx_ptr_, 0));
- auto env_callee = llvm::FunctionCallee(ftype_tvm_get_func_from_env_,
RuntimeTVMGetFuncFromEnv());
- llvm::Value* retcode = builder_->CreateCall(env_callee, {ctx_load,
GetConstString(fname), out});
- init_block = CheckCallSuccess(retcode);
- llvm::Value* loaded_handle =
- builder_->CreateAlignedLoad(t_tvm_func_handle_, out, llvm::Align(align));
- // Store the handle
- builder_->CreateStore(loaded_handle, hptr);
- builder_->CreateBr(end_block);
- // end block
- builder_->SetInsertPoint(end_block);
- llvm::PHINode* phi = builder_->CreatePHI(t_tvm_func_handle_, 2);
- phi->addIncoming(handle, pre_block);
- phi->addIncoming(loaded_handle, init_block);
- return phi;
+
+ auto init_once_callee =
+ llvm::FunctionCallee(ftype_tvm_ffi_handle_init_once_,
RuntimeTVMFFIHandleInitOnce());
+ CheckCallSuccess(builder_->CreateCall(init_once_callee, {hptr, init_func}));
+ return builder_->CreateAlignedLoad(hptr->getValueType(), hptr,
llvm::Align(align));
}
CodeGenCPU::PackedCall CodeGenCPU::MakeCallPackedLowered(const
ffi::Array<PrimExpr>& args,
@@ -899,10 +876,18 @@ llvm::Value* CodeGenCPU::RuntimeTVMFFIFunctionCall() {
return GetContextPtr(gv_tvm_ffi_func_call_);
}
-llvm::Value* CodeGenCPU::RuntimeTVMGetFuncFromEnv() {
- if (f_tvm_get_func_from_env_ != nullptr) return f_tvm_get_func_from_env_;
- return GetContextPtr(gv_tvm_get_func_from_env_);
+llvm::Value* CodeGenCPU::RuntimeTVMFFIEnvModLookupFromImports() {
+ if (f_tvm_ffi_env_mod_lookup_from_imports_ != nullptr) {
+ return f_tvm_ffi_env_mod_lookup_from_imports_;
+ }
+ return GetContextPtr(gv_tvm_ffi_env_mod_lookup_from_imports_);
+}
+
+llvm::Value* CodeGenCPU::RuntimeTVMFFIHandleInitOnce() {
+ if (f_tvm_ffi_handle_init_once_ != nullptr) return
f_tvm_ffi_handle_init_once_;
+ return GetContextPtr(gv_tvm_ffi_handle_init_once_);
}
+
llvm::Value* CodeGenCPU::RuntimeTVMFFIErrorSetRaisedFromCStr() {
if (f_tvm_ffi_set_raised_by_c_str_ != nullptr) return
f_tvm_ffi_set_raised_by_c_str_;
return GetContextPtr(gv_tvm_ffi_set_last_error_c_str_);
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index d090687b7e..d448d381e7 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -92,15 +92,14 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::FunctionType* ftype_tvm_ffi_c_func_{nullptr};
llvm::FunctionType* ftype_tvm_parallel_lambda_{nullptr};
llvm::FunctionType* ftype_tvm_ffi_func_call_{nullptr};
- llvm::FunctionType* ftype_tvm_get_func_from_env_{nullptr};
+ llvm::FunctionType* ftype_tvm_ffi_env_mod_lookup_from_imports_{nullptr};
+ llvm::FunctionType* ftype_tvm_ffi_handle_init_callback_{nullptr};
+ llvm::FunctionType* ftype_tvm_ffi_handle_init_once_{nullptr};
llvm::FunctionType* ftype_tvm_ffi_error_set_raised_by_c_str_{nullptr};
llvm::FunctionType*
ftype_tvm_ffi_error_set_raised_from_c_str_parts_{nullptr};
llvm::FunctionType* ftype_tvm_parallel_launch_{nullptr};
llvm::FunctionType* ftype_tvm_parallel_barrier_{nullptr};
llvm::FunctionType* ftype_tvm_register_system_symbol_{nullptr};
- // Lazy entry for function call.
- llvm::FunctionType* ftype_tvm_static_init_callback_{nullptr};
- llvm::FunctionType* ftype_tvm_static_init_{nullptr};
private:
// the parallel group information
@@ -117,7 +116,8 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::GlobalVariable* InitContextPtr(llvm::Type* type, std::string name);
llvm::Value* GetContextPtr(llvm::GlobalVariable* gv);
llvm::Value* RuntimeTVMFFIFunctionCall();
- llvm::Value* RuntimeTVMGetFuncFromEnv();
+ llvm::Value* RuntimeTVMFFIEnvModLookupFromImports();
+ llvm::Value* RuntimeTVMFFIHandleInitOnce();
llvm::Value* RuntimeTVMFFIErrorSetRaisedFromCStr();
llvm::Value* RuntimeTVMFFIErrorSetRaisedFromCStrParts();
// Create a temp function to simplify error raising.
@@ -125,6 +125,7 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::Value* RuntimeTVMParallelLaunch();
llvm::Value* RuntimeTVMParallelBarrier();
llvm::Value* CreateStaticHandle();
+ llvm::Function* CreatePackedFuncInit(const std::string& fname);
llvm::Value* GetPackedFuncHandle(const std::string& str);
TypedPointer PackClosureData(const ffi::Array<Var>& fields, uint64_t*
num_bytes,
std::string struct_name = "");
@@ -143,8 +144,6 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::Value* CreateCallPacked(const CallNode* op);
// Create trace call into tvm packed function.
llvm::Value* CreateCallTracePacked(const CallNode* op);
- // Create static initialization
- void CreateStaticInit(const std::string& init_fname, const Stmt& body);
// Create parallel launch
void CreateParallelLaunch(const Stmt& body, int num_task, std::string name =
"");
// Create a new compute scope.
@@ -161,14 +160,16 @@ class CodeGenCPU : public CodeGenLLVM {
// Context for injection lookup
llvm::GlobalVariable* gv_mod_ctx_{nullptr};
llvm::GlobalVariable* gv_tvm_ffi_func_call_{nullptr};
- llvm::GlobalVariable* gv_tvm_get_func_from_env_{nullptr};
+ llvm::GlobalVariable* gv_tvm_ffi_env_mod_lookup_from_imports_{nullptr};
+ llvm::GlobalVariable* gv_tvm_ffi_handle_init_once_{nullptr};
llvm::GlobalVariable* gv_tvm_ffi_set_last_error_c_str_{nullptr};
llvm::GlobalVariable* gv_tvm_parallel_launch_{nullptr};
llvm::GlobalVariable* gv_tvm_parallel_barrier_{nullptr};
std::unordered_map<ffi::String, llvm::GlobalVariable*> gv_func_map_;
// context for direct dynamic lookup
llvm::Function* f_tvm_ffi_func_call_{nullptr};
- llvm::Function* f_tvm_get_func_from_env_{nullptr};
+ llvm::Function* f_tvm_ffi_env_mod_lookup_from_imports_{nullptr};
+ llvm::Function* f_tvm_ffi_handle_init_once_{nullptr};
llvm::Function* f_tvm_ffi_set_raised_by_c_str_{nullptr};
llvm::Function* f_tvm_ffi_set_raised_from_c_str_parts_{nullptr};
llvm::Function* f_tvm_parallel_launch_{nullptr};
@@ -180,6 +181,7 @@ class CodeGenCPU : public CodeGenLLVM {
std::unordered_map<int, llvm::Function*> set_raised_helpers_;
// global to packed function handle
std::unordered_map<std::string, llvm::GlobalVariable*> func_handle_map_;
+ std::unordered_map<std::string, llvm::Function*> func_handle_init_map_;
// List of symbols to be exported to TVM system lib.
std::vector<std::pair<std::string, llvm::Constant*>> export_system_symbols_;
// List of functions to be registered in the FuncRegistry, if generated.
diff --git a/src/target/source/codegen_c_host.cc
b/src/target/source/codegen_c_host.cc
index c48e1b1e65..cd5e24ddf6 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -46,6 +46,7 @@ void CodeGenCHost::Init(bool output_ssa, bool emit_asserts,
bool emit_fwd_func_d
declared_globals_.clear();
decl_stream << "// tvm target: " << target_str << "\n";
decl_stream << "#define TVM_EXPORTS\n";
+ decl_stream << "#include \"tvm/ffi/extra/c_env_api.h\"\n";
decl_stream << "#include \"tvm/runtime/base.h\"\n";
decl_stream << "#include \"tvm/runtime/c_backend_api.h\"\n";
decl_stream << "#include \"tvm/ffi/c_api.h\"\n";
@@ -209,8 +210,8 @@ void CodeGenCHost::PrintGetFuncFromBackend(const
std::string& func_name,
this->stream << "if (" << packed_func_name << " == NULL) {\n";
int packed_func_if_scope = this->BeginScope();
this->PrintIndent();
- this->stream << "if (TVMBackendGetFuncFromEnv(" << module_name_ << ", \"" <<
func_name << "\""
- << ", &" << packed_func_name << ") != 0) {\n";
+ this->stream << "if (TVMFFIEnvModLookupFromImports(" << module_name_ << ",
\"" << func_name
+ << "\", &" << packed_func_name << ") != 0) {\n";
int get_func_env_scope = this->BeginScope();
this->PrintIndent();
this->stream << "return -1;\n";
