将 AMD GPU 引入 TVM 堆栈和 NNVM 编译器（使用 ROCm）


Aditya Atluri，美国超微半导体公司

Masahiro Masuda，Ziosoft 公司

我们很高兴地宣布 TVM 堆栈的一个新的 GPU 后端 - 用于 AMD GPU 的 ROCm 后端。如果您不熟悉 TVM，可以先参考之前的公告。简而言之，TVM 堆栈是一个端到端的编译堆栈，用于将深度学习工作负载部署到所有硬件后端。今天的公告重点介绍对 AMD GPU 的代码生成器支持。具体来说，我们为 AMD GPU 开发了一个新的 GPU 代码生成器。它将以 TVM 前端编写的高级计算描述编译为优化的原生 GPU 代码。它通过使用 TVM 中现有的 LLVM 代码生成器和 LLVM 的 AMDGPU 后端来实现这一点。我们还为 TVM 构建了一个 ROCm 运行时，以支持编译后的 ROCm 模块的原生部署。由于 TVM 堆栈对 NNVM 编译器的支持，我们现在可以直接编译来自深度学习框架的描述，并将它们编译为在 AMD GPU 上运行的裸机代码。

image

TVM 堆栈是由一个开源社区在 Apache-2.0 许可下开发的。ROCm 后端支持是在社区的帮助下完成的。Aditya 首先实现了代码生成和运行时。后来 Masahiro 也加入了进来。Masahiro 的全职工作与 TVM 或 AMD GPU 无关。尽管如此，TVM 让他感到兴奋，并且他一直参与修复错误、解决所有失败的单元测试以及向代码生成添加数学函数支持。

ROCm 堆栈

Radeon Open Compute 是 AMD 的一项开源计划，旨在利用当前和未来几代 GPU 的计算能力。ROCm 软件堆栈是表达和运行最常用的 GPU 编程模型并实现峰值性能的绝佳工具。ROCm 不仅是一个开源堆栈，而且是一个开放堆栈，这意味着所有的 ISA 和硬件特性都有完善的文档记录，并且可以由开发人员编程。开发人员可以尝试不同的编程模型，并尝试多种方法来实现其算法的峰值吞吐量和带宽。

TVM 通过使用 LLVM AMDGPU 后端代码生成器来利用 ROCm 堆栈的开源特性。TVM 从其中间表示 (IR) 转换为 LLVM 中间表示。这就是 ROCm 堆栈开源特性发挥作用的地方。TVM 的 LLVM AMDGPU CodeGen 过程将 LLVM IR 转换为 GPU 汇编代码和目标代码，然后调用这些代码来运行整个网络、层组或单层。

在 ROCm 堆栈上，没有虚拟 ISA，您得到的是您所要求的，不多也不少。因此，可以在内核中以单条指令的粒度调度操作，而无需担心指令重排序和您未要求的其他优化。

将 NNVM 编译器与 ROCm 后端结合使用

借助 TVM 堆栈，我们今天可以直接使用 NNVM 编译器将来自流行的深度学习框架（如 MXNet 和 PyTorch）的模型编译为 AMD GPU 汇编代码。使用 ROCm 后端，通用工作流程如下。

image

我们整理了一些使用 NNVM 编译来自 MXNet 和 PyTorch 的模型并在 AMD GPU 上使用 ROCm 后端运行它们的工作示例。通过 NNVM 编译器堆栈支持更多框架。该存储库可在此处获得。

脚本 mxnet_imagenet_inference.py 演示了在使用最近引入的 MXNet-Gluon 模型在 AMD GPU 上进行 Imagenet 推理。它执行以下操作

  • 从 Gluon 模型库加载 Resnet 50 模型
  • 使用 nnvm.frontend.from_mxnet (...) 将 Gluon Resnet 50 模型转换为 NNVM 图格式
  • 使用 ROCm 后端编译和执行图

该示例附带以下猫的图像。

image

运行我们的网络，它将此图像预测为 “虎斑猫 (tigar cat)”，在 1000 个类别中。

$ python mxnet_imagenet_inference.py
Testing model resnet50_v1
x (1, 3, 224, 224)
TVM prediction top-1: 282 tiger cat

