TVM图编译器Relay简单探究
- 1. 前言
  - 3.1 Part I

TVM图编译器Relay简单探究

## Relay 探究

1. 前言

Relay 是 TVM 中用来替代 NNVM 的模块，其本身被认为是 NNVM 第二代。在设计上，Relay 被认为相对 NNVM 有以下优势：1. 有文本形式中间表示，便于开发和 debug2. 支持子图函数、联合模块，便于联合优化3. 前端用户友好其介绍信息可以在这里找到。相比于最初的 NNVM ，Relay 融合进了编程语言领域的知识，带来了许多新的特色（如 let 表达式，支持递归等等），这也使得 Relay 既具有一个编程语言的特色，也具有常规深度学习网络构建的能力。关于其语言特色的介绍，将在后续文章中展现。本文的重点放在 Relay 如何构建神经网络计算图并在计算图上做图优化。## 2. 使用在构建神经网络的时候，Relay使用起来与 NNVM 非常相似，比如下面的例子就是一个Conv2d + BatchNorm + ReLU 的简单网络：def batch_norm_infer(data, gamma=None, beta=None, moving_mean=None, moving_var=None, kwargs): name = kwargs.get(“name”) kwargs.pop(“name”) if not gamma: gamma = relay.var(name + “_gamma”) if not beta: beta = relay.var(name + “_beta”) if not moving_mean: moving_mean = relay.var(name + “_moving_mean”) if not moving_var: moving_var = relay.var(name + “_moving_var”) return relay.nn.batch_norm(data, gamma=gamma, beta=beta, moving_mean=moving_mean, moving_var=moving_var, kwargs)[0]

def conv2d(data, weight=None, kwargs): name = kwargs.get(“name”) kwargs.pop(“name”) if not weight: weight = relay.var(name + “_weight”) return relay.nn.conv2d(data, weight, kwargs)

def conv_block(data, name, channels, kernel_size=(3, 3), strides=(1, 1), padding=(1, 1), epsilon=1e-5): conv = conv2d( data=data, channels=channels, kernel_size=kernel_size, strides=strides, padding=padding, data_layout=’NCHW’, name=name+’_conv’) bn = batch_norm_infer(data=conv, epsilon=epsilon, name=name + ‘_bn’) act = relay.nn.relu(data=bn) return act
为了运行上述网络，可以写下述代码：data_shape = (1, 3, 224, 224)kernel_shape = (32, 3, 3, 3)dtype = “float32”data = relay.var(“data”, shape=data_shape, dtype=dtype)act = conv_block(data, “graph”, 32, strides=(2, 2))func = relay.Function(relay.ir_pass.free_vars(act), act)
target = “cuda”ctx = tvm.context(target, 0)
net = relay.ir_pass.infer_type(func)shape_dict = { v.name_hint : v.checked_type for v in net.params}params = {}for k, v in shape_dict.items(): if k == “data”: continue init_value = np.random.uniform(-1, 1, v.concrete_shape).astype(v.dtype) params[k] = tvm.nd.array(init_value, ctx=ctx)
with relay.build_config(opt_level=3): graph, lib, params = relay.build(net, target, params=params)
module = runtime.create(graph, lib, ctx)data_tvm = tvm.nd.array((np.random.uniform(-1, 1, size=data_shape)).astype(dtype))module.set_input(‘data’, data_tvm)module.set_input(**params)
module.run()output = module.get_output(0)
这里面大部分代码都在做数据准备工作，我们关注的调用是relay.build，这里面含有如何从 Relay 程序编译出可执行的代码的过程。再继续探究之前，我们先窥探一下优化前的 IR 的样子：fn (%data: Tensor[(1,3,224,224), float32],%graph_conv_weight,%graph_bn_gamma,%graph_bn_beta,%graph_bn_moving_mean,%graph_bn_moving_var) {%0=nn.conv2d(%data,%graph_conv_weight, strides=[2,2], padding=[1,1], channels=32, kernel_size=[3,3])%1=nn.batch_norm(%0,%graph_bn_gamma,%graph_bn_beta,%graph_bn_moving_mean,%graph_bn_moving_var)%2=%1.0%3=nn.relu(%2)%3}## 3. 追踪 relay.build

