深度学习编译中间件之NNVM(十二)NNVM源代码阅读1

2023-10-29 05:08
文章标签 编译 学习 深度 阅读 源代码 十二 中间件 nnvm

本文主要是介绍深度学习编译中间件之NNVM(十二)NNVM源代码阅读1，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

参考文档

对于阅读NNVM源代码而言，建议从最外层使用的nnvm.compiler.build函数开始阅读，逐渐深入．

这里先展示一个最简单的NNVM编译器的使用过程：

# 从本地文件加载mxnet模型
mx_sym, args, auxs = mx.model.load_checkpoint('mobilenet', 0)
nnvm_sym, nnvm_params = nnvm.frontend.from_mxnet(mx_sym, args, auxs)# 设置输入数据的shape
batch_size = 1
image_shape = (3, 224, 224)
data_shape = (batch_size,) + image_shape# 进行NNVM编译
with nnvm.compiler.build_config(opt_level = 3):graph, lib, params = nnvm.compiler.build(nnvm_sym, tvm.target.rasp(), shape={"data": data_shape}, params = nnvm_params)# 保存生成的执行so库
lib.export_library("mobilenet_deploy.so")

可以将nnvm.compiler.build的执行过程总结为如下步骤：

  1. 校正Layout
  2. 初始化Pass(指定shape)
  3. 初始化所有变量(_all_var_init)
  4. 应用优化
  5. 预计算裁剪
  6. 融合相邻运算并生成最终so
  7. 保存变量的初始化值到params参数文件中

在介绍具体的步骤之前先介绍graph.apply这个函数：

展示python/nnvm/graph.py的部分代码

class Graph(object):def apply(self, passes):"""Apply passes to the graphParameters----------passes : str or list of strThe passes to be appliedReturns-------g : GraphThe transformed graph."""if isinstance(passes, string_types):passes = [passes]cpass = c_array(ctypes.c_char_p, [c_str(key) for key in passes])ghandle = GraphHandle()npass = nn_uint(len(passes))check_call(_LIB.NNGraphApplyPasses(self.handle, npass, cpass, ctypes.byref(ghandle)))return Graph(ghandle)

从上述代码可以看到graph.apply用于调用后端Pass返回转换之后的图．具体通过NNGraphApplyPasses接口来实现调用．

融合相邻运算并生成最终so

展示python/nnvm/build_module.py的部分代码

# 代码可能有缩减，只展示核心代码
graph = graph_attr.set_shape_inputs(graph, shape)
graph = graph.apply("InferShape")graph = graph_attr.set_dtype_inputs(graph, dtype)graph._set_json_attr("target", str(target), "str")
graph._set_json_attr("target_host", str(target_host), "str")
graph._set_json_attr("opt_level", 1, "int")graph = graph.apply("InferShape").apply("InferType")
graph = graph.apply("GraphFusePartition").apply("GraphFuseCompile")libmod = graph_attr._move_out_module(graph, "module")

上述代码中的graph.apply属于核心代码，这些代码用于调用后端Pass．

graph_attr._move_out_module的定义位于python/nnvm/graph_attr.py

_move_out_module = tvm.get_global_func("nnvm.graph._move_module")

nnvm.graph._move_module的定义位于src/compiler/packed_func_ext.cc

TVM_REGISTER_GLOBAL("nnvm.graph._move_module").set_body([](TVMArgs args, TVMRetValue *rv) {const nnvm::Graph& g = args[0].AsExtension<Graph>();*rv = const_cast<nnvm::Graph*>(&g)->MoveCopyAttr<tvm::runtime::Module>(args[1]);});

Graph.MoveCopyAttr的定义位于include/nnvm/top/graph.h

template<typename T>
inline T Graph::MoveCopyAttr(const std::string& attr_name) {auto it = attrs.find(attr_name);CHECK(it != attrs.end())<< "Cannot find attribute " << attr_name << " in the graph";std::shared_ptr<any> sptr = it->second;attrs.erase(it);if (sptr.unique()) {return std::move(nnvm::get<T>(*sptr));} else {return nnvm::get<T>(*sptr);}
}

从上述代码可以看到graph_attr._move_out_module(graph, "module")访问的是一个tvm::runtime::Module的对象．但是还不清楚这个Module对象是如何生成的，所以需要继续看下去．

在NNVM代码工程中搜索attrs["module"]得到如下代码：

/src/compiler/graph_fuse.cc

