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

作者/机构: 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)

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的无缝集成。

核心贡献与影响:
* TC能够快速合成性能可靠的基线版本，有效解决大型训练中的瓶颈问题，从而避免机器学习研究因性能问题而减速。
* 该工作解决了机器学习领域的生产力和效率差距，将更强的表达能力和控制权交给了领域专家，减少了ML框架对高度优化的供应商库的依赖，同时不牺牲性能。
* TC自动化了许多在各大深度学习框架中被重复实现的样板代码，并建立在一个与其它领域（如图像处理）和通用编译器（LLVM/Polly）共享的通用多面体中间表示之上。

未来展望:
本文为未来的研究和应用开辟了众多方向，包括：
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则通过其统一和富有表现力的语言避免了这些问题。