深度学习编译中间件之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。

相关文章

IDEA编译报错“java: 常量字符串过长”的原因及解决方法

IDEA编译报错“java: 常量字符串过长”的原因及解决方法

《IDEA编译报错“java:常量字符串过长”的原因及解决方法》今天在开发过程中,由于尝试将一个文件的Base64字符串设置为常量,结果导致IDEA编译的时候出现了如下报错java:常量字符串过长,... 目录一、问题描述二、问题原因2.1 理论角度2.2 源码角度三、解决方案解决方案①：StringBui
阅读更多...
Java深度学习库DJL实现Python的NumPy方式

Java深度学习库DJL实现Python的NumPy方式

《Java深度学习库DJL实现Python的NumPy方式》本文介绍了DJL库的背景和基本功能,包括NDArray的创建、数学运算、数据获取和设置等,同时,还展示了如何使用NDArray进行数据预处理... 目录1 NDArray 的背景介绍1.1 架构2 JavaDJL使用2.1 安装DJL2.2 基本操
阅读更多...
最长公共子序列问题的深度分析与Java实现方式

最长公共子序列问题的深度分析与Java实现方式

《最长公共子序列问题的深度分析与Java实现方式》本文详细介绍了最长公共子序列（LCS）问题,包括其概念、暴力解法、动态规划解法,并提供了Java代码实现,暴力解法虽然简单,但在大数据处理中效率较低,... 目录最长公共子序列问题概述问题理解与示例分析暴力解法思路与示例代码动态规划解法DP 表的构建与意义动
阅读更多...
解决IDEA使用springBoot创建项目,lombok标注实体类后编译无报错,但是运行时报错问题

解决IDEA使用springBoot创建项目,lombok标注实体类后编译无报错,但是运行时报错问题

《解决IDEA使用springBoot创建项目,lombok标注实体类后编译无报错,但是运行时报错问题》文章详细描述了在使用lombok的@Data注解标注实体类时遇到编译无误但运行时报错的问题,分析... 目录问题分析问题解决方案步骤一步骤二步骤三总结问题使用lombok注解@Data标注实体类，编译时
阅读更多...
Go中sync.Once源码的深度讲解

Go中sync.Once源码的深度讲解

《Go中sync.Once源码的深度讲解》sync.Once是Go语言标准库中的一个同步原语,用于确保某个操作只执行一次,本文将从源码出发为大家详细介绍一下sync.Once的具体使用,x希望对大家有... 目录概念简单示例源码解读总结概念sync.Once是Go语言标准库中的一个同步原语，用于确保某个操
阅读更多...
五大特性引领创新! 深度操作系统 deepin 25 Preview预览版发布

五大特性引领创新! 深度操作系统 deepin 25 Preview预览版发布

《五大特性引领创新!深度操作系统deepin25Preview预览版发布》今日，深度操作系统正式推出deepin25Preview版本，该版本集成了五大核心特性：磐石系统、全新DDE、Tr... 深度操作系统今日发布了 deepin 25 Preview，新版本囊括五大特性：磐石系统、全新 DDE、Tree
阅读更多...
Python3中Sanic中间件的使用

Python3中Sanic中间件的使用

《Python3中Sanic中间件的使用》Sanic框架中的中间件是一种强大的工具,本文就来介绍Python3中Sanic中间件的使用,具有一定的参考价值,感兴趣的可以了解一下... 目录Sanic 中间件的工作流程中间件的使用1. 全局中间件2. 路由中间件3. 异常处理中间件4. 异步中间件5. 优先级
阅读更多...
Node.js 中 http 模块的深度剖析与实战应用小结

Node.js 中 http 模块的深度剖析与实战应用小结

《Node.js中http模块的深度剖析与实战应用小结》本文详细介绍了Node.js中的http模块,从创建HTTP服务器、处理请求与响应,到获取请求参数,每个环节都通过代码示例进行解析,旨在帮... 目录Node.js 中 http 模块的深度剖析与实战应用一、引言二、创建 HTTP 服务器：基石搭建（一
阅读更多...
HarmonyOS学习(七)——UI（五）常用布局总结

HarmonyOS学习(七)——UI（五）常用布局总结

自适应布局 1.1、线性布局（LinearLayout） 通过线性容器Row和Column实现线性布局。Column容器内的子组件按照垂直方向排列，Row组件中的子组件按照水平方向排列。 属性说明space通过space参数设置主轴上子组件的间距，达到各子组件在排列上的等间距效果alignItems设置子组件在交叉轴上的对齐方式，且在各类尺寸屏幕上表现一致，其中交叉轴为垂直时，取值为Vert
阅读更多...
Ilya-AI分享的他在OpenAI学习到的15个提示工程技巧

Ilya-AI分享的他在OpenAI学习到的15个提示工程技巧

Ilya（不是本人，claude AI）在社交媒体上分享了他在OpenAI学习到的15个Prompt撰写技巧。 以下是详细的内容： 提示精确化:在编写提示时,力求表达清晰准确。清楚地阐述任务需求和概念定义至关重要。例:不用"分析文本",而用"判断这段话的情感倾向:积极、消极还是中性"。 快速迭代:善于快速连续调整提示。熟练的提示工程师能够灵活地进行多轮优化。例:从"总结文章"到"用
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2