Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

发表时间: 2018-02 · arXiv:1802.04730

原文: https://arxiv.org/abs/1802.04730

作者/机构: Nicolas Vasilache (Facebook AI Research), Oleksandr Zinenko (Inria & ENS, DI), Theodoros Theodoridis (ETH Zurich), Priya Goyal (Facebook AI Research), Zachary DeVito (Facebook AI Research), William S. Moses (MIT CSAIL), Sven Verdoolaege (Polly Labs & Facebook AI Research), Andrew Adams (Facebook AI Research), Albert Cohen (Inria & ENS, DI & Facebook AI Research)

速读

一句话结论 本文提出了一种名为 Tensor Comprehensions (TC) 的高级语言及端到端多面体编译流程，能够将数学表示的张量计算自动降级并即时编译为高性能 GPU 核函数，在真实生产模型上相比 NVIDIA 官方库实现了高达 4 倍的加速。

要解决什么问题 深度学习框架在执行计算图时高度依赖 CUDNN 等底层高性能库，这导致了两个核心卡点。首先，当研究人员发明新算子时，手动开发高性能自定义算子的工程成本极高，且极易掉速，严重拖慢了算法创新。其次，即使使用现有的高度优化库，由于库函数是预编译的通用实现，它们无法针对用户特定的张量尺寸（如非标准的小批量大小或特殊的特征图维度）进行特化调优，同时也错失了跨算子融合优化的机会。在低延迟场景或非标准网络架构下，频繁的核函数启动和内存同步开销会取代纯计算，成为极其严重的性能瓶颈。

怎么做的 核心思路是引入一种接近纯数学表达的领域特定语言 Tensor Comprehensions (TC)，并配合基于多面体模型（Polyhedral Model）的即时编译（JIT）流程，将高层算子自动融合并特化为单个底层 GPU 核函数。这种机制通过将多个操作折叠进同一个核函数，彻底绕开了多算子调用带来的启动与同步开销。该方法由四个关键部件构成：首先是 TC 语言前端，它借鉴了爱因斯坦求和约定，开发者只需写出逐元素的数学定义，例如标准矩阵乘法可极简表达为 $C(i,j) +=! A(i,k) * B(k,j)$，编译器会自动推断索引变量的最大合法范围，保证不越界。其次是多面体调度器，它将 TC 降级为多面体中间表示，把循环嵌套转化为多维空间中的迭代域。调度器利用 isl 库构建数据依赖图，通过下沉无依赖的并行循环来实现不完美嵌套循环的分块（Tiling），专门针对深度学习中长依赖和非均匀重用模式进行优化。第三是 GPU 映射与内存提升部件，它将最外层并行的调度带映射为 GPU 块，最内层映射为线程，并通过贪心启发式算法，将频繁访问的张量瓦块自动提升到 GPU 共享内存中，以满足合并访存需求。最后是自动调优框架，由于 JIT 编译较慢，系统利用遗传算法在后台并行搜索最佳的分块大小、展开因子和网格维度，并将针对特定输入尺寸调优后的最优 CUDA 代码序列化存入编译缓存，供后续直接命中调用。

效果如何 实验在配备 Tesla M40 (Maxwell 架构) 和 Tesla P100 (Pascal 架构) 的多 GPU 节点上进行，对比基线包括代表极致手工优化的通用底层库 CUBLAS 和 CUDNN，以及代表极速生产环境的 Caffe2 和 ATen 框架。在通用核函数测试中，对于非标准尺寸的转置批处理矩阵乘法，TC 在 Pascal 架构上相比 CUBLAS 实现了 3.7 倍的加速；在 ResNeXt 等模型使用的分组卷积任务上，TC 生成的核函数相比 Caffe2 调用的 CUDNN Winograd 算法实现了 1.9 倍到 8.8 倍的巨大加速。在真实的推荐系统生产模型测试中，针对包含连续 3 个低延迟全连接层的 MLP3 模块，TC 成功将其 JIT 融合成单一核函数，相比混合调用 CUDNN 和 CUBLAS 的 Caffe2 参考实现，运行速度提升了 3.6 倍。此外，在极小尺寸的低延迟转置矩阵乘法中，TC 耗时仅 4.6µs，远快于 CUBLAS 专用核函数的 15.9µs。不过作者也承认了该方法的局限性：由于目前缺乏寄存器级别的分块和汇编级别的极致优化，TC 在大尺寸的标准矩阵乘法上比 CUBLAS 慢 3.4 到 4.2 倍；同时，当前版本的范围推断机制尚不支持实现 RNN 所需的递归定义。