// 代码段位于GraphFuseCompile函数中static const PackedFunc& fbuild = GetPackedFunc("nnvm.compiler.build_target");
tvm::runtime::Module module = fbuild(func_list, target, target_host);
ret.attrs["module"] = std::make_shared<any>(std::move(module));

上述代码中fbuild函数是使用GetPackedFunc获得，根据深度学习编译中间件之NNVM(四)TVM设计理念与开发者指南中提到的，此处是使用了C++调用Python函数的方法．

通过全局搜索可以得到nnvm.compiler.build_target的定义位于python/nnvm/build_module.py：

@tvm.register_func("nnvm.compiler.build_target")
def _build(funcs, target, target_host):if target_host == "":target_host = Nonereturn tvm.build(funcs, target=target, target_host=target_host)

nnvm.compiler.build_target调用了tvm.build

tvm.build的定义位于tvm/python/tvm/build_module.py，执行到这里表示对于整个编译过程而言已经完成了NNVM图优化的阶段，进入到TVM代码生成的阶段．

在介绍TVM具体的代码生成过程前，先了解NNVM传送给TVM进行代码生成的数据结构为：

Array<tvm::LoweredFunc> func_list;
// tvm::LoweredFunc数组

这个数据结构包含了被lower的TVM函数的相关信息，是代码生成前的最终数据结构(IR表示)。这里将介绍这个IR表示是如何生成的。

展示nnvm::Graph GraphFuseCompile函数中和lower相关的部分代码：

src/compiler/graph_fuse.cc

fe.compiled_func = GraphLower(fe.subgraph, inputs, target, sub_master_idx);
for (LoweredFunc f : fe.compiled_func->funcs) {if (!func_set.count(f.get())) {func_set.insert(f.get());func_list.push_back(f);}
}

src/compiler/compile_engine.cc

GraphFunc GraphLower(Graph graph,const Array<tvm::Tensor>& inputs,const std::string& target,int master_idx) {return CompileEngine::Global()->Lower(graph, inputs, target, master_idx);
}// CompileEngine::Global()->Lower最终调用了CompileEngine::DoLower函数
// run the actual lowering process
GraphFunc DoLower(Graph graph,const Array<tvm::Tensor>& inputs,const std::string& target,int master_idx) {std::string readable_name;Array<tvm::Tensor> all_args;Array<tvm::Tensor> outputs;Schedule sch;std::tie(sch, all_args, graph) = GetScheduleArgs(graph, inputs, target, master_idx,&readable_name, &outputs);std::shared_ptr<GraphFuncNode> gf = std::make_shared<GraphFuncNode>();gf->target = target;gf->func_name = GetUniqeName(readable_name);gf->inputs = inputs;gf->outputs = outputs;static const PackedFunc& flower = GetPackedFunc("nnvm.compiler.lower");gf->funcs = flower(sch, all_args, gf->func_name, graph);return GraphFunc(gf);
}// DoLower函数中比较重要的有两点
// 1. GetScheduleArgs函数用于生成Schedule参数
// 2. GetPackedFunc("nnvm.compiler.lower")重新调用了TVM的Python接口