脚本 advanced_superres_onnx.py 给出了加载使用 PyTorch 训练的模型的示例。该模型以 ONNX 格式存储。在本例中，我们的网络以低分辨率图像作为输入，并输出 4 倍高分辨率图像。我们将问题设置和网络架构的详细信息参考 原始论文。该网络有 37 个卷积层，因此它比 NNVM 教程中简单的 4 层网络复杂得多。使用最新 Pytorch 包中的 ONNX 导出接口，我们导出了一个训练好的模型（可在 此处获得）为 ONNX 格式，以便在本例中使用。我们感谢该存储库的作者公开了他的代码和训练好的模型。

为了在 NNVM 中使用 ONNX 格式的模型，我们首先使用 ONNX 库将 ONNX 模型加载到 Protocol buffer 对象中。然后，我们可以使用 nnvm.frontend.from_onnx(...) 来获得等效的 NNVM 图。有了 NNVM 图，我们可以按照上面概述的通用编译和图执行工作流程进行操作。

image

网络的输入是左侧的 64 x 64 图像，输出是右侧的 256 x 256 图像。中间是通过使用双三次插值简单地调整输入图像大小而获得的 256 x 256 图像。网络输出的图像质量要好得多。

输入图像取自原始论文，可在此处获得。

关于性能的说明

当前 ROCm 的支持侧重于功能覆盖范围。通过简单地采用现有的用于 CUDA 后端的 TVM 调度，我们已经看到了有希望的性能结果。例如，您可以尝试运行 TVM 存储库中的 gemm 测试脚本 并查看结果。对于我们测试的两种类型的卡，当前用于方阵乘法的 gemm 配方（尚未针对 AMD GPU 进行专门优化）已经实现了 60% 到 65% 的峰值性能。这是一个很有希望的开始，因为优化性能以达到峰值非常困难，而且我们尚未应用 AMD GPU 特定的优化。我们正在开始研究性能优化，并且我们期望会有更多的改进。

ROCm 后端演练

在本文的以下部分，我们将重点解释在直接使用 TVM 时如何使用 ROCm 后端。您只需在 “rocm” 目标下构建您的 TVM 函数并为其创建一个运行时上下文。在这里，我们展示一个 ROCm 后端用法的示例，遵循 TVM 入门教程中的 “向量加法示例”。

我们首先为向量加法内核设置计算操作和调度。此步骤与后端无关。

from __future__ import absolute_import, print_function
import tvm
import numpy as np

n = tvm.var("n")
A = tvm.placeholder((n,), name='A')
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")
s = tvm.create_schedule(C.op)
bx, tx = s[C].split(C.op.axis[0], factor=64)
s[C].bind(bx, tvm.thread_axis("blockIdx.x"))
s[C].bind(tx, tvm.thread_axis("threadIdx.x"))

接下来，要使用 ROCm 后端，我们在 “rocm” 目标下构建我们的内核。这将导致 TVM 使用我们的新代码生成器。我们还需要一个 ROCm 后端的运行时上下文。

target = "rocm"
fadd_rocm = tvm.build(s, [A, B, C], target, target_host="llvm", name="myadd")
ctx = tvm.rocm(0)

在构建内核并设置运行时上下文后，我们可以启动我们的向量加法内核。

n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)

fadd_rocm(a, b, c)
np.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())

我们可以通过以下方式查看 TVM 生成的 LLVM IR

dev_module = fadd_rocm.imported_modules[0]
print(dev_module.get_source("llvm"))

您应该看到类似这样的内容

; ModuleID = 'myadd__kernel0'
source_filename = "myadd__kernel0"
target datalayout = "e-p:32:32-p1:64:64-p2:64:64-p3:32:32-p4:64:64-p5:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64"
target triple = "amdgcn-amd-amdhsa-hcc"


; Function Attrs: nounwind
define dllexport amdgpu_kernel void @myadd__kernel0(float addrspace(1)* noalias nocapture, float addrspace(1)* noalias nocapture readonly, float addrspace(1)* noalias nocapture readonly, i32) local_unnamed_addr #0 {
entry:
  %4 = tail call i32 @llvm.amdgcn.workgroup.id.x()
  %5 = tail call i32 @llvm.amdgcn.workitem.id.x()
  %6 = add nsw i32 %3, -127
  %7 = ashr i32 %6, 6
  %8 = icmp slt i32 %4, %7
  br i1 %8, label %if_then, label %if_else