A1 主要贡献

本文针对深度学习领域中新算子开发成本高、性能差以及现有库无法针对特定网络架构和数据进行优化的问题，提出了一套完整的解决方案。

核心问题与研究目标:
深度学习框架（如TensorFlow, PyTorch）依赖于高性能库（如CUDNN）来执行计算图中的算子。当研究人员发明新算子或现有算子无法满足特定网络和数据的性能需求时，面临两大挑战：

  1. 工程成本高昂：开发自定义的高性能算子需要大量工程投入，且通常性能不佳，这限制了创新速度。
  2. 优化缺失：即使使用现有库，也常常错失跨算子的优化机会，并且库函数未针对用户特定的数据大小和形状进行调优。

为了解决这些问题，本文旨在设计一种编程语言和编译流程，它必须同时满足两个要求：

创新点与主要贡献:
本文的主要贡献包括：

  1. Tensor Comprehensions (TC)：一种接近深度学习数学表示的高级语言，其语法概括了爱因斯坦求和约定。它支持形状和大小推断，语法简洁且安全，能避免“差一错误”（off-by-one errors），并允许布局转换和特化。
  2. 端到端的编译流程：一个能将TC表达式降级为高效GPU代码（CUDA核函数）的完整流程。该流程利用多面体编译技术，实现了算子融合和针对特定尺寸的特化，同时将内存管理和同步委托给上层框架。
  3. 领域定制的多面体编译算法：与通用并行编译器不同，本文的编译器专注于通过核函数融合来减少启动和同步开销，并优先考虑多级并行性和数据向内存层次结构更深层的提升（promotion）。
  4. 自动调优框架：结合即时编译（JIT）和代码缓存，该框架能对非标准尺寸进行特化，消除控制流和地址生成逻辑，并全面控制从机器学习框架到代码生成器的所有优化选项。
  5. 主流框架集成：该系统已集成到两个主流的机器学习框架中：面向生产的Caffe2和面向研究的PyTorch，通过ATen异步张量库实现。

通过这些贡献，本文的流程在机器学习社区关注的核心算子以及Facebook生产环境中使用的真实模型上，相比NVIDIA库实现了高达4倍的加速。

A2 相关工作

A3 方法细节

3 张量推导式（Tensor Comprehensions）

def mv(float(M,K) A, float(K) x) → (C) { 
  C(i) = 0 
  C(i) += A(i,k) * x(k)
}

这里定义了函数 mv，输入为张量 A 和 x，输出为 C。语句引入了两个索引变量 i 和 k。未在任何地方定义的变量会隐式成为索引变量。它们的范围根据其在索引中的用法推断得出（此处 i 的范围为 [0,M)，k 的范围为 [0,K)）。由于 k 只出现在右侧，因此对 C 的存储操作将使用 + 归约算子对 k 进行归约。归约可以跨多个变量，但它们都共享同一种关联和交换算子（如 +=），以确保求值顺序不影响计算值。

tensor C({M}).zero(); // 填充0的一维张量
parallel for (int i = 0; i < M; i++)
  reduction for (int k = 0; k < K; k++)
    C(i) += A(i,k) * x(k);

值得注意的是，嵌套顺序（i 然后 k）是任意的：张量推导式的语义总是通过循环置换保持不变。

