Unverified Commit fdf035e8 by Tianqi Chen Committed by GitHub

[CODEGEN][LLVM] Cache packed func ptr, lift alloca (#2070)

parent 292ab02d
...@@ -47,8 +47,10 @@ class CodeGenAMDGPU : public CodeGenLLVM { ...@@ -47,8 +47,10 @@ class CodeGenAMDGPU : public CodeGenLLVM {
if (info.scope.rank == runtime::StorageRank::kLocal) { if (info.scope.rank == runtime::StorageRank::kLocal) {
// const int local_address_space = 5; // const int local_address_space = 5;
// TODO(tqchen): for higher version of LLVM, local address space can be set. // TODO(tqchen): for higher version of LLVM, local address space can be set.
llvm::AllocaInst* alloca = builder_->CreateAlloca( llvm::AllocaInst* alloca = WithFunctionEntry([&]() {
LLVMType(op->type), ConstInt32(constant_size)); return builder_->CreateAlloca(
LLVMType(op->type), ConstInt32(constant_size));
});
if (alloca->getAlignment() < static_cast<uint32_t>(info.alignment)) { if (alloca->getAlignment() < static_cast<uint32_t>(info.alignment)) {
alloca->setAlignment(info.alignment); alloca->setAlignment(info.alignment);
} }
......
...@@ -503,7 +503,9 @@ llvm::Value* CodeGenCPU::GetPackedFuncHandle(const std::string& fname) { ...@@ -503,7 +503,9 @@ llvm::Value* CodeGenCPU::GetPackedFuncHandle(const std::string& fname) {
handle_not_null, end_block, init_block, md_very_likely_branch_); handle_not_null, end_block, init_block, md_very_likely_branch_);
// Initialize the handle if needed. // Initialize the handle if needed.
builder_->SetInsertPoint(init_block); builder_->SetInsertPoint(init_block);
llvm::Value* out = builder_->CreateAlloca(t_tvm_func_handle_); llvm::Value* out = WithFunctionEntry([&]() {
return builder_->CreateAlloca(t_tvm_func_handle_);
});
llvm::LoadInst* ctx = builder_->CreateAlignedLoad( llvm::LoadInst* ctx = builder_->CreateAlignedLoad(
gv_mod_ctx_, gv_mod_ctx_->getAlignment()); gv_mod_ctx_, gv_mod_ctx_->getAlignment());
ctx->setMetadata( ctx->setMetadata(
...@@ -513,6 +515,8 @@ llvm::Value* CodeGenCPU::GetPackedFuncHandle(const std::string& fname) { ...@@ -513,6 +515,8 @@ llvm::Value* CodeGenCPU::GetPackedFuncHandle(const std::string& fname) {
RuntimeTVMGetFuncFromEnv(), {ctx, GetConstString(fname), out}); RuntimeTVMGetFuncFromEnv(), {ctx, GetConstString(fname), out});
init_block = CheckCallSuccess(retcode); init_block = CheckCallSuccess(retcode);
llvm::Value* loaded_handle = builder_->CreateAlignedLoad(out, align); llvm::Value* loaded_handle = builder_->CreateAlignedLoad(out, align);
// Store the handle
builder_->CreateStore(loaded_handle, hptr);
builder_->CreateBr(end_block); builder_->CreateBr(end_block);
// end block // end block
builder_->SetInsertPoint(end_block); builder_->SetInsertPoint(end_block);
...@@ -637,19 +641,23 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const Call* op) { ...@@ -637,19 +641,23 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const Call* op) {
} else if (op->is_intrinsic(intrinsic::tvm_stack_alloca)) { } else if (op->is_intrinsic(intrinsic::tvm_stack_alloca)) {
CHECK_EQ(op->args.size(), 2U); CHECK_EQ(op->args.size(), 2U);
const std::string& type = op->args[0].as<StringImm>()->value; const std::string& type = op->args[0].as<StringImm>()->value;
llvm::Value* num = MakeValue(op->args[1]); return WithFunctionEntry([&]() -> llvm::AllocaInst* {
if (type == "shape") { const int64_t* pval = as_const_int(op->args[1]);
return builder_->CreateAlloca(t_tvm_shape_index_, num); CHECK(pval) << "require stack alloca to contain constant value";
} else if (type == "arg_value") { llvm::Value* num = ConstInt32(pval[0]);
return builder_->CreateAlloca(t_tvm_value_, num); if (type == "shape") {
} else if (type == "arg_tcode") { return builder_->CreateAlloca(t_tvm_shape_index_, num);
return builder_->CreateAlloca(t_int_, num); } else if (type == "arg_value") {
} else if (type == "array") { return builder_->CreateAlloca(t_tvm_value_, num);
return builder_->CreateAlloca(t_tvm_array_, num); } else if (type == "arg_tcode") {
} else { return builder_->CreateAlloca(t_int_, num);
LOG(FATAL) << "Unknown stack alloca type " << type; } else if (type == "array") {
return nullptr; return builder_->CreateAlloca(t_tvm_array_, num);
} } else {
LOG(FATAL) << "Unknown stack alloca type " << type;
return nullptr;
}
});
} else { } else {
return CodeGenLLVM::CreateIntrinsic(op); return CodeGenLLVM::CreateIntrinsic(op);
} }
......
...@@ -1049,8 +1049,10 @@ void CodeGenLLVM::VisitStmt_(const Allocate* op) { ...@@ -1049,8 +1049,10 @@ void CodeGenLLVM::VisitStmt_(const Allocate* op) {
if (info.alignment > 16) { if (info.alignment > 16) {
info.alignment = 16; info.alignment = 16;
} }
llvm::AllocaInst* alloca = builder_->CreateAlloca( llvm::AllocaInst* alloca = WithFunctionEntry([&]() {
LLVMType(op->type), ConstInt32(constant_size)); return builder_->CreateAlloca(
LLVMType(op->type), ConstInt32(constant_size));
});
if (alloca->getAlignment() < static_cast<uint32_t>(info.alignment)) { if (alloca->getAlignment() < static_cast<uint32_t>(info.alignment)) {
alloca->setAlignment(info.alignment); alloca->setAlignment(info.alignment);
} }
......
...@@ -132,6 +132,26 @@ class CodeGenLLVM : ...@@ -132,6 +132,26 @@ class CodeGenLLVM :
/*! \brief The alignment of allocation */ /*! \brief The alignment of allocation */
int alignment{0}; int alignment{0};
}; };
/*!
* \brief Execute falloca at the beginning of the
* currrent function and obtain its return value.
*
* This is a helper function to make sure that
* alloca always happen in the beginning of the function.
*
* \param falloca The allocation function to be executed.
* \tparam F The function to be executed.
* \return The result.
*/
template<typename F>
inline llvm::AllocaInst* WithFunctionEntry(F falloca) {
llvm::BasicBlock* current = builder_->GetInsertBlock();
llvm::BasicBlock* entry = &(function_->getEntryBlock());
builder_->SetInsertPoint(entry, entry->begin());
llvm::AllocaInst* res = falloca();
builder_->SetInsertPoint(current);
return res;
}
// create intrinstic given call // create intrinstic given call
virtual llvm::Value* CreateIntrinsic(const Call* op); virtual llvm::Value* CreateIntrinsic(const Call* op);
// create extern function call // create extern function call
......
...@@ -49,8 +49,10 @@ class CodeGenNVPTX : public CodeGenLLVM { ...@@ -49,8 +49,10 @@ class CodeGenNVPTX : public CodeGenLLVM {
if (info.scope.rank == runtime::StorageRank::kLocal) { if (info.scope.rank == runtime::StorageRank::kLocal) {
// const int local_address_space = 5; // const int local_address_space = 5;
// TODO(tqchen): for higher version of LLVM, local address space can be set. // TODO(tqchen): for higher version of LLVM, local address space can be set.
llvm::AllocaInst* alloca = builder_->CreateAlloca( llvm::AllocaInst* alloca = WithFunctionEntry([&]() {
LLVMType(op->type), ConstInt32(constant_size)); return builder_->CreateAlloca(
LLVMType(op->type), ConstInt32(constant_size));
});
if (alloca->getAlignment() < static_cast<uint32_t>(info.alignment)) { if (alloca->getAlignment() < static_cast<uint32_t>(info.alignment)) {
alloca->setAlignment(info.alignment); alloca->setAlignment(info.alignment);
} }
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment