tqchen commented on code in PR #19939:
URL: https://github.com/apache/tvm/pull/19939#discussion_r3523501870
##########
src/target/llvm/codegen_cpu.cc:
##########
@@ -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()));
Review Comment:
Audited in 9a85bb49ee and intentionally kept the explicit aligned load.
`GetContextPtr` unconditionally attaches the `ctx_ptr` TBAA tag, which would
reintroduce the metadata explicitly removed in response to owner review
3523446757. Focused IR inspection confirms the module-context load remains
untagged while the generated initializer and init-once path are unchanged.
##########
src/target/llvm/codegen_cpu.cc:
##########
@@ -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(t_tvm_func_handle_, hptr,
llvm::Align(align));
Review Comment:
Addressed in 9a85bb49ee. The final handle load now uses
`hptr->getValueType()`. This is supported across TVM’s configured LLVM range
and restores the exact pattern used by this code before the init-once refactor;
the LLVM build and 44 focused codegen tests pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