def sgemm(float a, float b, float(N,M) A, float(M,K) B) → (C) { 
  C(i,j) = b * C(i,j)            # 初始化
  C(i,j) += a * A(i,k) * B(k,j) # 累加
}

def mv(float(M,K) A, float(K) x) → (C) {
  C(i) +=! A(i,k) * x(k)
}

def fcrelu(float(B,I) in, float(O,I) weight, float(I) bias) → (out) { 
  out(i,j) = bias(j) 
  out(b,o) += in(b,i) * weight(o,i) 
  out(i,j) = fmaxf(out(i,j), 0)
}

这里选择在所有推导式中重用 out 张量，表明没有临时存储。

def conv2d(float(B,IP,H,W) in, float(OP,IP,KH,KW) weight) → (out) { 
  out(b,op,h,w) +=! in(b,ip, h + kh, w + kw) * weight(op,ip,kh,kw)
}

其中归约维度是 kh 和 kw。

def maxpool2x2(float(B,C,H,W) in) → (out) { 
  out(b,c,i,j) max=! in(b,c, 2 * i + kw, 2 * j + kh) 
  where kw in 0:2, kh in 0:2
}

在最大池化的情况下，索引 kw 和 kh 是欠约束的，因为它们无法从任何输入张量推断出来，若不提供额外信息，范围推断程序会报错。因此，我们提供 where 注解来告知推断算法这些变量的预期范围，并让它从上下文中推断其余范围。

def gather(float(N) X, int(A,B) I) → (Z) { 
  Z(i,j) = X(I(i,j))
}

def sconv2d(int sh, int sw, float(N,C,H,W) I, float(F,C,KH,KW) W, float(F) B) → (O) { 
  O(n,f,h,w) = B(f) 
  O(n,f,h,w) += I(n,c, sh * h + kh, sw * w + kw) * W(f,c,kh,kw)
}

值得注意的是，sconv2d 的高效实现可以利用部分求值（当 sh 和 sw 是常量时），从而得到仿射下标表达式。这是JIT编译TC的另一个优势。

3.1 范围推断
def conv1d(float(M) I, float(N) K) → (O) {
  O(i) = K(x) * I(i + x)
}

这里存在歧义。如果首先最大化 i，它将覆盖整个 I，迫使 x 只能取单个值，这不符合卷积的预期。如果首先最大化 x，使其覆盖整个 K，那么为避免越界，i 的范围必须更小，这才是期望的行为。为了在不需显式注解的情况下解决这种歧义，范围推断分轮进行。我们维护一个未解析变量集合。在每一轮中，我们考虑所有只包含一个未解析变量的张量参数表达式，并构造一个布尔表达式来确保访问不越界。然后使用Halide的工具（solve_for_inner_interval）来找到满足此条件的变量的最大范围。

3.2 数据布局转换

JIT编译流程将TC降级为Halide-IR，然后到Polyhedral-IR，接着进行优化、代码生成和执行
图 2: JIT编译流程将TC降级为Halide-IR，然后到Polyhedral-IR，接着进行优化、代码生成和执行

4 高层工作流描述

5 多面体即时编译（JIT Compilation）

def outerProductMM(float(P,Q,R) A, float(S,R,T) B) → (O) {
  O(p,s,q,t) +=! A(p,q,r) * B(s,r,t)
}

如果需要，还可以进一步转换布局，派生出QPTS版本（以输出维度顺序命名）而不是PSQT版本。

sgemm的优化步骤
图 3: 图1中sgemm的优化步骤。(a) 规范形式, (b) 融合后, (c) 融合并分块, (d) 融合、分块并下沉, (e) 融合、分块、下沉并映射

5.1 多面体变换与映射
5.2 调度
5.3 不完美嵌套循环分块
5.4 映射到块和线程
5.5 内存提升
5.6 映射到寄存器

6 自动调优与缓存

