【TVM 教程】在 Relay 中使用外部库
💡
原文中文,约49200字,阅读约需117分钟。
📝
内容提要
本文介绍了如何将cuDNN或cuBLAS等外部库与Relay一起使用。首先导入Relay和TVM,然后创建一个简单的网络进行演示。接下来使用CUDA后端构建和运行网络。最后生成了一个融合内核的伪代码表示。
🎯
关键要点
- 本文介绍如何将 cuDNN 或 cuBLAS 等外部库与 Relay 一起使用。
- Relay 内部用 TVM 来生成特定目标的代码,支持透明调用外部库。
- 使用外部库前,需要在 TVM 中构建相应的库。
- 示例中创建了一个包含卷积、批归一化和 ReLU 激活的简单网络。
- 使用 CUDA 后端构建和运行网络,并设置日志记录级别为 DEBUG。
- 生成的伪代码展示了卷积、批归一化和 ReLU 激活的融合过程。
- 使用 cuDNN 替换卷积核时,Relay 无法将卷积与后续层融合。
- 伪代码显示了 cuDNN 卷积与后续操作分为两个计算阶段。
- 支持 cuBLAS 的全连接层,使用时需设置相应的 target 字符串。
- 对于 ROCm 后端,支持 MIOpen 和 rocBLAS。
- 使用外部库可能限制 TVM 和 Relay 的使用,影响算子融合的可能性。
- 应继续为 TVM 和 Relay 开发更好的优化,同时在必要时使用外部库回退。
🏷️
标签
➡️