GetScheduleArgs函数定义位于src/compiler/compile_engine.cc

    // get schedule and its argsstd::tuple<Schedule, Array<tvm::Tensor>, Graph>GetScheduleArgs(Graph graph,const Array<tvm::Tensor> &inputs,const std::string &target,int master_idx,std::string *readable_name,Array<tvm::Tensor> *outputs) {// shape, type// 获取TVM计算函数和TVM调度函数static auto& fcompute =nnvm::Op::GetAttr<FTVMCompute>("FTVMCompute");static auto& fschedule =nnvm::Op::GetAttr<FTVMSchedule>("FTVMSchedule");// 获取并设置输入Shape和类型std::vector<TShape> ishape;std::vector<int> idtype;for (const tvm::Tensor t : inputs) {std::vector<dim_t> shape;for (Expr v : t->shape) {CHECK(v.as<tvm::ir::IntImm>());shape.push_back(v.as<tvm::ir::IntImm>()->value);}ishape.emplace_back(TShape(shape.begin(), shape.end()));idtype.emplace_back(GetTypeFlag(t->dtype));}graph = pass::InferShape(graph, ishape);graph = pass::InferType(graph, idtype);const ShapeVector& shape_vec = graph.GetAttr<ShapeVector>("shape");const DTypeVector& dtype_vec = graph.GetAttr<DTypeVector>("dtype");const IndexedGraph& idx = graph.indexed_graph();CHECK_EQ(inputs.size(), idx.input_nodes().size());// 设置输入Tensorstd::vector<tvm::Tensor> tensor_vec(idx.num_node_entries());for (size_t i = 0; i < idx.input_nodes().size(); ++i) {uint32_t nid = idx.input_nodes()[i];tensor_vec[idx.entry_id(nid, 0)] = inputs[i];}std::ostringstream readable_name_os;readable_name_os << "fuse";for (uint32_t nid = 0; nid < idx.num_nodes(); ++nid) {const auto& inode = idx[nid];if (inode.source->is_variable()) continue;Array<Tensor> op_inputs, out_info;readable_name_os << "_" << inode.source->op()->name;// input arrayfor (const IndexedGraph::NodeEntry& e : inode.inputs) {const tvm::Tensor& t = tensor_vec[idx.entry_id(e)];CHECK(t.defined());op_inputs.push_back(t);}// output hintfor (uint32_t i = 0; i < inode.source->num_outputs(); ++i) {Array<Expr> shape;for (int64_t x : shape_vec[idx.entry_id(nid, i)]) {CHECK_LE(x, static_cast<int64_t>(std::numeric_limits<int>::max()));shape.push_back(make_const(Int(32), x));}out_info.push_back(placeholder(shape,GetTVMType(dtype_vec[idx.entry_id(nid, i)])));}// 运行一次op，输入数据随机Array<Tensor> out = fcompute[inode.source->op()](inode.source->attrs, op_inputs, out_info);CHECK_EQ(out.size(), inode.source->num_outputs());// schedule on root node, and use master's schedulefor (uint32_t index = 0; index < inode.source->num_outputs(); ++index) {uint32_t eid = idx.entry_id(nid, index);tensor_vec[eid] = out[index];}}// Schedule on final output.Array<Tensor> all_args = inputs;Array<Tensor> outs;for (const IndexedGraph::NodeEntry& e : idx.outputs()) {const tvm::Tensor& t = tensor_vec[idx.entry_id(e)];CHECK(t.defined());outs.push_back(t);all_args.push_back(t);}Schedule sch = fschedule[idx[master_idx].source->op()](idx[master_idx].source->attrs, outs, target);// store extra return valuesif (readable_name != nullptr) {*readable_name = readable_name_os.str();}if (outputs != nullptr) {*outputs = outs;}return std::make_tuple(sch, all_args, graph);}

nnvm.compiler.lower的定义位于tvm/python/tvm/build_module.py

def lower(sch,args,name="default_function",binds=None,simple_mode=False):
TVM代码生成过程

这里先展示tvm.build的部分代码：

if fdevice:mdev = codegen.build_module(fdevice, str(target_device))mhost.import_module(mdev)
return mhost

tvm.build调用了codegen.build_module方法，位于tvm/python/tvm/codegen.py：

from ._ffi.function import _init_apidef build_module(lowered_func, target):"""Build lowered_func into Module.Parameters----------lowered_func : LoweredFuncThe lowered functiontarget : strThe target module type.Returns-------module : ModuleThe corressponding module."""return _Build(lowered_func, target)

codegen._Build的定义位于tvm/src/api/api_codegen.cc:

TVM_REGISTER_API("codegen._Build")
.set_body([](TVMArgs args, TVMRetValue *ret) {if (args[0].IsNodeType<LoweredFunc>()) {*ret = Build({args[0]}, args[1]);} else {*ret = Build(args[0], args[1]);}});

runtime::Module::Build位于tvm/src/codegen/codegen.cc:

runtime::Module Build(const Array<LoweredFunc>& funcs,const std::string& target) {std::string mode = target;size_t pos = mode.find(' ');if (pos != std::string::npos) {mode = mode.substr(0, pos);}std::string build_f_name = "codegen.build_" + mode;// the build function.const PackedFunc* bf = runtime::Registry::Get(build_f_name);CHECK(bf != nullptr)<< "Target " << target << " is not enabled";runtime::Module m = (*bf)(funcs, target);return m;
}

因为这里验证的ARM处理器，所以mode为llvm:

codegen.build_llvm的定义位于tvm/src/codegen/llvm/llvm_module.cc:

TVM_REGISTER_API("codegen.build_llvm")
.set_body([](TVMArgs args, TVMRetValue* rv) {std::shared_ptr<LLVMModuleNode> n = std::make_shared<LLVMModuleNode>();n->Init(args[0], args[1]);*rv = runtime::Module(n);
});

