[OSDI'18] TVM: An Automated End-to-End Optimizing Compiler for Deep Learning 阅读笔记

发表于 2026-02-01 分类于 mlsys 阅读次数：本文字数： 10k 阅读时长 ≈ 38 分钟

#0. 摘要

将机器学习应用于广泛多样的硬件设备的需求日益增长. 当前框架依赖于特定供应商的算子库, 并针对狭窄范围的服务器级 GPU 进行优化. 将工作负载部署到新平台–例如手机、嵌入式设备和加速器 (如 FPGA、ASIC) --需要大量手动工作. 我们提出了 TVM, 一种编译器, 它通过暴露图级别和算子级别的优化, 为深度学习工作负载在不同硬件后端提供性能可移植性. TVM 解决了深度学习特有的优化挑战, 例如高级算子融合、映射到任意硬件原语以及内存延迟隐藏. 它还通过采用一种新颖的基于学习的成本建模方法来自动化低级程序的硬件特性优化, 以快速探索代码优化. 实验结果表明, TVM 在低功耗 CPU、手机 GPU 和服务器级 GPU 等硬件后端上提供的性能, 与最先进的、手工调优的库具有竞争力. 我们也展示了 TVM 针对新型加速器后端的能力, 例如基于 FPGA 的通用深度学习加速器. 该系统已开源, 并在多家大型公司内部署使用.

#1. 引言

深度学习 (DL) 模型现在可以识别图像、处理自然语言, 并在具有挑战性的策略游戏中击败人类. 目前, 存在将智能应用程序部署到广泛设备的需求, 这些设备范围从云服务器到自动驾驶汽车和嵌入式设备. 将 DL 工作负载映射到这些设备因硬件特性的多样性而变得复杂, 包括嵌入式 CPU、GPU、FPGA 和 ASIC (例如 TPU). 这些硬件目标在内存组织、计算功能单元等方面存在差异, 如图 1 所示.

divergence

图 1: CPU、GPU 和类似 TPU 的加速器需要不同的片上内存架构 (on-chip memory architectures) 和计算原语 (compute primitives). 在生成优化代码时必须解决这种差异.

当前的深度学习框架, 如 TensorFlow、MXNet、Caffe 和 PyTorch, 依赖于计算图中间表示来实现优化, 例如自动微分和动态内存管理. 然而, 图级别的优化通常过于抽象, 难以处理特定于硬件后端的算子级别转换. 这些框架大多专注于狭隘的服务器级 GPU 设备, 并将特定目标的优化委托给高度工程化且供应商特定的算子库. 这些算子级别库需要大量的手动调优, 因此过于专业化和不透明, 难以在不同硬件设备间轻松移植. 目前, 为各种深度学习框架提供对多样化硬件后端的支持需要大量的工程工作. 即使对于已支持的后端, 框架也必须在以下两者之间做出艰难的选择: (1) 避免会产生预定义算子库中不存在的算子的图优化, 以及 (2) 使用这些算子, 但是是未优化过的实现.

为了实现对不同硬件后端的图级和算子级优化, 我们采用了一种根本不同的端到端方法. 我们构建了 TVM, 这是一个编译器, 它从现有框架中获取深度学习程序的高级规范, 并为多种硬件后端生成低级优化代码. 为了吸引用户, TVM 需要提供与跨不同硬件后端的手动优化算子库相媲美的性能. 这一目标需要解决下面描述的关键挑战.

利用特定硬件特性和抽象 深度学习加速器引入了优化的张量计算原语, 而 GPU 和 CPU 则持续改进其处理单元. 这给针对给定算子描述生成优化代码带来了重大挑战. 硬件指令的输入是多维的, 具有固定或可变长度; 它们决定了不同的数据布局; 并且对内存层次结构有特殊要求. 系统必须有效利用这些复杂原语以实现加速. 此外, 加速器设计通常也更倾向于精简控制, 并将大部分调度复杂性卸载到编译器栈中. 对于专用加速器, 系统现在需要生成显式控制流水线依赖关系的代码, 以隐藏内存访问延迟–这是硬件为 CPU 和 GPU 执行的任务.

优化的大搜索空间 另一个挑战是在不手动调整算子的情况下生成高效代码. 内存访问、线程模式和新型硬件原语的综合选择为生成的代码 (例如, 循环 tiles 和排序、缓存、循环展开) 创建了一个巨大的配置空间, 如果我们实现黑盒自动调优, 这将导致巨大的搜索成本. 人们可以采用预定义的成本模型来指导搜索, 但由于现代硬件复杂性的增加, 建立精确的成本模型很困难. 此外, 这种方法将要求我们为每种硬件类型建立单独的成本模型.

TVM 通过三个关键模块来解决这些挑战.

我们引入了一种 张量表达式语言 来构建算子, 并提供程序变换原语, 这些原语能够生成具有各种优化的程序的不同版本. 这一层扩展了 Halide 的计算/调度分离概念, 通过将目标硬件内建函数 (hardware intrinsics) 与变换原语分离, 从而支持新型加速器及其对应的新内建函数. 此外, 我们引入了新的变换原语来解决与 GPU 相关的挑战, 并支持在专用加速器上部署. 然后, 我们可以应用不同序列的程序变换, 为给定的算子声明形成丰富的有效程序空间.
我们引入了一个 自动化的程序优化框架 来寻找优化的张量算子. 优化器在基于机器学习的成本模型的指导下进行工作, 该模型会随着我们从硬件后端收集更多数据而自适应和改进.
在自动代码生成器的基础上, 我们引入了一个 图重写器, 该重写器充分利用了高层 (high-) 和算子层 (operator-level) 的优化.

通过结合这三个模块, TVM 可以从现有的深度学习框架中获取模型描述, 执行高低级联合优化, 并为后端 (例如 CPU、GPU 和基于 FPGA 的专用加速器) 生成硬件特定的优化代码.

本文做出了以下贡献:

我们识别了在为不同硬件后端上的深度学习工作负载提供性能可移植性时面临的主要优化挑战.
我们引入了新的调度原语, 这些原语利用了跨线程内存重用、新的硬件内建函数和延迟隐藏.
我们提出并实现了一个基于机器学习的优化系统, 用于自动探索和搜索优化的张量算子.
我们构建了一个端到端的编译和优化栈, 允许在高级框架 (包括 TensorFlow、MXNet、PyTorch、Keras、CNTK) 中指定的深度学习工作负载部署到不同的硬件后端 (包括 CPU、服务器 GPU、移动 GPU 和基于 FPGA 的加速器). 开源的 TVM 已经在几家大型公司中投入生产使用. 我们在服务器级 GPU、嵌入式 GPU、嵌入式 CPU 和自定义通用基于 FPGA 的加速器上使用真实世界工作负载评估了 TVM. 实验结果表明, TVM 在不同的后端上提供了可移植的性能, 并实现了比现有基于手工优化库的框架快 1.2x 到 3.8x 的速度提升.

#2. 概述

alt text

图 2: TVM 系统概述. 当前栈支持从许多深度学习框架和交换格式 (如 CoreML 和 ONNX)进行描述, 以针对主要的 CPU、GPU 和专用加速器.

本节通过示例介绍 TVM 的各个组成部分. 图 2 总结了 TVM 的执行步骤及其在论文中对应的章节. 系统首先从现有框架中获取模型, 并将其转换为计算图表示. 然后执行高级数据流重写以生成优化后的图. 算子级优化模块必须为该图中的每个融合算子生成高效代码. 算子通过声明式张量表达式语言指定; 执行细节不需要. TVM 为给定硬件目标的算子识别一系列可能的代码优化方案. 可能的优化方案构成一个庞大的空间, 因此我们使用基于机器学习的成本模型来寻找优化后的算子. 最后, 系统将生成的代码打包成可部署的模块.

#终端用户示例

只需几行代码, 用户即可从现有的深度学习框架中获取模型, 并调用 TVM API 来获取可部署的模块:

import tvm as t
# Use keras framework as example, import model
graph, params = t.frontend.from_keras(keras_model)
target = t.target.cuda()
graph, lib, params = t.compiler.build(graph, target, params)

该编译运行时模块包含三个组件: 最终优化的计算图 (graph)、生成的算子 (lib)和模块参数 (params). 这些组件可用于将模型部署到目标后端:

import tvm.runtime as t
module = runtime.create(graph, lib, t.cuda(0))
module.set_input(**params)
module.run(data=data_array)
output = tvm.nd.empty(out_shape, ctx=t.cuda(0))
module.get_output(0, output)

TVM 支持在 C++、Java 和 Python 等语言中的多个部署后端. 本文的其余部分描述了 TVM 的架构以及系统程序员如何扩展它以支持新的后端.

#3. 优化计算图

alt text

图 3: 一个两层卷积神经网络的示例计算图.

图中的每个节点代表一个操作, 该操作消耗一个或多个张量并产生一个或多个张量. 张量操作可以通过属性进行参数化以配置其行为 (例如, 填充 padding 或步长 strides).

计算图是表示深度学习框架中程序的一种常见方式. 图 3 展示了一个两层卷积神经网络的计算图表示示例. 这种高级表示与低级编译器中间表示 (如 LLVM) 的主要区别在于, 中间数据项是大型、多维的张量. 计算图提供了操作符的全局视图, 但它们避免了指定每个操作符必须如何实现. 与 LLVM 中间表示类似, 计算图可以转换为功能等价的图以应用优化. 我们还利用了常见深度学习工作负载中的形状特定性来针对一组固定的输入形状进行优化.

TVM 利用计算图表示来应用高级优化: 节点表示张量或程序输入上的操作, 边表示操作之间的数据依赖关系. 它实现了许多图级优化, 包括: 算子融合 (operator fusion), 将多个小操作融合在一起; 常量折叠 (constant-folding), 预先计算可以静态确定的图部分, 节省执行成本; 静态内存规划阶段 (static memory planning pass), 预先分配内存来存储每个中间张量; 以及数据布局转换 (data layout transformations), 将内部数据布局转换为后端友好的形式. 我们现在讨论算子融合和数据布局转换.

alt text

图 4: 融合操作与非融合操作的性能比较. TVM 生成这两种操作. 在 NVIDIA Titan X 上测试.

#算子融合

算子融合将多个算子组合成一个内核, 而无需将中间结果保存在内存中. 这种优化可以显著减少执行时间, 特别是在 GPU 和专用加速器中. 具体来说, 我们识别了四种图算子类别: (1) 单射 (一对一映射, 例如加法), (2) 归约 (例如求和), (3) 复杂输出可融合 (可以融合元素级映射到输出, 例如 conv2d), 以及 (4) 不透明 (无法融合, 例如排序). 我们提供了通用的融合规则, 如下所示.

多个单射算子可以融合成另一个单射算子.
一个归约算子可以与输入的单射算子融合 (例如, 融合 scale 和 sum). 像 conv2d 这样的算子是复杂输出可融合的, 我们可以将其元素级算子融合到其输出.

我们可以应用这些规则将计算图转换为融合版本. 图 4 展示了这种优化对不同工作负载的影响. 我们发现, 融合算子通过减少内存访问, 最多可产生 1.2x 到 2x 的加速比.

#数据布局转换

对于一个给定的张量, 在计算图中存储的方式有多种. 最常见的数据布局选择是列主序 (column major) 和行主序 (row major). 在实践中, 我们可能更喜欢使用更复杂的数据布局. 例如, 一个深度学习加速器可能会利用4x4矩阵运算, 需要将数据分块为4x4块以优化访问局部性.

数据布局优化将计算图转换为可以在目标硬件上使用更好的内部数据布局执行的形式. 它首先根据内存层次结构所规定的约束, 为每个算子指定首选的数据布局. 然后, 如果生产者和消费者的首选数据布局不匹配, 我们就会在它们之间执行适当的布局转换.

alt text

图 5: 示例调度转换, 优化在专用加速器上的矩阵乘法.

虽然高级图优化可以极大地提高深度学习工作负载的效率, 但其效果仅取决于算子库所提供的内容. 目前, 少数支持算子融合的深度学习框架需要算子库提供融合模式的实现. 随着网络算子的不断引入, 可能的融合核的数量可能会急剧增长. 当面向越来越多的硬件后端时, 所需的融合模式实现数量会随着必须支持的数据布局、数据类型和加速器内联函数数量的组合呈指数增长. 手动为程序所需的各种操作以及每个后端设计算子核是不可行的. 为此, 我们接下来提出一种代码生成方法, 该方法可以为给定模型的算子生成各种可能的实现.

#4. 生成张量操作

TVM 通过为每个算子生成多种有效的硬件后端实现并选择最优实现来生成高效代码. 该过程基于 Halide 将描述与计算规则 (或调度优化) 解耦的思想 , 并将其扩展以支持新的优化 (嵌套并行、张量化和延迟隐藏) 以及广泛的硬件后端. 我们现在重点介绍 TVM 特有的功能.

#4.1 张量表达式与调度空间

alt text

图 6: TVM 调度 lowering 和代码生成过程. 该表列出了现有的 Halide 和新的 TVM 调度原语, 用于优化 CPU、GPU 和加速器后端的调度. 张量化对于加速器至关重要, 但也可以用于 CPU 和 GPU. 特殊的内存作用域能够实现 GPU 中的内存重用, 以及在加速器中对片上内存进行显式管理. 延迟隐藏是针对类似 TPU 的加速器的.

我们介绍一种张量表达式语言, 以支持自动代码生成. 与高级计算图表示不同, 其中张量操作的实现是不透明的, 每个操作都在一个索引公式表达式语言中描述. 以下代码展示了一个计算转置矩阵乘法的张量表达式示例:

m, n, h = t.var('m'), t.var('n'), t.var('h')
A = t.placeholder((m, h), name='A')
B = t.placeholder((n, h), name='B')
k = t.reduce_axis((0, h), name='k')
C = t.compute((m, n), lambda y, x:
t.sum(A[k, y] * B[k, x], axis=k))

每个计算操作都指定了输出张量的形状以及描述如何计算其每个元素的算术表达式. 我们的张量表达式语言支持常见的算术和数学运算, 涵盖了常见的深度学习算子模式. 该语言不指定循环结构和其他许多执行细节, 并为针对不同后端添加硬件感知优化提供了灵活性. 我们采用 Halide 中的计算/调度解耦原则, 使用调度来表示从张量表达式到低级代码的特定映射. 有许多可能的调度可以实现这一功能.

我们通过逐步应用保留程序逻辑等价性的基本转换 (调度原语) 来构建调度. 图 5 展示了一个在专用加速器上调度矩阵乘法的例子. 内部, TVM 使用一种数据结构来跟踪循环结构和其他信息, 随着我们应用调度转换. 这些信息随后可以帮助为给定的最终调度生成低级代码.

我们的张量表达式借鉴了 Halide、Darkroom 和 TACO 的设计. 其主要增强功能包括支持下面讨论的新调度优化. 为了在许多后端上实现高性能, 我们必须支持足够的调度原语, 以涵盖不同硬件后端上的各种优化. 图 6 总结了 TVM 支持的指令代码生成过程和调度原语. 我们从 Halide 重用了有用的原语和低级循环程序抽象语法树, 并引入了新的原语来优化 GPU 和加速器的性能. 这些新原语对于实现最佳 GPU 性能是必要的, 对于加速器也是必不可少的. CPU、GPU、TPU 类加速器是深度学习的三种重要硬件类型. 本节描述了针对 CPU、GPU 和 TPU 类加速器的新优化原语, 而第 5 节则解释了如何自动推导高效的调度.

#4.2 基于合作的嵌套并行

并行性是提高深度学习工作负载中计算密集型内核效率的关键. 现代 GPU 提供了巨大的并行性, 要求我们将并行模式嵌入到调度转换中. 大多数现有解决方案采用一种称为 嵌套并行 (nested parallelism) 的模型, 这是一种形式的 fork-join. 该模型需要一个并行调度原语来并行化数据并行任务; 每个任务可以进一步递归地细分为子任务, 以利用目标架构的多级线程层次结构 (例如 GPU 中的线程组). 我们称此模型为 无共享嵌套并行 (shared-nothing nested parallelism), 因为同一并行计算阶段中的一个工作线程不能查看其兄弟线程的数据.

alt text

图 7: 在矩阵乘法工作负载上, 带有和未带有协同共享内存获取的 TVM 性能比较. 在 NVIDIA Titan X 上测试.

一种 shared-nothing 的替代方法是协同获取数据. 具体而言, 一组线程可以协同获取它们所需的所有数据, 并将其放置到一个共享内存空间中.

这种优化可以利用 GPU 内存层次结构, 并通过共享内存区域实现跨线程的数据重用. TVM 使用调度原语支持这种著名的 GPU 优化, 以实现最佳性能. 以下 GPU 代码示例优化了矩阵乘法.

for thread_group (by, bx) in cross(64, 64):
  for thread_item (ty, tx) in cross(2, 2):
    local CL[8][8] = 0
    shared AS[2][8], BS[2][8]
    for k in range(1024):
      # 所有线程协作式地通过不同的 pattern 加载 AS 和 BS
      for i in range(4):
        AS[ty][i*4+tx] = A[k][by*64+ty*8+i*4+tx]
      for each i in 0..4:
        BS[ty][i*4+tx] = B[k][bx*64+ty*8+i*4+tx]
      # 编译器自动插入的 Barrier
      memory_barrier_among_threads()
      for yi in range(8):
        for xi in range(8):
          CL[yi][xi] += AS[yi] * BS[xi]
      for yi in range(8):
        for xi in range(8):
          C[yo*8+yi][xo*8+xi] = CL[yi][xi]

图 7 展示了这种优化的影响. 我们将内存作用域的概念引入调度空间, 以便可以将计算阶段 (代码中的 AS 和 BS) 标记为共享. 在没有显式内存作用域的情况下, 自动作用域推断会将计算阶段标记为线程局部. 共享任务必须计算组中所有工作线程的依赖关系. 此外, 必须正确插入内存同步屏障, 以确保共享加载的数据对消费者可见. 最后, 除了对 GPU 有用外, 内存作用域还允许我们标记特殊内存缓冲区, 并在针对专用 DL 加速器时创建特殊优化规则.

#4.3 张量化

深度学习工作负载具有高算术强度, 通常可以分解为矩阵-矩阵乘法或 1D 卷积等张量运算. 这些自然分解导致了近期添加张量计算原语的趋势. 这些新原语为基于调度的编译带来了机遇和挑战; 虽然使用它们可以提高性能, 但编译框架必须无缝集成它们. 我们将此过程称为张量化 (tensorization): 它类似于 SIMD 架构的向量化 (vectorization), 但存在显著差异. 指令输入是多维的, 具有固定或可变长度, 并且每个输入具有不同的数据布局. 更重要的是, 我们不能支持一组固定的原语, 因为新的加速器正在涌现, 它们具有各自不同的张量指令变体. 因此, 我们需要一个可扩展的解决方案.

我们通过使用张量内建函数声明机制, 将目标硬件内建函数与调度分离, 使张量化过程可扩展. 我们使用相同的张量表达式语言来声明每个新硬件内建函数的行为及其相关的降低规则. 以下代码展示了如何声明一个 8x8 张量硬件内建函数.

w, x = t.placeholder((8, 8)), t.placeholder((8, 8))
k = t.reduce_axis((0, 8))
y = t.compute((8, 8), lambda i, j: t.sum(w[i, k] * x[j, k], axis=k))

def gemm_intrin_lower(inputs, outputs):
  ww_ptr = inputs[0].access_ptr("r")
  xx_ptr = inputs[1].access_ptr("r")
  zz_ptr = outputs[0].access_ptr("w")
  compute = t.hardware_intrin("gemm8x8", ww_ptr, xx_ptr, zz_ptr)
  reset = t.hardware_intrin("fill_zero", zz_ptr)
  update = t.hardware_intrin("fuse_gemm8x8_add", ww_ptr, xx_ptr, zz_ptr)
  return compute, reset, update

gemm8x8 = t.decl_tensor_intrin(y.op, gemm_intrin_lower)

此外, 我们引入了一种张量化调度原语, 用以将一个计算单元替换为相应的内建函数. 编译器将计算模式与硬件声明进行匹配, 并将其转换为相应的硬件内建函数.

张量化将调度与特定的硬件原语解耦, 使得扩展 TVM 以支持新的硬件架构变得容易. 张量化调度生成的代码与高性能计算中的实践一致: 将复杂操作分解为一系列微内核调用. 我们还可以使用 张量化 (tensorize) 原语来利用手工制作的微内核, 这在某些平台上可能有益. 例如, 我们通过利用位串行矩阵向量乘法微内核, 为在宽度为一或两位的数据类型上运行的移动 CPU 实现了超低精度运算. 该微内核将结果累积到逐渐增大的数据类型中, 以最小化内存占用. 将微内核作为张量内置于 TVM 中, 与非张量化版本相比, 可加速高达 1.5x.

#4.4 显式内存延迟隐藏

延迟隐藏 (Latency hiding) 是指将内存操作与计算重叠的过程, 以最大化内存和计算资源的利用率. 它需要根据目标硬件后端采用不同的策略. 在 CPU 上, 内存延迟隐藏通过同时多线程或硬件预取隐式实现. GPU 依赖于多条线程束的快速上下文切换. 相比之下, TPU 等专用深度学习加速器通常倾向于更精简的控制, 采用 解耦访问-执行 (decoupled access-execute, DAE) 架构, 并将细粒度同步问题卸载到软件.

alt text

图 9: 硬件中的解耦访问-执行隐藏了大部分内存访问延迟, 通过允许内存和计算重叠实现. 执行正确性通过依赖令牌入队/出队形式的低级同步来强制执行, 编译器栈必须在指令流中插入这些操作.

图 9 展示了一个能够减少运行时延迟的 DAE 硬件流水线. 与单体式硬件设计相比, 该流水线能够隐藏大部分内存访问开销, 并几乎完全利用计算资源. 为了实现更高的利用率, 指令流必须增加细粒度的同步操作. 如果没有这些操作, 依赖关系无法得到保证, 从而导致执行错误. 因此, DAE 硬件流水线需要在流水线阶段之间进行细粒度的依赖入队/出队操作, 以确保正确执行, 如图 9 中的指令流所示.

alt text

图 8: TVM 虚拟线程 lowering 将虚拟线程并行程序转换为单指令流; 该流包含硬件可解释的显式 low-level 同步, 以恢复隐藏内存访问延迟所需的流水线并行性.

为 DAE 加速器编程需要显式的低级同步, 这很困难. 为了减轻编程负担, 我们引入了一种虚拟线程调度原语, 允许程序员像指定硬件后端一样指定高级数据并行程序, 并支持多线程. TVM 随后自动将程序降低为具有低级显式同步的单指令流, 如图 8 所示. 该算法从高级多线程程序调度开始, 然后插入必要的低级同步操作, 以保证每个线程内的正确执行. 接下来, 它将所有虚拟线程的操作交错到一个单指令流中. 最后, 硬件从指令流中的低级同步操作中恢复可用的流水线并行性.

#延迟隐藏的硬件评估

alt text

图 10: 基于 FPGA 的深度学习加速器运行 ResNet 推理时的屋顶线. 通过 TVM 启用的延迟隐藏功能, 基准测试的性能更接近屋顶线, 展示了更高的计算和内存带宽效率.

我们现在展示延迟隐藏在基于定制 FPGA 的加速器设计上的有效性, 该设计在 6.4 小节中进行了详细描述. 我们在加速器上运行了 ResNet 的每一层, 并使用 TVM 生成了两个调度方案: 一个带有延迟隐藏, 一个没有. 带有延迟隐藏的调度方案通过虚拟线程并行化程序, 以暴露流水线并行性并因此隐藏内存访问延迟. 结果如图 10 所示, 以屋顶线图的形式呈现; 屋顶线性能图提供了关于给定系统如何利用计算和内存资源为不同基准测试提供见解. 总体而言, 延迟隐藏提高了所有 ResNet 层的性能. 峰值计算利用率从没有延迟隐藏时的 70%增加到带有延迟隐藏时的 88%.

#5. 自动化优化

鉴于丰富的调度原语, 我们剩余的问题是为深度学习模型的每一层找到最优算子实现. 在此, TVM 为与每一层相关的特定输入形状和布局创建专用算子. 这种专用化提供了显著的性能优势 (与针对形状和布局多样性较小的手工代码相比), 但也带来了自动化挑战. 系统需要选择调度优化–例如修改循环顺序或针对内存层次结构进行优化–以及调度特定参数, 如tiles大小和循环展开因子. 此类组合选择为每个硬件后端创造了大量的算子实现搜索空间. 为应对这一挑战, 我们构建了一个自动调度优化器, 包含两个主要组件:

一个提出有前景新配置的调度探索器
一个预测给定配置性能的机器学习成本模型

本节将介绍这些组件以及 TVM 的自动优化流程 (图 11).

#5.1 调度空间规范

我们构建了一个调度模板规范 API, 允许开发者在调度空间中声明调节器. 模板规范在指定可能的调度时, 允许开发者根据需要融入其领域特定知识. 我们还为每个硬件后端创建了一个通用主模板, 该模板能够根据使用张量表达式语言表达的计算描述自动提取可能的调节器. 从高层次来看, 我们希望尽可能多地考虑配置, 并让优化器管理选择负担. 因此, 优化器必须针对我们实验中使用的真实世界深度学习工作负载搜索数十亿种可能的配置.

#5.2 基于机器学习的成本模型

一种从大型配置空间中寻找最佳调度方案的方法是通过黑盒优化, 即自动调优. 该方法用于调优高性能计算库. 然而, 自动调优需要大量实验来识别良好配置.

另一种方法是构建预定义的成本模型来指导特定硬件后端的搜索, 而不是运行所有可能性并测量其性能. 理想情况下, 完美的成本模型考虑了所有影响性能的因素: 内存访问模式、数据重用、流水线依赖关系和线程模式等. 不幸的是, 由于现代硬件复杂性的增加, 这种方法负担沉重. 此外, 每个新的硬件目标都需要一个新的 (预定义的)成本模型.

alt text

图 11: 自动化优化框架概述. 调度探索器使用基于机器学习的成本模型检查调度空间, 并通过 RPC 选择在分布式设备集群上运行的实验. 为了提高其预测能力, 机器学习模型使用数据库中记录的收集数据定期更新.

方法类别	数据成本	模型偏差	需要硬件信息	从历史中学习
黑盒自动调优 Blackbox auto-tuning	高	无	否	否
预定义成本模型 Predefined cost model	无	高	是	不
基于机器学习的成本模型 ML based cost model	低	低	不	是

表 1: 自动化方法的比较. 模型偏差是指由于建模导致的 inaccuracy.

我们采用统计方法来解决这个问题. 在这种方法中, 一个调度探索器提出可能提高算子性能的配置. 对于每个调度配置, 我们使用一个机器学习模型, 该模型以降级后的循环程序为输入, 并预测其在给定硬件后端的运行时间. 该模型使用在探索过程中收集的运行时测量数据训练, 不需要用户输入详细的硬件信息. 在优化过程中, 我们随着探索更多配置而定期更新模型, 这提高了与其他相关工作负载的准确性. 通过这种方式, 随着更多实验试验的增加, 机器学习模型的质量也随之提高. 表 1 总结了自动化方法之间的主要差异. 基于机器学习的成本模型在自动调优和预定义成本建模之间取得了平衡, 并且可以从相关工作负载的历史性能数据中受益.

#机器学习模型设计选择

alt text

图 12: 在 TITAN X 上对 ResNet-18 中 conv2d 算子不同自动化方法的比较. 基于 ML 的模型从无训练数据开始, 并使用收集到的数据来提升自身性能. Y 轴表示相对于 cuDNN 的速度提升. 我们观察到其他工作负载也呈现出类似的趋势.

在选择调度探索器将使用的机器学习模型时, 我们必须考虑两个关键因素: 质量和速度. 调度探索器频繁查询成本模型, 这会导致由于模型预测时间和模型重新拟合时间而产生的开销. 为了有效, 这些开销必须小于在真实硬件上测量性能所需的时间, 该时间可能因具体的工作负载/硬件目标而异, 通常为秒级. 这一速度要求将我们的问题与传统的超参数调整问题区分开来, 在后者中, 测量成本相对于模型开销非常高, 可以使用更昂贵的模型. 除了模型的选择外, 我们还需要选择一个用于训练模型的目标函数, 例如配置预测运行时间的误差. 然而, 由于探索器仅根据预测的相对顺序选择最佳候选 (A 比 B 运行更快), 因此我们无需直接预测绝对执行时间. 相反, 我们使用排序目标来预测运行成本相对顺序.

alt text

图 13: ML 成本模型的示例工作流程. XGBoost 根据循环程序特征预测成本. TreeRNN 直接总结 AST.

我们在我们的机器学习优化器中实现了多种类型的模型. 我们采用了一种梯度树提升模型 (基于 XGBoost), 该模型基于从循环程序中提取的特征进行预测; 这些特征包括每个循环级别中每个内存缓冲区的内存访问次数和重用率, 以及循环注解 (如"向量化"、“循环展开” 和"并行") 的独热编码. 我们还评估了一种使用 TreeRNN 对循环程序的抽象语法树 (AST) 进行总结的神经网络模型, 而无需进行特征工程. 图 13 总结了成本模型的流程图. 我们发现树提升和 TreeRNN 具有相似的预测质量. 然而, 前者预测速度是后者的两倍, 且训练时间成本更低. 因此, 在实验中我们选择了梯度树提升作为默认的成本模型. 尽管如此, 我们相信这两种方法都具有价值, 并期待未来对此问题的更多研究.

平均而言, 树提升模型进行预测需要 0.67 毫秒, 比运行真实测量快数千倍. 图 12 比较了基于机器学习的优化器与黑盒自动调优方法; 前者比后者更快地找到更好的配置.

#5.3 调度探索

一旦我们选择了一个成本模型, 就可以使用它来选择有潜力的配置, 并在这些配置上迭代运行实际测量. 在每次迭代中, 探索器使用机器学习模型的预测来选择一批候选者进行测量. 收集到的数据随后被用作训练数据来更新模型. 如果不存在初始训练数据, 探索器会随机选择候选者进行测量.

最简单的探索算法通过成本模型枚举并运行每种配置, 选择预测表现最佳的配置. 然而, 这种策略在大搜索空间中变得难以处理. 相反, 我们运行一个并行模拟退火算法. 探索器从随机配置开始, 在每一步随机走到附近的一个配置. 如果成本按照成本模型的预测降低, 则该转换成功. 如果目标配置的成本更高, 则很可能失败 (被拒绝). 随机游走倾向于收敛到成本模型预测成本较低的配置. 探索状态在成本模型更新时持续存在; 在这些更新后, 我们从最后一个配置继续.

#5.4 分布式设备池和远程过程调用

分布式设备池扩展了硬件上的试验运行, 并支持多个优化任务之间的细粒度资源共享. TVM 实现了一个基于 RPC 的定制化分布式设备池, 允许客户端在特定类型的设备上运行程序. 我们可以使用此接口在主机编译器上编译程序、请求远程设备、远程运行函数, 并在主机上的同一脚本中访问结果. TVM 的 RPC 支持动态上传, 并能运行使用其运行时约定的交叉编译模块和函数. 因此, 相同的底层架构可以执行单个工作负载优化和端到端图推理. 我们的方法自动化了跨多个设备的编译、运行和分析步骤. 该基础设施对于嵌入式设备尤为重要, 传统上嵌入式设备需要繁琐的手动交叉编译、代码部署和测量工作.

#7. 相关工作

深度学习框架为用户提供便捷的接口, 以表达深度学习工作负载, 并轻松地在不同的硬件后端上部署它们. 尽管现有的框架目前依赖于特定供应商的张量运算库来执行其工作负载, 但它们可以利用 TVM 的栈为更多硬件设备生成优化代码.

高级计算图 DSL 是表示和执行高级优化的典型方式. Tensorflow 的 XLA 和最近引入的 DLVM 都属于这一类别. 这些工作中的计算图表示是相似的, 本文也使用了高级计算图 DSL. 虽然图级表示非常适合高级优化, 但它们对于在多样化的硬件后端下优化张量运算来说过于高级. 先前的工作依赖于特定的降低规则来直接生成低级 LLVM, 或者采用供应商定制的库. 这些方法需要对每个硬件后端和运算变体组合投入大量的工程工作.

alt text

图 21: 我们将 ResNet 工作负载中的卷积卸载到基于 FPGA 的加速器上. 灰色的条形对应那些无法被 FPGA 加速且因此必须在 CPU 上运行的层. 与 Cortex A9 相比, FPGA 为卸载的卷积层提供了 40x 的加速.

Halide 引入了计算与调度分离的思想. 我们借鉴了 Halide 的见解, 并在我们的编译器中重用了其现有的有用调度原语. 我们的张量算子调度也与 GPU 的 DSL 相关工作以及基于多边形的循环变换有关. TACO 引入了一种生成 CPU 上稀疏张量算子的通用方法. Weld 是一种用于数据处理任务的 DSL. 我们特别关注解决 GPU 和专用加速器上 DL 工作负载的新调度挑战. 我们的新原语有可能被这些工作中的优化管道所采用.

高性能库如 ATLAS 和 FFTW 采用自动调优以获得最佳性能. Tensor comprehension 结合了黑盒自动调优和多面体优化来优化 CUDA 内核. OpenTuner 和现有的超参数调优算法采用领域无关的搜索. Halide 使用预定义的成本模型自动调度图像处理管道. TVM 的机器学习模型采用有效的领域感知成本建模, 考虑程序结构. 基于分布式调度的优化器可扩展到更大的搜索空间, 并能在支持的后端范围内找到最先进的内核. 更重要的是, 我们提供了一个端到端的栈, 可以直接从深度学习框架获取描述, 并与图级栈联合优化.

尽管深度学习加速器正逐渐流行, 但如何构建编译栈以有效针对这些设备仍不明确. 我们在评估中使用的 VDLA 设计提供了一种通用的方法来总结类似 TPU 加速器的特性, 并支持对如何为加速器编译代码的具体案例研究. 我们的方法也可能惠及现有的深度学习编译至 FPGA 的系统. 本文通过张量化和编译器驱动的延迟隐藏, 提供了一种通用的解决方案, 以有效针对加速器.

#8. 结论

我们提出了一种端到端的编译栈, 以解决深度学习在不同硬件后端上的基本优化挑战. 我们的系统包括自动端到端优化, 这在历史上是一项劳动密集且高度专业化的任务. 我们希望这项工作能够鼓励对端到端编译方法的研究, 并为深度学习系统软硬件协同设计技术开辟新的机遇.