TVM, Relay, 深度学习编译器

TVM/Relay graph compiler

Posted by Sz Zheng on 2019-11-10

Relay 探究

1. 前言

Relay 是 TVM 中用来替代 NNVM 的模块，其本身被认为是 NNVM 第二代。在设计上，Relay 被认为相对 NNVM 有以下优势：

  1. 有文本形式中间表示，便于开发和 debug
  2. 支持子图函数、联合模块，便于联合优化
  3. 前端用户友好

其介绍信息可以在这里找到。相比于最初的 NNVM ，Relay 融合进了编程语言领域的知识，带来了许多新的特色（如 let 表达式，支持递归等等），这也使得 Relay 既具有一个编程语言的特色，也具有常规深度学习网络构建的能力。关于其语言特色的介绍，将在后续文章中展现。本文的重点放在 Relay 如何构建神经网络计算图并在计算图上做图优化。

2. 使用

在构建神经网络的时候，Relay使用起来与 NNVM 非常相似，比如下面的例子就是一个Conv2d + BatchNorm + ReLU 的简单网络：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
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

为了运行上述网络，可以写下述代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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 的样子：

1
2
3
4
5
6
7
8
9
10
11
12
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：

1
2
3
4
5
6
7
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
  if isinstance(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

下面进行目标无关优化：

1
func = optimize(func, target, params)

跟踪进 optimize 函数，首先进行的优化是:

1
2
3
if cfg.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 为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
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
}

其次进行的优化是去除公共子表达式：

1
2
3
4
5
6
7
8
9
if cfg.pass_enabled("EliminateCommonSubexpr"):
  def fskip(expr):
      if isinstance(expr, _expr.Call) and expr.op.name == 'cast' and \
         expr.attrs.dtype == 'int32':
          return True
      return False

  func = ir_pass.infer_type(func)
  func = ir_pass.eliminate_common_subexpr(func, fskip)

然后是分支卷积优化：

1
2
3
if cfg.pass_enabled("CombineParallelConv2D"):
  func = ir_pass.infer_type(func)
  func = ir_pass.combine_parallel_conv2d(func)

这部分优化会将具有相同输入的卷积合并成一个大的卷积运算。

接着是常量传播优化，与 NNVM 相似：

1
2
3
4
5
6
7
8
9
if cfg.pass_enabled("FoldConstant"):
  func = ir_pass.fold_constant(func)

if cfg.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 是：

1
2
3
4
5
6
7
8
9
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
}

下面是进行规范化，想法在于将一些特殊运算转换成等价的常规算子运算，便于后续的分析优化：

1
2
3
if cfg.pass_enabled("CanonicalizeOps"):
  func = ir_pass.infer_type(func)
  func = ir_pass.canonicalize_ops(func)

其实 Relay 也只实现了将 bias_add 转换为 expand_dim + broadcast_add 这一种形式的转换。

最后是 layout 变换和常量传播：

1
2
3
4
5
6
7
8
9
10
11
if cfg.pass_enabled("AlterOpLayout"):
  if isinstance(target, _target.Target):
      func = ir_pass.infer_type(func)
      with target:
          func = ir_pass.alter_op_layout(func)
  elif isinstance(target, dict):
      warnings.warn("AlterOpLayout pass is not enabled for heterogeneous"
                    " execution yet.")

if cfg.pass_enabled("FoldConstant"):
  func = ir_pass.fold_constant(func)

目前的实现里，layout 变换有些 bug ，不支持异构编译。

3.3 Part III

1
2
3
4
5
6
if isinstance(target, dict):
   func, target = _run_device_annotation_passes(func, target,
                                                fallback_device)
# Fuse ops before running code gen
func = 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 是：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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

1
2
3
4
5
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 的实现对于程序语言领域的研究者来说更加优美和易懂。

FEATURED TAGS
TVM, Relay, 深度学习编译器
FRIENDS