LLVMModuleNode::Init的定义位于tvm/src/codegen/llvm/llvm_module.cc:

void Init(const Array<LoweredFunc>& funcs, std::string target) {InitializeLLVM();tm_ = GetLLVMTargetMachine(target);bool system_lib = (target.find("-system-lib") != std::string::npos);CHECK_NE(funcs.size(), 0U);ctx_ = std::make_shared<llvm::LLVMContext>();std::unique_ptr<CodeGenLLVM> cg = CodeGenLLVM::Create(tm_);entry_func_ = funcs[0]->name;cg->Init(funcs[0]->name, tm_, ctx_.get(), system_lib, system_lib);for (LoweredFunc f :  funcs) {cg->AddFunction(f);}cg->AddMainFunction(funcs[0]->name);module_ = cg->Finish();module_->addModuleFlag(llvm::Module::Warning, "tvm_target",llvm::MDString::get(*ctx_, target));target_ = target;mptr_ = module_.get();
}

LLVMModuleNode::Init函数中和代码生成相关的主要代码调用接口为：

CodeGenLLVM::Create
CodeGenLLVM::Init
CodeGenLLVM::AddFunction
CodeGenLLVM::AddMainFunction
CodeGenLLVM::Finish/*!* \brief Compile and add function f to the current module.* \param f The function to be added.*/
virtual void AddFunction(const LoweredFunc& f);

CodeGenLLVM::AddFunction即是负责编译每一个函数并添加到当前module的函数。

CodeGenLLVM::AddFunction的定义位于tvm/src/codegen/llvm/codegen_llvm.cc:

void CodeGenLLVM::AddFunction(const LoweredFunc& f) {this->AddFunctionInternal(f, false);
}void CodeGenLLVM::AddFunctionInternal(const LoweredFunc& f, bool ret_void) {this->InitFuncState();std::vector<llvm::Type*> arg_types;is_restricted_ = f->is_restricted;for (Var arg : f->args) {Type t = arg.type();if (t.is_handle()) {auto it = f->handle_data_type.find(arg);if (it != f->handle_data_type.end()) {arg_types.push_back(LLVMType((*it).second.type())->getPointerTo(GetGlobalAddressSpace()));} else {arg_types.push_back(t_int8_->getPointerTo(GetGlobalAddressSpace()));}if (!is_restricted_) {alias_var_set_.insert(arg.get());}} else {arg_types.push_back(LLVMType(arg.type()));}}llvm::FunctionType* ftype = llvm::FunctionType::get(ret_void ? t_void_ : t_int_, arg_types, false);CHECK(module_->getFunction(f->name) == nullptr)<< "Function " << f->name << " already exist in module";function_ = llvm::Function::Create(ftype, llvm::Function::ExternalLinkage,f->name, module_.get());function_->setCallingConv(llvm::CallingConv::C);function_->setDLLStorageClass(llvm::GlobalValue::DLLStorageClassTypes::DLLExportStorageClass);// set var map and align informationauto arg_it = function_->arg_begin();for (size_t i = 0; i < f->args.size(); ++i, ++arg_it) {llvm::Argument* v = &(*arg_it);const Var& var = f->args[i];var_map_[var.get()] = v;if (is_restricted_) {if (var.type().is_handle() && !alias_var_set_.count(var.get())) {// set non alias.
#if TVM_LLVM_VERSION >= 50function_->addParamAttr(i, llvm::Attribute::NoAlias);
#elsefunction_->setDoesNotAlias(i + 1);
#endif}}}llvm::BasicBlock* entry = llvm::BasicBlock::Create(*ctx_, "entry", function_);builder_->SetInsertPoint(entry);this->VisitStmt(f->body);if (ret_void) {builder_->CreateRetVoid();} else {builder_->CreateRet(ConstInt32(0));}
}

CodeGenLLVM::AddFunctionInternal函数的主要内部实现细节为：

  1. 确定函数参数和返回值类型，以此确定函数类型
  2. 创建函数(llvm::Function::Create)
  3. 设置函数选项(调用约定、DLL存储类型)
  4. 填充函数参数

这篇关于深度学习编译中间件之NNVM(十二)NNVM源代码阅读1的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！


http://www.chinasem.cn/article/298423。

相关文章

深度解析Java DTO(最新推荐)

深度解析Java DTO(最新推荐)

