作者：Masahiro Masuda，Truman Tian

本文介绍如何将 cuDNN 或 cuBLAS 等外部库与 Relay 一起使用。

Relay 内部用 TVM 来生成 target-specific 的代码。例如，TVM 使用 CUDA 后端为用户提供的网络中的所有层生成 CUDA 内核。有时也可将各个供应商开发的外部库合并到 Relay 中，TVM 有一种机制可以透明地调用这些库——对于 Relay 用户，只需要设置一个适当的 target 字符串。

使用 Relay 的外部库前，用你要用的库构建 TVM。例如，要用 cuDNN，需启用 cmake/config.cmake 中的 USE_CUDNN 选项，必要时要指定 cuDNN 头文件和库目录。

首先导入 Relay 和 TVM。

创建一个简单网络

下面创建一个简单网络进行演示，它由 convolution，batch normalization 和 ReLU activation 组成。

使用 CUDA 后端构建和运行

正常使用 CUDA 后端构建和运行这个网络。设置日志记录级别为 DEBUG，Relay 计算图编译的结果将作为伪代码转储。

输出结果：

生成的伪代码应如下。注意 bias add，batch normalization 和 ReLU activation 是如何融合到卷积核中的。 TVM 从这个表示中生成一个单一的融合内核。

将 cuDNN 用于卷积层

将选项 "-libs=cudnn" 附加到 target 字符串，从而用 cuDNN 将卷积核替换为 cuDNN。

输出结果：

注意，若用 cuDNN，Relay 无法将卷积与其后面的层融合。因为层融合发生在 TVM internal representation（IR）级别。 Relay 将外部库视为黑盒，因此无法将它们与 TVM IR 融合。

下面的伪代码显示了 cuDNN 卷积 + bias add + batch norm + ReLU 变成了两个计算阶段，一个用于 cuDNN 调用，另一个用于其余操作。

验证结果

检查两次运行的结果是否匹配。

结论

本教程介绍了 cuDNN 与 Relay 的使用，此外还支持 cuBLAS。若启用了 cuBLAS，它将在全连接层（relay.dense）内使用。若要用 cuBLAS，请将 target 字符串设置为 "cuda -libs=cublas"。也可以将 cuDNN 和 cuBLAS 与 "cuda -libs=cudnn,cublas" 一起使用。

对于 ROCm 后端，支持 MIOpen 和 rocBLAS。将 target 设置为 "rocm -libs=miopen,rocblas" 以启用它们。

使用外部库的注意事项：

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

其次，外部库限制了计算图编译期间算子融合的可能性，如上所示。TVM 和 Relay 旨在通过联合算子级别和计算图级别优化，在各种硬件上实现最佳性能。为了实现这个目标，应该继续为 TVM 和 Relay 开发更好的优化，同时在必要时使用外部库回退到现有实现。

下载 Python 源代码：using_external_lib.py

下载 Jupyter Notebook：using_external_lib.ipynb