if_then:                                          ; preds = %entry
  %9 = shl nsw i32 %4, 6
  br label %if_end.sink.split


if_end.sink.split:                                ; preds = %if_else, %if_then
  %.pre-phi = phi i32 [ %21, %if_else ], [ %9, %if_then ]
  %10 = add nsw i32 %.pre-phi, %5
  %11 = add nsw i32 %.pre-phi, %5
  %12 = sext i32 %11 to i64
  %13 = getelementptr inbounds float, float addrspace(1)* %2, i64 %12
  %14 = load float, float addrspace(1)* %13, align 4, !tbaa !2
  %15 = getelementptr inbounds float, float addrspace(1)* %1, i64 %12
  %16 = load float, float addrspace(1)* %15, align 4, !tbaa !6
  %17 = fadd float %14, %16
  %18 = sext i32 %10 to i64
  %19 = getelementptr inbounds float, float addrspace(1)* %0, i64 %18
  store float %17, float addrspace(1)* %19, align 4, !tbaa !9
  br label %if_end


if_end:                                           ; preds = %if_end.sink.split, %if_else
  ret void


if_else:                                          ; preds = %entry
  %20 = sub nsw i32 %3, %5
  %21 = shl nsw i32 %4, 6
  %22 = icmp slt i32 %21, %20
  br i1 %22, label %if_end.sink.split, label %if_end, !prof !12
}

我们还可以查看 ROCm 后端生成的 GPU 汇编代码。这是在您的 GPU 上运行的真实代码。

print(dev_module.get_source("asm"))

汇编代码应如下所示（省略不必要的细节）

        s_load_dword s1, s[4:5], 0x18
        v_mov_b32_e32 v2, -1
        v_mov_b32_e32 v1, 0
        s_waitcnt lgkmcnt(0)
        s_add_i32 s0, s1, 0xffffff81
        s_ashr_i32 s0, s0, 6
        s_cmp_ge_i32 s6, s0
        s_cbranch_scc0 BB0_2
        v_sub_i32_e32 v1, vcc, s1, v0
        s_lshl_b32 s0, s6, 6
        v_cmp_lt_i32_e32 vcc, s0, v1
        v_mov_b32_e32 v2, 0
        v_cndmask_b32_e64 v1, 0, -1, vcc
BB0_2:
        v_cmp_ne_u32_e32 vcc, 0, v2
        v_cndmask_b32_e64 v2, 0, 1, vcc
        v_cmp_ne_u32_e32 vcc, 1, v2
        s_and_b64 vcc, exec, vcc
        s_cbranch_vccnz BB0_4
        s_lshl_b32 s0, s6, 6
        v_mov_b32_e32 v1, -1
BB0_4:
        v_cmp_ne_u32_e32 vcc, 0, v1
        v_mov_b32_e32 v1, s0
        s_and_saveexec_b64 s[0:1], vcc
        s_xor_b64 s[0:1], exec, s[0:1]
        s_cbranch_execz BB0_6
BB0_5:
        s_load_dwordx2 s[2:3], s[4:5], 0x0
        s_load_dwordx2 s[6:7], s[4:5], 0x8
        v_add_i32_e32 v0, vcc, v1, v0
        s_load_dwordx2 s[4:5], s[4:5], 0x10
        v_ashrrev_i32_e32 v1, 31, v0
        v_lshlrev_b64 v[0:1], 2, v[0:1]
        s_waitcnt lgkmcnt(0)
        v_add_i32_e32 v2, vcc, s4, v0
        v_mov_b32_e32 v3, s5
        v_addc_u32_e32 v3, vcc, v3, v1, vcc
        flat_load_dword v2, v[2:3]
        v_add_i32_e32 v4, vcc, s6, v0
        v_mov_b32_e32 v3, s7
        v_addc_u32_e32 v5, vcc, v3, v1, vcc
        flat_load_dword v4, v[4:5]
        v_mov_b32_e32 v3, s3
        v_add_i32_e32 v0, vcc, s2, v0
        v_addc_u32_e32 v1, vcc, v3, v1, vcc
        s_waitcnt vmcnt(0) lgkmcnt(0)
        v_add_f32_e32 v2, v2, v4
        flat_store_dword v[0:1], v2
BB0_6:
        s_or_b64 exec, exec, s[0:1]
        s_endpgm

链接