gemini-code-assist[bot] commented on code in PR #19939:
URL: https://github.com/apache/tvm/pull/19939#discussion_r3523491994
##########
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:

Instead of manually loading the context pointer from `gv_mod_ctx_` with
`CreateAlignedLoad`, we can use the existing helper method
`GetContextPtr(gv_mod_ctx_)`. This is cleaner, avoids code duplication, and
ensures that the TBAA (Type-Based Alias Analysis) metadata is correctly
attached to the load instruction, which can help LLVM optimize the generated
code.
```c
llvm::Value* module_ctx = GetContextPtr(gv_mod_ctx_);
```
##########
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:

Using `hptr->getValueType()` instead of the hardcoded `t_tvm_func_handle_`
is more robust and consistent with modern LLVM opaque pointer patterns, as well
as other load instructions in this file.
```suggestion
return builder_->CreateAlignedLoad(hptr->getValueType(), hptr,
llvm::Align(align));
```
--
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]