《深度解析JavaDTO(最新推荐)》DTO（DataTransferObject）是一种用于在不同层（如Controller层、Service层）之间传输数据的对象设计模式,其核心目的是封装数据,... 目录一、什么是DTO？DTO的核心特点：二、为什么需要DTO？（对比Entity）三、实际应用场景解析
阅读更多...
深度解析Java项目中包和包之间的联系

深度解析Java项目中包和包之间的联系

《深度解析Java项目中包和包之间的联系》文章浏览阅读850次，点赞13次，收藏8次。本文详细介绍了Java分层架构中的几个关键包：DTO、Controller、Service和Mapper。_jav... 目录前言一、各大包1.DTO1.1、DTO的核心用途1.2. DTO与实体类(Entity)的区别1
阅读更多...
Visual Studio 2022 编译C++20代码的图文步骤

Visual Studio 2022 编译C++20代码的图文步骤

《VisualStudio2022编译C++20代码的图文步骤》在VisualStudio中启用C++20import功能,需设置语言标准为ISOC++20,开启扫描源查找模块依赖及实验性标... 默认创建Visual Studio桌面控制台项目代码包含C++20的import方法。右键项目的属性：
阅读更多...
深度解析Python装饰器常见用法与进阶技巧

深度解析Python装饰器常见用法与进阶技巧

《深度解析Python装饰器常见用法与进阶技巧》Python装饰器（Decorator）是提升代码可读性与复用性的强大工具,本文将深入解析Python装饰器的原理,常见用法,进阶技巧与最佳实践,希望可... 目录装饰器的基本原理函数装饰器的常见用法带参数的装饰器类装饰器与方法装饰器装饰器的嵌套与组合进阶技巧
阅读更多...
深度解析Spring Boot拦截器Interceptor与过滤器Filter的区别与实战指南

深度解析Spring Boot拦截器Interceptor与过滤器Filter的区别与实战指南

《深度解析SpringBoot拦截器Interceptor与过滤器Filter的区别与实战指南》本文深度解析SpringBoot中拦截器与过滤器的区别,涵盖执行顺序、依赖关系、异常处理等核心差异,并... 目录Spring Boot拦截器（Interceptor）与过滤器（Filter）深度解析：区别、实现
阅读更多...
深度解析Spring AOP @Aspect 原理、实战与最佳实践教程

深度解析Spring AOP @Aspect 原理、实战与最佳实践教程

《深度解析SpringAOP@Aspect原理、实战与最佳实践教程》文章系统讲解了SpringAOP核心概念、实现方式及原理,涵盖横切关注点分离、代理机制（JDK/CGLIB）、切入点类型、性能... 目录1. @ASPect 核心概念1.1 AOP 编程范式1.2 @Aspect 关键特性2. 完整代码实
阅读更多...
SpringBoot开发中十大常见陷阱深度解析与避坑指南

SpringBoot开发中十大常见陷阱深度解析与避坑指南

《SpringBoot开发中十大常见陷阱深度解析与避坑指南》在SpringBoot的开发过程中,即使是经验丰富的开发者也难免会遇到各种棘手的问题,本文将针对SpringBoot开发中十大常见的“坑... 目录引言一、配置总出错？是不是同时用了.properties和.yml？二、换个位置配置就失效？搞清楚加
阅读更多...
Go学习记录之runtime包深入解析

Go学习记录之runtime包深入解析

《Go学习记录之runtime包深入解析》Go语言runtime包管理运行时环境,涵盖goroutine调度、内存分配、垃圾回收、类型信息等核心功能,：本文主要介绍Go学习记录之runtime包的... 目录前言：一、runtime包内容学习1、作用：① Goroutine和并发控制：② 垃圾回收：③ 栈和
阅读更多...
Python中文件读取操作漏洞深度解析与防护指南

Python中文件读取操作漏洞深度解析与防护指南

《Python中文件读取操作漏洞深度解析与防护指南》在Web应用开发中,文件操作是最基础也最危险的功能之一,这篇文章将全面剖析Python环境中常见的文件读取漏洞类型,成因及防护方案,感兴趣的小伙伴可... 目录引言一、静态资源处理中的路径穿越漏洞1.1 典型漏洞场景1.2 os.path.join()的陷
阅读更多...
Android学习总结之Java和kotlin区别超详细分析

Android学习总结之Java和kotlin区别超详细分析

《Android学习总结之Java和kotlin区别超详细分析》Java和Kotlin都是用于Android开发的编程语言,它们各自具有独特的特点和优势,：本文主要介绍Android学习总结之Ja... 目录一、空安全机制真题 1：Kotlin 如何解决 Java 的 NullPointerExceptio
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2