TVM学习之从relay到TOPI

lower操作完成从高级算子(relay)到低级算子(topi)的转化。lower开始于以下代码(src/relay/backend/graph_runtime_codegen.cc):
loweredoutput codegen(relay::function func) { auto pf = getpackedfunc(relay.backend.graphplanmemory); storage_device_map_ = (*pf)(func); // first we convert all the parameters into input nodes. for (auto param : func->params) { auto node_ptr = graphinputnode::make_node_ptr(param->name_hint(), graphattrs()); var_map_[param.get()] = addnode(node_ptr, param); } heads_ = visitexpr(func->body); std::ostringstream os; dmlc::jsonwriter writer(&os); getjson(&writer); loweredoutput ret; ret.graph_json = os.str(); ret.params = params_; for (auto& kv : lowered_funcs_) { if (ret.lowered_funcs.count(kv.first) == 0) { ret.lowered_funcs.set(kv.first, irmodule()); } auto& mod = ret.lowered_funcs[kv.first]; mod->update(kv.second); ret.lowered_funcs.set(kv.first, mod); } ret.external_mods = compile_engine_->lowerexternalfunctions(); return ret; } 在完成内存申请优化之后,visitexpr对图进行遍历并lower每个relay算子。我们来看callnode节点的处理。主要代码如下:
auto pf0 = getpackedfunc(relay.backend._make_ccachekey); auto pf1 = getpackedfunc(relay.backend._compileenginelower); target target; // handle external function if (func->getattr(attr::kcompiler).defined()) { target = tvm::target::ext_dev(); ccachekey key = (*pf0)(func, target); cachedfunc ext_func = (*pf1)(compile_engine_, key);这一步是当存在外部compiler的时候,使用外部compiler进行lower。ccachekey将function和target打包到一起,可能是方便后边compiler的调用。而lower函数会调用src/relay/backend/compile_engine.cc中compileengineimpl类中的lowerinternal函数,在这个函数中实现了外部编译器lower和内部lower的代码,如果是有外部compiler参与,其将function,target等打包成ccachevalue返回,等待后边外部编译器统一处理。如果没有外部编译器,那么tvm将对relay算子转换到topi库中算子。cachedfunc lowered_func = (*pf1)(compile_engine_, key); if (!lowered_funcs_.count(target->str())) { lowered_funcs_[target->str()] = irmodule(); } lowered_funcs_[target->str()]->update(lowered_func->funcs);return graphaddcallnode(op, _getuniquename(lowered_func->func_name), lowered_func->func_name); 同样会调用lowerinternal函数,首先建立schedule:
cachedfunc createschedule(const function& source_func, const target& target) { return schedulegetter(target).create(source_func); } 在create函数中,首先将inputs都转换成te的算子表示:
for (var param : prim_func-> params) { array inputs; if (const auto* ttype = param->checked_type().as()) { tvm::te::tensor tensor = tvm::te::placeholder(getshape(ttype-> shape), ttype->dtype); cache_node-> inputs.push_back(tensor); inputs.push_back(tensor); } else { // flatten tuple of tensor type. const auto* tuple_type = param-> type_as (); for (type field : tuple_type-> fields) { const auto* ttype = field.as (); // todo(@icemelon): allow recursive tuple check(ttype != nullptr); tvm::te::tensor tensor = tvm::te::placeholder(getshape(ttype-> shape), ttype-> dtype); cache_node-> inputs.push_back(tensor); inputs.push_back(tensor); } } memo_[param] = inputs;} 然后遍历其它node来实现lower操作。
我们还是来看callnode的访问。
array visitexpr_(const callnode* call_node) final { static auto fpattern = op::getattrmap(toppattern); static auto flower_call = tvm::runtime::registry::get(relay.backend.lower_call); check(flower_call) checked_type().as()) { ++count_tuple; } for (te::tensor tensor : visitexpr(arg)) { inputs.push_back(tensor); } } if (count_tuple) { check_eq(call_node-> args.size(), 1u) opnode> ()) >> primitive function only allows call into primitive ops; op op = downcast>op>(call_node-> op); array>te::tensor> outputs; opimplementation impl; // skip fcompute for device copy operators as it is not registered. if (op == device_copy_op_) { const auto* copy_input = inputs[0].operator->(); outputs.push_back(te::tensor(copy_input->shape, copy_input->dtype, te::operation(), 0)); } else { loweredoutput lowered_out = (*flower_call)(getref>call>(call_node), inputs, target_); outputs = lowered_out->outputs; 这里lower操作会去调用python中注册的lower_call函数,这个函数位于python/tvm/relay/backend/compile_engine.py中。在这个函数中最主要的是select_implementation。
select_implementation是去选择relay算子的一个topi层级的实现方式。同一个relay算子在不同target上有不同实现方式,具体采用哪种方式要依据target的属性。在select_implementation中首先通过gat_valid_implementation获得所有已经注册的实现方式。
fstrategy = op.get_attr(ftvmstrategy) assert fstrategy is not none, %s doesn't have ftvmstrategy registered % op.name with target: strategy = fstrategy(attrs, inputs, out_type, target) analyzer = tvm.arith.analyzer() ret = [] for spec in strategy.specializations: if spec.condition: # check if all the clauses in the specialized condition are true flag = true for clause in spec.condition.clauses: clause = analyzer.canonical_simplify(clause) if isinstance(clause, tvm.tir.intimm) and clause.value: continue flag = false break if flag: for impl in spec.implementations: ret.append(impl) else: for impl in spec.implementations: ret.append(impl)return ret fstrategy指向的是op attr的ftvmstrategy对应的函数。比如con2d注册的策略有:
def conv2d_strategy(attrs, inputs, out_type, target): conv2d generic strategy logger.warning(conv2d is not optimized for this platform.) strategy = _op.opstrategy() data, kernel = inputs dilation = get_const_tuple(attrs.dilation) groups = attrs.groups layout = attrs.data_layout kernel_layout = attrs.kernel_layout (dilation_h, dilation_w) = dilation if dilation_h > 1 or dilation_w > 1: raise valueerror(dilation should be positive value) if groups == 1: if layout == nchw: assert kernel_layout == oihw strategy.add_implementation( wrap_compute_conv2d(topi.nn.conv2d_nchw), wrap_topi_schedule(topi.generic.schedule_conv2d_nchw), nam) 可见一个conv2d即使同一个target也会注册不同的策略。add_implementation将会把compute,schedule的具体函数注册到strategy中。strategy是一个包含了一个relay算子implementation方式的数据结构。它包含了很多opspecialization,每个opspecialization中包含一些列opimplementation,opimplementation中就对应着schedule和compute的具体方式,schedule是一个算子计算的排布,compute是对应了topi库算子。
获得了所有有效implementation之后,会依据两种方式选择,一种是通过auto tvm来自动化搜索最优的实现方式,另外一种在不适用auto tvm工具情况下,会选择plevel最大的implementation。选择好了implementation之后,就调用src/relay/backend/compile_engine.cc中的loweredoutput类建立一个实例。可以看出,lower_call实现了将relay算子统一用更底层的的抽象进行了表示。这种表示中包含了relay算子,以及这个算子的计算方式以及schedule信息。这样就方便后边对其进行schedule优化了。
然后将这些loweredoutput进行打包成cachedfuncnode。cachedfuncnode会作为后边schedule优化的入参。


智慧工厂和工业物联网对中国12个产业GDP的影响
常用保护电器故障分析
什么是COB封装?有哪些优劣势?
基于一款PWM控制型的高效恒流型LED驱动IC设计
电动机转速的微机测量系统
TVM学习之从relay到TOPI
泛海微FS5054 单节锂电池充电IC
解析三相PWM逆变器的主电源电路设计
让存储器运行速度更快
MCU设计低功耗时需要注意的五点
如何调试设计中的时钟域交汇问题
Google拆分智能家居设备商Nest,为什么Android独霸布局?
华为供应商兴森科技五大问题“缠身”,国资解救或成黄粱一梦
浅谈一下汽车连接器中接线端子的性能介绍
一加5什么时候上市?一加5最新消息:一加5骁龙835+8G,抗衡三星S8
专设Z-Wave线上学习中心,助力启动长距离、低功耗Sub-GHz无线设计
全球化趋势明显 2016全球云计算市场发展情况总汇
稳压二极管工作原理和注意事项
一张图就足以看透5G颜值
悬臂梁称重传感器的工作原理