TVM Compiler中文教程：TVM.Relay使用外部库

梦里梦外; 2022-01-23 12:59 563阅读 0赞

### 文章目录 ###

*  TVM.Relay使用外部库
 *   *  创建一个简单的网络
     *  使用cuda后端构建和运行
     *  卷积层使用cuDNN实现
     *  验证结果
     *  结论

# 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开发更好的优化操作，同时在必要时（优化操作没有外部库性能好的情况下）使用外部库也是一种好方法。