3.1 Part I

与 NNVM 相似，Relay 先会寻找 AutoTVM 是否有预先 tune 好的参数记录，如果没有就使用 fallback 的 context：ifisinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):ifisinstance(target,dict):tophub_context=autotvm.tophub.context(list(target.values()))else:tophub_context=autotvm.tophub.context(target)else:tophub_context=autotvm.util.EmptyContext()接下来所有的操作都在 tophub_context 的 scope 之下(with tophub_context:)。值得一提的是 Relay 考虑了异构情景下的代码生成，用户可以指定多个生成代码的目标(target)。### 3.2 Part II 下面进行目标无关优化：func=optimize(func, target, params)跟踪进 optimize 函数，首先进行的优化是:ifcfg.pass_enabled(“SimplifyInference”):func=ir_pass.infer_type(func)func=ir_pass.simplify_inference(func)这里做的是针对 inference 情景下的特殊优化，回忆我们在对NNVM 的探究中，NNVM 在目标无关优化中第一步做的是 layout 的变换，而 Relay 则是调换了优化顺序，将 layout 变换放在了后面。这一步做的优化仍然是将 BatchNorm 展开以及去掉 dropout，与 NNVM 是相似的。优化后的 IR 为：fn (%data: Tensor[(1,3,224,224), float32])->Tensor[(1,32,112,112), float32] {%0=meta[relay.Constant][0]# ty=Tensor[(32, 3, 3, 3), float32]%1=nn.conv2d(%data,%0, strides=[2,2], padding=[1,1], channels=32, kernel_size=[3,3])# ty=Tensor[(1, 32, 112, 112), float32]%2=meta[relay.Constant][1]# ty=Tensor[(32,), float32]%3=add(%2,1e-05f)%4=sqrt(%3)%5=divide(1f,%4)%6=meta[relay.Constant][2]# ty=Tensor[(32,), float32]%7=multiply(%5,%6)%8=expand_dims(%7, axis=1, num_newaxis=2)%9=multiply(%1,%8)%10=meta[relay.Constant][3]# ty=Tensor[(32,), float32]%11=negative(%10)%12=multiply(%11,%7)%13=meta[relay.Constant][4]# ty=Tensor[(32,), float32]%14=add(%12,%13)%15=expand_dims(%14, axis=1, num_newaxis=2)%16=add(%9,%15)%17=nn.relu(%16)%17}其次进行的优化是去除公共子表达式：ifcfg.pass_enabled(“EliminateCommonSubexpr”):deffskip(expr):ifisinstance(expr, _expr.Call)andexpr.op.name==’cast’and\expr.attrs.dtype==’int32’:returnTruereturnFalse
func=ir_pass.infer_type(func)func=ir_pass.eliminate_common_subexpr(func, fskip)然后是分支卷积优化：ifcfg.pass_enabled(“CombineParallelConv2D”):func=ir_pass.infer_type(func)func=ir_pass.combine_parallel_conv2d(func)这部分优化会将具有相同输入的卷积合并成一个大的卷积运算。接着是常量传播优化，与 NNVM 相似：ifcfg.pass_enabled(“FoldConstant”):func=ir_pass.fold_constant(func)
ifcfg.pass_enabled(“FoldScaleAxis”):func=ir_pass.infer_type(func)func=ir_pass.backward_fold_scale_axis(func)func=ir_pass.infer_type(func)func=ir_pass.forward_fold_scale_axis(func)func=ir_pass.fold_constant(func)常量传播优化后的 IR 是：fn (%data: Tensor[(1,3,224,224), float32])->Tensor[(1,32,112,112), float32] {%0=meta[relay.Constant][0]%1=nn.conv2d(%data,%0, strides=[2,2], padding=[1,1], channels=32, kernel_size=[3,3])%2=meta[relay.Constant][1]# ty=Tensor[(32, 1, 1), float32]%3=add(%1,%2)%4=nn.relu(%3)%4}下面是进行规范化，想法在于将一些特殊运算转换成等价的常规算子运算，便于后续的分析优化：ifcfg.pass_enabled(“CanonicalizeOps”):func=ir_pass.infer_type(func)func=ir_pass.canonicalize_ops(func)其实 Relay 也只实现了将 bias_add 转换为 expand_dim + broadcast_add 这一种形式的转换。最后是 layout 变换和常量传播：ifcfg.pass_enabled(“AlterOpLayout”):ifisinstance(target, _target.Target):func=ir_pass.infer_type(func)withtarget:func=ir_pass.alter_op_layout(func)elifisinstance(target,dict):warnings.warn(“AlterOpLayout pass is not enabled for heterogeneous”” execution yet.”)
ifcfg.pass_enabled(“FoldConstant”):func=ir_pass.fold_constant(func)目前的实现里，layout 变换有些 bug ，不支持异构编译。### 3.3 Part III ifisinstance(target,dict):func, target=_run_device_annotation_passes(func, target,fallback_device)# Fuse ops before running code genfunc=ir_pass.infer_type(func)func=ir_pass.fuse_ops(func, cfg.opt_level)接下来首先为异构编译执行一个指定 pass ，根据异构 target 改写计算图，在需要的位置插入数据搬移节点(Device Copy Node)。之后便是进行图融合优化。其优化内容几乎与 NNVM 一样，都是基于算子的 pattern (kElemWise, kBroadcast,kInjective, kCommReduce, kOutEWiseFusable, kOpaque)和融合规则 rule (kUknown, kFuseToMaster, kRealize)来运行融合算法的，可以参考上一篇关于NNVM的文章，这里不再赘述。优化后的 IR 是：fn (%data: Tensor[(1,3,224,224), float32])->Tensor[(1,32,112,112), float32] {%0=meta[relay.Constant][0]# ty=Tensor[(32, 3, 3, 3), float32]%1=meta[relay.Constant][1]# ty=Tensor[(32, 1, 1), float32]%2=fn(%p0: Tensor[(1,3,224,224), float32],%p1: Tensor[(32,3,3,3), float32],%p2: Tensor[(32,1,1), float32])->Tensor[(1,32,112,112), float32] {%3=nn.conv2d(%p0,%p1, strides=[2,2], padding=[1,1], channels=32, kernel_size=[3,3])%4=add(%3,%p2)%5=nn.relu(%4)%5}%6=%2(%data,%0,%1)%6}### 3.4 Part IV func=ir_pass.infer_type(func)graph_gen=_graph_gen.GraphRuntimeCodegen(mod=None, target=target)graph_json, lowered_funcs, params=graph_gen.codegen(func)mod=_tvm_build_module(lowered_funcs, target=target, target_host=target_host)最后是生成代码并返回。生成代码过程中，对于 kernel 的生成，会将所有 Call Node 对应的 Relay IR 先转换为 TVM IR 然后找到 TOPI 中注册的 Compute 与 Schedule，利用 TVM 生成代码。## 4. 总结 Relay 作为 NNVM 的进阶版本，同时具有编程语言的特点和深度学习图构造的能力，借助 TVM 代码生成工具以及 TOPI 中丰富的算子库，可以完成一系列深度学习编译部署工作。对于 Relay 的掌握应该从两个方面入手，本文的内容只涉及第二点。仅从第二点上分析，Relay 和 NNVM 的优化思路别无二致，多出来的优化有两点：并行分支卷积的合并以及算子规范化。此外，Relay 还多考虑了异构场景下的代码生成，这一点在 NNVM 中是没有的。当然，在实现这些算法时，Relay 的技术和 NNVM 差别很大，Relay 将优化过程看作对 IR 的变换，使用 ExprMutator 完成一系列 IR 变换，这一点和 TVM 很像。而 NNVM 将计算图独立地看作了一个结构，在其上运行优化算法，并没有引入编程语言的特点。因此 Relay 的实现对于程序语言领域的研究者来说更加优美和易懂。