6.1 编译缓存
6.2 自动调优

用于核函数的多线程自动调优管道
图 4: 用于核函数的多线程自动调优管道

A4 实验环境

A4 实验结果

所有实验均在GPU上运行1000次，报告了p0（最小值）、p50（中位数）和p90（90百分位）的端到端执行时间（包括CPU开销）。

通用与研究性核函数 (图5)

通用核函数的执行时间（µs）
生产模型 (图6)

生产模型的执行时间（µs）
生产模型的执行时间（µs）

总体性能加速比 (图7)

该图表以Caffe2-CUBLAS为基准，展示了在最小数据集上，TC（基线和自动调优版本）以及ATen+CUBLAS的p50（中位数）运行时间加速比。结果清晰地显示，TC的自动调优版本在多个基准测试（如2LUT, MLP3, gconv）中都取得了显著的性能领先。

p50运行时间加速比，以Caffe2-CUBLAS为基准
p50运行时间加速比，以Caffe2-CUBLAS为基准

A5 结论

本文成功展示了一个从高层语言Tensor Comprehensions (TC) 到在GPU上自动生成高性能核函数的端到端流程。TC语言的设计接近于深度学习的数学模型，使得推理、交流和手动调整计算与存储权衡变得容易。该流程利用了数十年来在多面体编译领域的进展，并实现了领域特定的优化、代码生成、带编译缓存的自动调优，以及通过ATen张量库与Caffe2和PyTorch的无缝集成。

核心贡献与影响:

未来展望:
本文为未来的研究和应用开辟了众多方向，包括：

  1. 通过protobuf分发和共享不同架构下的最佳实现和自动调优历史。
  2. 移植到更多硬件架构，并结合用于高单线程性能的基元库。
  3. 系统地利用数据布局转换，将其作为自动调优的一部分，并支持新的数据类型（如低精度格式）。
  4. 开发自动的DAG划分算法。
  5. 结合Halide风格的变换以支持算法分块和模型并行。
  6. 直接在TC上提供符号自动微分。
  7. 支持更动态的控制流和稀疏、向量化及混合精度数据表示。
  8. 总体上，通过提升性能和简化从数学规范到实际实现的转换过程，加速机器学习研究。

A6 附录

A.1 机器学习框架简介

A.2 ML框架接口API

构建执行引擎
图 11: 构建执行引擎

JIT编译、调优或命中编译缓存，然后运行
图 12: JIT编译、调优或命中编译缓存，然后运行

A.3 多面体编译背景

TC核心语法的EBNF表示
图 14: TC核心语法的EBNF表示

A.4 TC示例的详细结果

tc::IslKernelOptions::makeDefaultMappingOptions()
    .scheduleSpecialize(true)
    .tile({32, 32, 32})
    .mapToThreads({32, 32})
    .mapToBlocks({M / 32, N / 32})
    .useSharedMemory(true)
    .usePrivateMemory(true)
    .unrollCopyShared(true)
    .unroll(256);

tc::IslKernelOptions::makeDefaultMappingOptions()
    .tile({1})
    .mapToThreads({128})
    .mapToBlocks({B})
    // ...

auto threads = (W >= 10) ? std::vector<size_t>{W / 4, H / 2} : std::vector<size_t>{4, 8, 4};
auto options = tc::IslKernelOptions::makeDefaultMappingOptions()
    .tile({1, 1, 1})
    .mapToThreads(threads)
    // ...

// MLP3 baseline
IslKernelOptions options = makeDefaultMappingOptions()
    .set_fusion_strategy(FusionStrategy::Max)
    .tile({1})
    .mapToThreads({128})
    .mapToBlocks({128});

对于像CUDNN和NNPACK这样的库中新增加的简单逐点融合功能，其API与ML框架集成时存在阻抗失配问题，给最终用户和研究人员带来了认知开销和使用上的不便。TC则通过其统一和富有表现力的语言避免了这些问题。