TVM.Relay使用外部库

这篇教程介绍怎么在Relay中使用cuDNN、cuBlas这样的外部库。

Relay在内部使用TVM生成特定目标的代码。例如，使用cuda作为后端，TVM为用户提供的网络生成所有层的cuda kernel代码。但是有时，将设备厂商提供的外部库合并到Relay中也很有用。幸运地，TVM提供一种透明地调用这些库的机制。对于Relay用户，我们需要做的只是设置适当的目标字符串。

在我们使用Relay的外部库之前，TVM需要使用我们想要使用的库来构建。例如，要使用cuDNN，需要启用cmake / config.cmake中的USE_CUDNN选项，并且必要时需要指定cuDNN include和library目录。

首先，导入Relay和TVM

import tvm
import numpy as np
from tvm.contrib import graph_runtime as runtime
from tvm import relay
from tvm.relay import testing

创建一个简单的网络

构建一个包含卷积、批归一化、ReLU激活的简单网络

out_channels = 16
batch_size = 1
#变量定义
data = relay.var("data",relay.TensorType((batch_size,3,224,224),"float32"))
weight = relay.var("weight")
bn_gamma = relay.var("bn_gamma")
bn_beta = relay.var("bn_beta")
bn_mmean = relay.var("bn_mean")
bn_mvar = relay.var("bn_var")
simple_net = relay.nn.conv2d(data=data, weight=weight, kernel_size=(3,3),channels=out_channels,padding=(1,1))
simple_net = relay.nn.batch_norm(simple_net, bn_gamma, bn_beta, bn_mmean, bn_mvar)[0]
simple_net = relay.nn.relu(simple_net)
simple_net = relay.Function(relay.ir_pass.free_vars(simple_net),simple_net)
data_shape = (batch_size, 3, 224, 224)
net, params = testing.create_workload(simple_net)

使用cuda后端构建和运行

像往常一样，我们使用cuda后端构建和运行网络。通过将日志记录级别设置为DEBUG，Relay graph编译的结果将作为伪代码转储。

import logging
logging.basicConfig(level=logging.DEBUG)
#构建
target = "cuda"
graph, lib, params = relay.build_module.build(
    net, target, params=params)
ctx = tvm.context(target, 0)
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
#创建运行时
module = runtime.create(graph, lib, ctx)
#设置权重和输入数据
module.set_input(**params)
module.set_input("data",data)
#运行推理
module.run()
out_shape = (batch_size, out_channels, 224,224)
out = module.get_output(0,tvm.nd.empty(out_shape))
out_cuda = out.asnumpy()

生成的伪代码应如下所示。请注意bias add，batch_norm和ReLU激活如何融合到卷积内核中。 TVM从这个表示来生成单个融合内核。

produce tensor{
....
}

卷积层使用cuDNN实现

net, params = testing.create_workload(simple_net)
#只需要修改target，其他代码与上面一样
target = "cuda -libs=cudnn" # use cudnn for convolution
graph, lib, params = relay.build_module.build(
        net, target, params=params)
ctx = tvm.context(target, 0)
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
module = runtime.create(graph, lib, ctx)
module.set_input(**params)
module.set_input("data", data)
module.run()
out_shape = (batch_size, out_channels, 224, 224)
out = module.get_output(0, tvm.nd.empty(out_shape))
out_cudnn = out.asnumpy()

注意，如果使用cuDNN，Relay不能将卷积与后面层融合。这是因为层融合发生在TVM IR（中间表示）的层面上。Relay对待外部库是当一个黑匣子，所以没有办法使用TVM IR去融合它们。

下面的伪代码表明cuDNN卷积+Bias Add+batch_norm+ ReLU变成了两个计算阶段，一个用于cuDNN调用，另一个用于其余操作。

// attr [y] storage_scope = "global"
allocate y[float32 * 802816]
produce y { 
  // attr [0] extern_scope = 0
  tvm_call_packed("tvm.contrib.cudnn.conv2d.forward", 1, 0, 1, 1, 1, 1, 1, 1, 1, tvm_stack_make_array(placeholder, tvm_stack_make_shape(1, 3, 224, 224), 0, 4, 0.000000f, 0), tvm_stack_make_array(placeholder, tvm_stack_make_shape(16, 3, 3, 3), 0, 4, 0.000000f, 0), tvm_stack_make_array(y, tvm_stack_make_shape(1, 16, 224, 224), 0, 4, 0.000000f, 0))
}
produce tensor { 
  // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 256
  // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 512
  for (ax0.ax1.fused.ax2.fused.ax3.fused.outer, 0, 7) { 
    if (likely(((blockIdx.x*512) < ((802816 - (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072)) - threadIdx.x)))) { 
      tensor[(((((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/802816)*802816) + (((((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/224) % 224)*224) + ((((blockIdx.x*64) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*32)) % 224))) + ((((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/50176) % 16)*50176))] = max(((y[(((((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/802816)*802816) + (((((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/224) % 224)*224) + ((((blockIdx.x*64) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*32)) % 224))) + ((((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/50176) % 16)*50176))]*placeholder[(((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/50176) % 16)]) + placeholder[(((((blockIdx.x*512) + threadIdx.x) + (ax0.ax1.fused.ax2.fused.ax3.fused.outer*131072))/50176) % 16)]), 0.000000f)
    }
  }
}

验证结果

tvm.testing.assert_allclose(out_cuda, out_cudnn, rtol=1e-5)

结论

教程介绍了cuDNN在Relay的使用。我们也支持cuBLAS。如果启用了cuBLAS，它将在全连接层（relay.dense）中使用。要使用cuBLAS，请将目标字符串设置为“cuda -libs = cublas”。您可以将cuDNN和cuBLAS用于“cuda -libs = cudnn，cublas”。

对于ROCm后端，我们支持MIOpen和rocBLAS。可以使用目标“rocm -libs = miopen，rocblas”启用它们。

能够使用外部库是很好的，但我们需要记住一些注意事项。

首先，使用外部库可能会限制你对TVM和Relay的使用。例如，MIOpen目前仅支持NCHW布局和fp32数据类型，因此您无法在TVM中使用其他布局或数据类型。

其次，更重要的是，外部库限制了图形编译期间操作符融合的可能性，如上所示。TVM和Relay旨在通过联合操作员级别和图形级别优化，在各种硬件上实现最佳性能。为了实现这一目标，我们应该继续为TVM和Relay开发更好的优化操作，同时在必要时（优化操作没有外部库性能好的情况下）使用外部库也是一种好方法。