TVM代码生成流程

runtime::ModuleRelayBuildCreate(){
autoexec=make_object<RelayBuildModule>();
returnruntime::Module(exec);
}
TVM_REGISTER_GLOBAL("relay.build_module._BuildModule")
.set_body([](TVMArgsargs,TVMRetValue*rv){
*rv=RelayBuildCreate();
});

PackedFuncGetFunction(conststd::string&name,
constObjectPtr<Object>&sptr_to_self)final{
// ...
if(name=="build"){
returnPackedFunc([sptr_to_self,this](TVMArgsargs,TVMRetValue*rv){
CHECK_EQ(args.num_args,3);
this->Build(args[0],args[1],args[2]);
});
// ...
}

voidBuildRelay(
Functionfunc,
conststd::unordered_map<std::string,tvm::runtime::NDArray>&params){
// Optimize input Relay Function and returns Relay Module
relay::Modulerelay_module=Optimize(func,targets_,params);
// Get the updated function.
func=relay_module->Lookup("main");
// Generate code for the updated function.
graph_codegen_=std::unique_ptr<GraphCodegen>(newGraphCodegen());
graph_codegen_->Init(nullptr,targets_);
graph_codegen_->Codegen(func);
ret_.graph_json=graph_codegen_->GetJSON();
ret_.params=graph_codegen_->GetParams();
autolowered_funcs=graph_codegen_->GetLoweredFunc();
if(lowered_funcs.size()==0){
LOG(WARNING)<<"no lowered funcs exist in the compiled module";
}else{
ret_.mod=tvm::build(
lowered_funcs,
target_host_,
BuildConfig::Current());
}
}

relay::ModuleOptimize(
Functionfunc,
constTargetsMap&targets,
conststd::unordered_map<std::string,runtime::NDArray>&params){
// BindParamsByName(func, params)
// Perform Module->Module optimizations.
relay::Modulerelay_module=relay::ModuleNode::FromExpr(func);
Array<Pass>pass_seqs;
// Run all dialect legalization passes.
// ...
pass_seqs.push_back(transform::SimplifyInference());
//
// ...fskip
//
pass_seqs.push_back(transform::EliminateCommonSubexpr(fskip));
pass_seqs.push_back(transform::CombineParallelConv2D(3));
pass_seqs.push_back(transform::CombineParallelDense(3));
pass_seqs.push_back(transform::FoldConstant());
pass_seqs.push_back(transform::FoldScaleAxis());
pass_seqs.push_back(transform::CanonicalizeCast());
pass_seqs.push_back(transform::CanonicalizeOps());
// ...AlterOpLayout
pass_seqs.push_back(transform::FoldConstant());
// Create a sequential pass and perform optimizations.
transform::Passseq=transform::Sequential(pass_seqs);
// ... judge & do
relay_module=seq(relay_module);
// Handle heterogeneous compilation.
transform::PassContextpass_ctx=PassContext::Current();
if(targets_.size()>1){
relay_module=
RunDeviceAnnotationPass(relay_module,pass_ctx->fallback_device);
}
// Fuse the operations if it is needed.
relay_module=transform::FuseOps()(relay_module);
relay_module=transform::InferType()(relay_module);
CHECK(relay_module.defined());
returnrelay_module;
}

// Build for heterogeneous execution.
runtime::Modulebuild(constMap<Target,Array<LoweredFunc>>&inputs,
constTarget&target_host,
constBuildConfig&config){
Array<LoweredFunc>fhost_all;
std::vector<runtime::Module>device_modules;
Targettarget_host_val=target_host;
if(!target_host.defined()){
for(constauto&it:inputs){
if(it.first->device_type==kDLCPU){
target_host_val=it.first;
break;
}
}
}
if(!target_host_val.defined()){
target_host_val=DefaultTargetHost(target_host_val);
}
for(constauto&it:inputs){
autohost_dev_funcs=
split_dev_host_funcs(it.second,it.first,target_host_val,config);
auto&fhost=host_dev_funcs[0];
auto&fdevice=host_dev_funcs[1];
// Get the module for a certain target.
runtime::Modulemdev=DeviceBuild(fdevice,it.first);
for(constauto&it:fhost){
fhost_all.push_back(it);
}
device_modules.push_back(mdev);
}
runtime::Modulemhost=codegen::Build(fhost_all,target_host_val->str());
// Import all modules
for(constauto&it:device_modules){
if(it.operator->()){
mhost.Import(it);
}
}
returnmhost;
}

runtime::ModuleBuild(constArray<LoweredFunc>&funcs,
conststd::string&target){
// do something
std::stringbuild_f_name="codegen.build_"+mode;
// the build function.
constPackedFunc*bf=runtime::Registry::Get(build_f_name);
runtime::Modulem=transformed_funcs.empty()?
(*bf)(funcs,target):
(*bf)(transformed_funcs,target);
returnm;
}