深度学习编译中间件之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 maven编译报错Java heap space的解决方法

idea maven编译报错Java heap space的解决方法

《ideamaven编译报错Javaheapspace的解决方法》这篇文章主要为大家详细介绍了ideamaven编译报错Javaheapspace的相关解决方法,文中的示例代码讲解详细,感兴趣的... 目录1.增加 Maven 编译的堆内存2. 增加 IntelliJ IDEA 的堆内存3. 优化 Mave
阅读更多...
Java编译生成多个.class文件的原理和作用

Java编译生成多个.class文件的原理和作用

《Java编译生成多个.class文件的原理和作用》作为一名经验丰富的开发者,在Java项目中执行编译后,可能会发现一个.java源文件有时会产生多个.class文件,从技术实现层面详细剖析这一现象... 目录一、内部类机制与.class文件生成成员内部类（常规内部类）局部内部类（方法内部类）匿名内部类二、
阅读更多...
SpringCloud动态配置注解@RefreshScope与@Component的深度解析

SpringCloud动态配置注解@RefreshScope与@Component的深度解析

《SpringCloud动态配置注解@RefreshScope与@Component的深度解析》在现代微服务架构中,动态配置管理是一个关键需求,本文将为大家介绍SpringCloud中相关的注解@Re... 目录引言1. @RefreshScope 的作用与原理1.1 什么是 @RefreshScope1.
阅读更多...
Python 中的异步与同步深度解析(实践记录)

Python 中的异步与同步深度解析(实践记录)

《Python中的异步与同步深度解析(实践记录)》在Python编程世界里,异步和同步的概念是理解程序执行流程和性能优化的关键,这篇文章将带你深入了解它们的差异,以及阻塞和非阻塞的特性,同时通过实际... 目录python中的异步与同步：深度解析与实践异步与同步的定义异步同步阻塞与非阻塞的概念阻塞非阻塞同步
阅读更多...
Redis中高并发读写性能的深度解析与优化

Redis中高并发读写性能的深度解析与优化

《Redis中高并发读写性能的深度解析与优化》Redis作为一款高性能的内存数据库,广泛应用于缓存、消息队列、实时统计等场景,本文将深入探讨Redis的读写并发能力,感兴趣的小伙伴可以了解下... 目录引言一、Redis 并发能力概述1.1 Redis 的读写性能1.2 影响 Redis 并发能力的因素二、
阅读更多...
最新Spring Security实战教程之表单登录定制到处理逻辑的深度改造(最新推荐)

最新Spring Security实战教程之表单登录定制到处理逻辑的深度改造(最新推荐)

《最新SpringSecurity实战教程之表单登录定制到处理逻辑的深度改造(最新推荐)》本章节介绍了如何通过SpringSecurity实现从配置自定义登录页面、表单登录处理逻辑的配置,并简单模拟... 目录前言改造准备开始登录页改造自定义用户名密码登陆成功失败跳转问题自定义登出前后端分离适配方案结语前言
阅读更多...
Java进阶学习之如何开启远程调式

Java进阶学习之如何开启远程调式

《Java进阶学习之如何开启远程调式》Java开发中的远程调试是一项至关重要的技能,特别是在处理生产环境的问题或者协作开发时,：本文主要介绍Java进阶学习之如何开启远程调式的相关资料,需要的朋友... 目录概述Java远程调试的开启与底层原理开启Java远程调试底层原理JVM参数总结&nbsMbKKXJx
阅读更多...
Redis 内存淘汰策略深度解析(最新推荐)

Redis 内存淘汰策略深度解析(最新推荐)

《Redis内存淘汰策略深度解析(最新推荐)》本文详细探讨了Redis的内存淘汰策略、实现原理、适用场景及最佳实践,介绍了八种内存淘汰策略,包括noeviction、LRU、LFU、TTL、Rand... 目录一、 内存淘汰策略概述二、内存淘汰策略详解2.1 ​noeviction（不淘汰）​2.2 ​LR
阅读更多...
Python与DeepSeek的深度融合实战

Python与DeepSeek的深度融合实战

《Python与DeepSeek的深度融合实战》Python作为最受欢迎的编程语言之一,以其简洁易读的语法、丰富的库和广泛的应用场景,成为了无数开发者的首选,而DeepSeek,作为人工智能领域的新星... 目录一、python与DeepSeek的结合优势二、模型训练1. 数据准备2. 模型架构与参数设置3
阅读更多...
IDEA编译报错“java: 常量字符串过长”的原因及解决方法

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

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

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2