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 相关工作

• 传统优化编译器的局限性
  尽管优化和并行化编译技术发展了数十年，但在计算密集型应用上，其性能仍远逊于手动优化的峰值性能。原因包括现代处理器的复杂性、编译器缺乏领域知识、变换收益难以评估，以及复杂变换难以组合等问题【索引23，CGO 2016】【索引25，LCPC 2005】【索引32，IJPP 2006】【索引15，USC TR 2008】【索引48，ARRAY 2014】【索引4，CGO 2016】。

• 领域特定程序生成器与主动库
  为解决上述问题，研究人员转向设计特定应用的程序生成器，即“主动库”（active libraries）【索引85，Dagstuhl Workshop 2003】。这类生成器常利用反馈导向优化来选择最佳生成策略，如用于密集矩阵运算的ATLAS【索引92，SC 1998】和BTO【索引9，SC 2009】，以及用于快速傅里叶变换的FFTW【索引31，ICASSP 1998】。SPIRAL项目【索引69，IJHPCA 2004】在此基础上实现了飞跃，它直接操作于数字信号处理公式的领域特定语言（DSL）。

• 面向图像处理和数值计算的DSL
  DSL编译器通过领域特定的构造来捕捉应用的内在并行性和局部性，取得了显著成果。例如，用于图像处理的Halide【索引72，PLDI 2013】。与Halide相比，TC更适合处理机器学习中固定尺寸、具有更高时间局部性的张量计算，且语法更简洁。OoLaLa【索引53，OOPSLA 2000】对线性代数采取了类似方法，而TACO【索引46，OOPSLA 2017】和Simit【索引47，TOG 2016】使用与TC类似的表示法，但主要生成用于数值求解器的稀疏矩阵代码。

• 通用代码生成与多面体框架
  本文不仅设计了新的DSL和编译器，还提出了一个更通用的代码生成与优化框架，整合了高性能计算领域数十年在循环嵌套优化和并行化方面的研究成果。该框架使用的多面体模型能够表示Halide【索引72，PLDI 2013】【索引57，TOG 2016】、Latte【索引82，PLDI 2016】或XLA【索引36，TensorFlow White Paper 2017】无法表示的复杂变换，如分层分块、映射、移位、融合、分布、交换等。

• 多面体编译工具与领域特定应用
  多面体框架是分析和变换循环嵌套的强大抽象，已催生了多种工具和库【索引28，IJPP 1992】【索引12，PLDI 2008】【索引88，TACO 2013】【索引11，TOPLAS 2016】【索引95，Inria TR 2017】，并被集成到GCC（Graphite）和LLVM（Polly）等生产编译器中。多面体技术也被用于特定领域，例如PolyMage DSL【索引59，ASPLOS 2015】用于图像处理，PENCIL方法【索引3，PACT 2015】【索引7，LCTES 2014】用于构建DSL的并行化编译器。本文编译器实现了针对深度学习模型中长而非均匀的重用模式和深度嵌套循环的特定优化，这些启发式策略在Halide及其变体中不可用。

• 与深度学习编译器的比较
  与Latte【索引82，PLDI 2016】相比，TC同样提供了高层次DSL和端到端流程，但在实现自定义层时表达力相当的同时，TC通过类型和形状推断更简洁，通过静态边界检查和图连接性更安全，并通过解耦索引与表示布局更灵活。TC还实现了比Latte更复杂的调度和映射变换，并被设计为JIT编译库以便与深度学习框架无缝集成。
  与XLA【索引36，TensorFlow White Paper 2017】相比，TC同样提供自动形状和大小推断，可作为JIT库“进程内”运行，并集成了生产级框架。但TC保持了与特定计算图框架的独立性，同时与生产框架紧密集成；并且TC没有采用嵌入式DSL方法，从而将用户与嵌入式DSL的复杂性和调试难题隔离开。

• 与R-Stream·TF的比较
  最近提出的R-Stream·TF【索引67，ESPT 2017】是R-Stream多面体编译器在TensorFlow算子自动优化上的概念验证。与本文方法类似，它将生成代码包装为TensorFlow的自定义算子。但R-Stream·TF在静态内存管理和核函数划分方面非常激进，它将大部分决策留给了编译器。相比之下，TC选择将这些决策留在张量代数层面，允许领域专家或优化器重写声明式的TC以优化容量和布局。TC在其领域专用的仿射调度和映射上更具雄心，旨在为神经网络中高维度、非均匀、长重用模式生成单个加速核函数。

A3 方法细节

3 张量推导式（Tensor Comprehensions）

• TC语言概述
  本文提供了一种使用张量推导式（Tensor Comprehensions, TC）来表达张量逐元素计算的语言。TC是一种在多维数组上进行计算的表示法，借鉴了爱因斯坦求和约定（Einstein notation）的特点：

  1. 索引变量通过在表达式中使用来隐式定义，其范围根据其索引的对象推断得出。
  2. 出现在表达式右侧但未出现在左侧的索引被假定为归约维度。
  3. 迭代空间中点的求值顺序不影响输出结果。

• 矩阵-向量乘法示例
  以一个简单的矩阵-向量乘积为例，该张量推导式包含两个语句：

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 进行归约。归约可以跨多个变量，但它们都共享同一种关联和交换算子（如 +=），以确保求值顺序不影响计算值。

• 与C风格伪代码的等价性
  直观上，一个张量推导式可以被看作是一个循环体，其控制流由上下文推断。其等价的C风格伪代码如下：

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）是任意的：张量推导式的语义总是通过循环置换保持不变。

• 原地更新与功能语义
  TC允许原地更新，但保留了在整个张量上原子性的功能语义：其语义是在对左侧（LHS）任何元素赋值之前，完整读取右侧（RHS）的表达式。这一规范在LHS张量也出现在RHS时非常重要【索引30，Parallel Computing 2012】。编译器负责检查原地更新在逐元素依赖上的因果关系。目前，本文实现了一个简单的句法检查，仅允许在逐点定义和张量收缩上进行原地更新。当检查失败时，编译器会因活跃性干扰而拒绝程序；例如，任何原地转置 a(i,j) = a(j,i) 都是不正确的（除非范围为空或单个元素），而 a(i,j) = b(j,i) 是一个有效的转置。这种显式重用和原子更新语义并不妨碍编译器做出其他调度和存储映射决策，只要这些决策保留了TC的逐元素依赖性。TC的这种混合声明式-命令式设计受到了Lush【索引13，Lush Manual 2002】的启发。

• SGEMM示例
  BLAS库中的SGEMM函数可以如下表示：

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) # 累加
}

• 初始化归约的快捷方式
  TC中的归约通常使用归约算子的单位元进行初始化；本文提供了一个初始化归约的快捷方式，即在归约符号后附加 !，例如用 +=! 代替 +=。下面是矩阵-向量乘积的一行定义：

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

• 全连接层与ReLU示例
  使用这些简单的属性，常见的机器学习核函数可以仅用几行代码写出。例如，一个全连接层后接一个修正线性单元（ReLU）可以表示为一个转置矩阵乘法，该乘法用一个广播的偏置项初始化，然后进行逐点钳位（即与0进行 fmaxf）：

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 张量，表明没有临时存储。

• 二维卷积示例
  一个二维卷积同样简单，其归约初始化为0（注意使用了 +=!）：

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 注解来告知推断算法这些变量的预期范围，并让它从上下文中推断其余范围。

• 下标表达式的灵活性
  下标表达式可以是迭代子的任意仿射函数，或下标的下标表达式，以及它们的组合。后者可以捕捉数据依赖的访问，例如一个gather操作：

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

• 步进卷积示例
  TC与算法表示法紧密匹配，这使得对层进行微小变动变得容易。例如，步进卷积可以通过对标准卷积进行微调来轻松创建。以下是一个沿 h 和 w 方向步长分别为 sh 和 sw 的卷积：

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 范围推断

• 范围推断的简洁性与必要性
  张量推导式的简洁性很大程度上源于循环范围可以从上下文中推断。类似于多态数据类型的推断，范围推断算法旨在根据使用模式准确推断范围。然而，由于非仿射表达式或欠约束问题（如最大池化），这并非总是可行。TC需要为这些情况提供额外的注解。

• 基本规则
  迭代变量总是非负的。除非在 where 子句中另有说明，否则每个迭代变量都假定从0开始。

• 推断方法
  推断过程是从输入参数到输出张量，它在TC函数的所有仿射数组访问中建立一个基于约束的分析问题。我们采用一种直接的方法以保证可用性：只在范围显而易见的情况下进行推断，并在其他情况下要求显式注解。

• 最大化范围规则
  我们推断的矩形范围是在不读取输入越界的情况下尽可能大。如果存在多种方式实现这一点，我们会抛出错误并要求使用 where 子句进行显式注解。这个规则足以理解 sgemm 的情况：i 的范围是 A 的行数，j 的范围是 B 的列数，k 的范围是 A 的列数和 B 的行数中较小的一个（由类型签名约束为相等）。

• 多轮推断处理歧义
  考虑一个简单的一维卷积：

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）来找到满足此条件的变量的最大范围。

• 推断过程示例
  对于上面的卷积例子，第一轮忽略 I(i+x) 因为它包含多个未解析变量。我们使用 K(x) 推断出 x 的范围。在第二轮中，I(i+x) 现在只包含一个未解析变量 i，我们使用已推断的 x 的范围来推断 i 的最大范围。具体来说，K(x) 约束 x 的范围是 0 到 N-1，而 I(i+x) 约束 i+x 的范围是 0 到 M-1，由此推断出 i 必须满足以下约束：
  
  这最终得出 i 的范围是 0 到 M-N。

• 通用量化与语义选择
  请注意 i 上的 ∀ 通用量化，它约束输出张量 O 在其整个定义域上由完全相同的输入 I 均匀定义。这与Alpha语言【索引51，J VLSI Signal Process Syst 1991】中采用的方法不同。选择通用量化迭代子的语义源自Halide，它保证了输入张量的形状总是（超）平行六面体，对机器学习专家很直观，并有助于生成紧凑且无条件控制流的代码。

• 复杂情况与临时张量
  更复杂的例子涉及在同一变量上交叉进行多轮推断，以及无法从一组更一般的约束中导出单个矩形形状的模糊情况。后者将无法通过类型检查，编译器会要求插入 where 子句以消除歧义。此外，由于内存管理外化给机器学习框架，TC内部定义的临时张量列表（及其形状和范围）必须明确，范围推断算法会自动推断此列表。

• 当前局限性
  在首个版本中，TC不支持实现循环神经网络（RNN）所需的递归定义。张量索引可以涉及任何整数值表达式，只要它不依赖于当前语句的LHS张量。编译器会尽力进行范围推断和静态验证，确保在命令式定义的维度中每个元素只被定义一次；如果范围推断失败且没有where子句提供信息，编译将失败。

3.2 数据布局转换

• TC对数据布局转换的简化作用
  TC使得全局数据布局转换变得显著容易。机器学习社区一直严重依赖这类转换，通过对张量元数据（形式为(dataPtr, offset, size[], stride[])的元组）进行操作组合来实现【索引9，SC 2009】。研究人员主要使用这些原语进行算法分块（algorithmic tiling）和层次化分解。前者与数据分块（data tiling）密切相关，在频域高性能卷积实现中无处不在，用于控制内存占用【索引83，CoRR 2014】或适应快速局部存储器【索引2，CoRR 2015】。Lush【索引13，Lush Manual 2002】和Torch【索引24，Book Chapter 2012】中的unfold操作完美匹配隐式数据分块（即无显式内存拷贝），当没有部分边界效应时。

• TC支持的布局转换类型
  TC支持广义的张量转置（即应用一个n > 2的n维置换矩阵），并且数据分块可以通过简单地重塑张量和调整TC索引表达式来实现。范围推断和检查保证了这种重塑在整个TC中始终保持一致。结构数组到数组结构的转换以及类似的短向量操作是其副产品：它们是维度交换和数据分块的特例。

• 当前实现状态
  目前，数据布局和TC转换由领域专家在源码级别完成，TC的推断过程保证了表达式具有正确的范围。未来，随着在TC中引入自动数据布局转换，这一假设将被重新评估。

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

4 高层工作流描述

• 在深度学习框架中的定位
  本文的工作是在TensorFlow、Caffe2和PyTorch等深度学习框架的背景下进行的。

• 与ML框架的集成方式
  TC表达式通过“进程内”（in process）实现与机器学习框架集成，简化了与计算图引擎和上层应用的交互，这对于一个全自动调度和映射流程来说是独一无二的。我们提供了一个轻量级API，将特定框架的张量对象模型转换为我们自己的模型。算子定义被重写以生成TC而不是框架的后端实现，用户也可以编写自己的TC。在这种情况下，单个TC可能对应ML框架中的一个算子有向无环图（DAG）。目前这种匹配是手动完成的，未来的工作将包括自动的DAG划分、匹配和重写。

• JIT编译流程
  TCs随后被即时编译（JIT），流程如图2所示。如果框架的默认后端实现性能更好，我们会回退到参考实现。编译流程从一个具有特定张量大小和步幅的TC开始，首先将其降级为一个参数化的Halide表达式。

• 从Halide-IR到多面体表示
  在系统的原型版本中，Halide-IR通过翻译到PENCIL【索引3，PACT 2015】再由pet【索引89，IMPACT 2012】和isl【索引86，ICMS 2010】库解析来降级为多面体表示。为了减少IR之间的阻抗失配并促进语义注解的传播，流程演变为将Halide-IR直接降级为多面体表示，绕过了PENCIL中间语言。

• CUDA核函数生成
  类似地，最初原型的CUDA核函数生成是从修改版的PPCG编译器【索引88，TACO 2013】引导的。修改包括将其用作库，以支持进程内、多线程操作（见第5节）。在放弃PENCIL后，PPCG的功能被从头开始重新实现，旨在提高模块化并符合现代C++环境。

• 自动调优与代码缓存
  作为该流程的补充，一个自动调优器和一个可序列化的编译缓存与调度和映射策略紧密交互，以搜索变换空间（见第6节）。

• 其他代码生成路径
  此外，我们提供了一个简单的“恒等”多面体映射选项，以生成一个简单、可读的CUDA参考实现，可在单个GPU线程上运行和检查正确性。我们计划增加一个简单的LLVM JIT编译通道，使该参考实现更普遍地适用于CPU，但这尚未实现。最后，我们将提供一个路径来生成一系列库调用，作为回退到由CUDNN支持的默认实现的有用方式。

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

• 逻辑布局到数据格式的转换
  编译器解决了高层张量操作的逻辑布局（维度顺序）与多面体代码生成器期望的数据格式（C99数组【索引3，PACT 2015】）之间的阻抗失配问题。降级步骤确保大小和步幅的组合对应于非混叠的数组和子数组语法；它还通过分析访问关系和推断的张量范围来确保没有越界访问；它生成张量声明并重排表达式以匹配目标语言的数据模型，即行主序数组。

• 处理别名和广播语义
  张量规范可能包含输入和输出参数的别名，用于原地计算和高维张量的隐式转换。我们认为这种规范应导致多版本化和对每种别名场景的特化。此外，TC采用的基于范围推断的语义，不同于XLA、PyTorch和MXNet等采用的NumPy风格的“广播”语义。TC不需要这种隐式语法糖。例如，对应于所谓的外积矩阵乘法 [p,q,r] matmul [1,s,r,t] → [p,s,q,t] 的TC非常简单：

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版本。

• 其他降级步骤
  额外的降级步骤包括卷积表达式的前向替换（存储/计算权衡）、零填充、镜像和裁剪。

• 多面体框架的适用性
  多面体框架为“大部分规则”的【索引29，Encyclopedia of Parallel Computing 2011】基于循环的程序提供了一种先进的内部编译器表示。它能描述由循环和分支包围的算术表达式，其中控制流可以静态确定。因此，多面体框架作用于程序的静态控制部分（SCoP）。具体来说，循环边界和分支条件是外层循环迭代子和静态未知常数值（被视为参数）的仿射表达式。本文证明了多面体框架特别适用于深度神经网络，这些网络通常与大型、深度嵌套的循环相关，具有长依赖链和非均匀或全连接模式。这些特性将优化问题推向了与Halide用于图像处理不同的启发式空间，以及一个远比单纯线性代数更广阔的空间。

• 命名关系与调度树
  我们使用iscc【索引87，IMPACT 2011】中引入的命名关系表示法。调度树【索引90，IMPACT 2014】可以将仿射映射组合起来，以方便将高层语言（TC）的属性传递给下游优化器，并沿途附加目标特定信息（如CUDA中的SPMD线程相对归纳、同步、数据传输指令）。调度树引入了特定节点类型：

  • Band节点: 通过一个或多个分段仿射函数定义部分执行顺序。
  • Filter节点: 划分迭代空间，将其子树绑定到迭代域的子集。
  • Set/Sequence节点: 组织Filter节点，取决于执行顺序是否需要被约束。
  • Context节点: 提供关于变量和参数的额外信息，如张量范围或GPU网格/块大小。
  • Extension节点: 引入不属于原始迭代域的辅助计算，例如用于在共享内存之间拷贝数据的语句。

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

• TC的规范调度树
  一个TC的规范调度树由一个外部Sequence节点定义，后跟每个TC语句的Filter节点。在每个过滤分支内部，Band节点定义了一个恒等调度。图3(a)展示了sgemm TC的规范调度树，它包含一个用于初始化语句的二维循环嵌套和一个用于更新语句的三维循环嵌套。通过观察到张量C在两个循环嵌套间被重用，可以构造如图3(b)所示的调度，通过循环融合来利用访问局部性。

• 越界访问检测
  多面体模型允许对张量访问进行关系编码。通过将这些关系与表示为集合的迭代域进行组合，可以计算出所有被访问的张量元素的集合，即语句的足迹（footprint），并检查其是否符合张量大小。具体来说，访问关系使得检测越界访问成为可能。属于足迹 F = D.A 但不属于从张量形状推断出的张量元素集合 T 的元素对应于越界访问。因此，(D.A)\T = ∅ 是没有越界访问的条件。

5.1 多面体变换与映射

• 变换与映射概述
  当面向并行架构时，程序变换涉及调度和映射策略的改变；这些必须在尊重依赖关系的同时，为目标特定属性进行优化。除了保证变换的有效性，依赖关系还可用于暴露并行性（独立实例可并行执行）或改善数据访问局部性（依赖实例在时间上紧密执行）。存在多种高效的调度算法，专注于并行性、局部性和向量化的组合【索引91，KU Leuven TR 2017】【索引12，PLDI 2008】【索引95，Inria TR 2017】【索引84，IMPACT 2012】【索引65，POPL 2011】。我们的变换引擎基于四个组件：
  1. 核心多面体调度：由 isl 提供，自动优化（外层）循环并行性和局部性；我们调整了仿射调度启发式，以将完整的TC程序折叠为单个GPU核函数。
  2. 分块（Tiling）：对调度进行进一步分块，以促进在GPU的深度并行和内存层次结构上的映射和时间重用【索引91，KU Leuven TR 2017】。
  3. GPU映射：借鉴了先前在PPCG【索引88，TACO 2013】中实现的算法，并进行了扩展以支持机器学习核函数中出现的更复杂和不完美嵌套的结构。
  4. 内存提升（Promotion）：处理到共享和私有内存空间的显式数据传输。

5.2 调度

• isl调度器及其扩展
  isl调度器的核心部分迭代地求解整数线性规划问题，以计算构成调度带（schedule bands）的分段仿射函数。它内部构建一个数据依赖图（DDG），其中节点对应语句，边表示它们之间的依赖关系。isl调度器【索引91，KU Leuven TR 2017】通过仿射聚类技术提高可扩展性，该技术基于为DDG的各个强连通分量分别计算调度带，然后迭代地聚类这些分量并相互调度。我们扩展了isl调度器，为调用者提供了对调度过程更细粒度的控制。对于仿射变换，用户可以提供额外的任意约束插入到线性规划中。对于聚类，用户可以提供一个决策函数来组合依赖图分量。这些配置点构成了调度策略的基础，允许限制仿射变换，控制融合目标（例如，保留至少三个嵌套的并行循环），并通过自动调优过程进行选择和配置。在所有情况下，我们都强制生成单个GPU核函数。

5.3 不完美嵌套循环分块

• 分块作为调度树变换
  循环分块（tiling）是在调度之后作为一个独立的步骤实现的，并作为调度树的变换来执行。它将一个可置换的调度带转换为两个带的链，外层带包含瓦块循环（tile loops），内层带包含固定行程次数的点循环（point loops）。例如，图3(c)展示了融合并分块后的sgemm的调度树。

• 不完美嵌套循环的分块技术
  除了常规的循环嵌套分块，调度树表示还允许对不完美嵌套的循环进行分块。该技术基于一个观察：如果不携带依赖关系的循环（即并行循环）可以被下沉到任何其他循环之下。并行循环不携带依赖，因此不影响依赖的满足或违反。因此，不完美嵌套分块的实现方式是，首先独立地对所有带进行分块，然后在树中下沉并行的点循环。在此过程中，点循环带在序列（或集合）节点下的每个子树中被复制，并更新其调度以使其仅包含相关的域点。

• sgemm示例中的应用
  sgemm的调度树特意设计为两个不完美嵌套的带。依赖分析表明循环 i 和 j 是并行的。因此，我们可以对它们进行分块，并将点循环下沉到归约循环 k 的带之下，得到如图3(d)所示的调度树。最内层的嵌套带中的点循环在检查可置换性后可以合并成一个单独的带。

5.4 映射到块和线程

• 调度树与GPU映射
  调度树也可以用来表示到加速器（特别是具有多个块和线程的GPU）的映射。这个操作通过将特定的调度带成员（以及相应的循环）与线程或块索引相关联来完成。多面体代码生成器随后会尽可能省略这些循环，并相应地重写索引表达式。我们的映射算法源自最初为PPCG设计的算法，其中网格和块的大小独立于瓦块大小指定。由于块和线程的语义，只有属于可置换调度带的并行循环可以被映射。

• 映射算法细节
  我们要求调度树至少有一个最外层的带，具有外部并行维度。与PPCG不同，我们的映射算法将单个带映射到块，以生成ML框架期望的单个核函数，而多个带可以映射到线程。单个最外层带的并行维度被映射到GPU块。在每个调度树分支中，最内层的可置换带（通常由点循环组成）被映射到GPU线程，但映射的维度数必须在所有分支中相等。

• 映射的实现
  映射本身是通过插入特殊名称（通过上下文节点传达给代码生成器）并将其与过滤节点中的带维度相关联来执行的。对于矩阵乘法示例，我们的映射策略产生了图3(e)中的调度树。我们引入了一个上下文节点来指明参数的有效大小以及网格和块的大小（分别表示为bx, by和tx, ty）。

5.5 内存提升

• 内存提升的目标与接口
  我们关注将部分张量提升到共享或私有GPU内存中。虽然提升决策由启发式算法做出，并且相应的命令式代码在后续阶段生成，但调度树为附加内存相关信息提供了方便的接口。内存提升基于数组瓦块（array tile）的概念，这是一种用于软件控制的局部存储器的数据分块形式。我们重访并扩展了PPCG对内存提升的支持【索引88，TACO 2013】【索引91，KU Leuven TR 2017】。

• 间接访问数组的提升
  我们的内存提升方法也处理间接访问的数组。考虑访问O[l+Idx[i][j]][k]，我们将O称为外部数组，Idx称为索引数组。对于嵌套间接访问，我们从内到外迭代处理外/索引对。虽然外部数组第一个索引表达式的值在静态时是未知的，但我们仍可将其本地缓存为shared_O[l][i][j][k] = O[l + Idx[i][j]][k]。由于某些值可能被复制，间接提升仅在外部和索引数组都只读时才可能。我们要求索引数组有一个数组瓦块，即只访问其固定大小的块。在计算外部数组的数组瓦块时，我们忽略下标的间接部分。然后在提升的外部数组中引入与索引数组中一样多的额外索引表达式。这些新维度上数组的范围与索引数组的数组瓦块大小完全对应。

• 提升启发式
  直接访问的数组在以下情况下被提升到共享内存：存在固定大小的数组瓦块；单个元素被多次访问；至少有一次访问不满足内存合并。后者可以从应用了调度的访问关系中看出：最后的访问维度应与映射到x线程的调度维度对齐。对于间接数组，由于存在额外的长内存依赖，合并要求被放宽。由于共享内存总量固定，我们应用简单的贪婪启发式，如果所需的共享内存量超过可用量，则拒绝提升。

5.6 映射到寄存器

• 当前策略与理由
  目前没有计划实现比PPCG先前支持的更复杂的寄存器提升策略。这是一个临时的、务实的选择，基于以下观察：

  • 除了少数情况【索引69，IJHPCA 2004】【索引80，IPDPS 2011】，没有经验证据表明从高层规范自动生成底层代码比汇编能带来显著的性能增益。此外，现代汇编生成器现已公开可用且可重定向【索引26，PyHPC 2013】。
  • 将寄存器优化遍（pass）推广并重定向到具有不同向量长度、对齐约束和模式的不同架构并非易事。
  • 存在许多不同的策略，包括寄存器级流水线、寄存器轮换、多缓冲、具有集合式语义的新ISA【索引63，NVIDIA Blog 2017】，其中一些甚至难以实现。
  • 这类策略通常首先用内部函数和汇编作为概念验证来实现，并且可以很容易地用像PeachPy【索引26，PyHPC 2013】这样的工具进行模板化和重用。

• 性能影响与未来计划
  这一临时实现状态对性能结果有轻微影响。我们还计划依赖外部库调用，例如，用于加速归约操作。

6 自动调优与缓存

• JIT编译与缓存
  虽然多面体调度和GPU映射与训练神经网络相比开销很小，但对于传统的JIT设置来说仍然太慢。我们利用神经网络大部分是静态结构的特点，来缓存和重用在相似条件下（输入形状、目标架构、优化选项）对一个操作/核函数进行自动调优编译的最佳结果。我们通过一个编译缓存来利用这种重用。

6.1 编译缓存

• 缓存内容与键
  编译缓存存储了给定TC生成的CUDA或PTX代码。生成的代码取决于输入形状、选择的优化选项（如瓦块大小、循环融合调度策略、映射决策）以及目标GPU架构引入的约束（如共享内存/寄存器大小等）。因此，每个缓存条目的键是一个元组：(templatized TC, input_shapes, optimization_options, gpu_arch)。为了缓存，我们使用TC的摘要模板化版本，使键与名称变化无关。每个缓存条目都持有已知的最快版本。

• 缓存操作与持久化

  在每次核函数优化之前，可以查询缓存。如果缓存未命中，则调用常规的JIT编译流程。缓存被序列化为protocol buffer交换格式【索引35，Google Dev Guide 2017】，以实现持久化和重用。为了避免长时间和不可预测的编译和自动调优时间，可以用参考实现预填充缓存。由于缓存序列化为字符串，我们增加了向缓存注入条目的功能，以支持低级手动调优有益的特殊情况。

• 缓存大小与参数化
  在当前实现中，编译缓存的大小可能成为一个问题，但在实践中，感兴趣的操作数量有限，输入形状的空间非常稀疏。我们选择尽早将运行时大小注入到多面体编译流程中，因为它为调度、分块、映射和消除非循环控制流提供了有价值的信息。小批量（minibatch）维度是一个很好的参数化候选项，它通常映射到CUDA网格维度，可以在不重新生成代码的情况下通过核函数启动调用来改变。

6.2 自动调优

• 自动调优器流程
  自动调优器通过编译缓存与环境的其余部分交互，存储最佳版本以供后续使用。它包括以下步骤：

  1. 设置一组起始配置。
  2. 设置调优空间维度和可接受的值范围。
  3. 设置搜索类型——目前是遗传算法或随机搜索。
     自动调优器在规定的时间内运行，沿途用更好的版本更新缓存。

• 遗传搜索算法
  它使用遗传搜索【索引33，Book 1989】，对候选配置的“代”进行操作。在每个调优步骤中，候选配置被编译和评测，然后根据其运行时间的倒数分配一个适应度值。每个新候选通过三亲本均匀交叉繁殖，并以低概率进行突变。

• 并行自动调优管道
  自动调优需要评估成百上千个版本。我们利用遗传算法的特性，通过一个通用的多线程、多GPU自动调优器（如图4）并行启动多个编译任务。搜索策略将多个候选配置排入“编译作业”队列；这些配置由多个CPU线程并行编译。编译完成后，结果被排入“评测作业”队列；每个评测作业在可用的GPU上评估；评测结果用于更新自动调优数据库并生成下一组候选。

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

• 调优的核函数选项
  调优的核函数选项多种多样：
  • 尺寸选择：瓦块、块和网格大小的搜索范围被缩小到有助于避免尾部效应的选择，考虑2的幂和问题大小的整数上取整除数。
  • 展开因子：从调度树的叶子到根的自底向上路径上可以展开的迭代次数，考虑高达4096的2的幂。
  • 离散选择：关于融合策略和可接受调度形状的离散选择，例如，防止循环倾斜，用并行性换取局部性。
  • 低级选项：如共享或私有内存的使用，以及在每个块中多积极地使用这些内存空间，这直接影响GPU占用率。

A4 实验环境

• 硬件配置:

  • Maxwell系统: 8个节点，每个节点配备2个12核Intel Xeon E5-2680 v3 CPU，8个Tesla M40 GPU（每个12GB显存）。
  • Pascal系统: 8个节点，每个节点配备2个14核Intel Xeon E5-2680 v4 CPU，8个Tesla P100-SXM2 GPU（每个16GB显存）。

• 软件配置:

  • CUDA v8.0, CUBLAS v8.0.45, CUDNN v6.0.21, NVRTC v8.0。
  • 对比框架：Caffe2（被认为是速度最快的ML框架之一）和ATen（PyTorch的底层库）。
  • 代码实现基于C++，通过JIT编译生成CUDA核函数。

• 数据集/模型与用途:
  实验评估了多种在深度学习中常见的核函数以及一个在Facebook大规模使用的生产模型。

  • 通用核函数:
    • 转置矩阵乘法 (tmm): 测试不同尺寸下的性能，包括低延迟场景（例如 (M, K, N) = (128, 32, 256)）。
    • 转置批处理矩阵乘法 (tbmm): 针对因子分解机等应用中的常见尺寸进行测试，例如 (B, N, M, K) = (500, 26, 72, 26)。
    • 分组卷积 (gconv): 在ResNeXt等SOTA模型中使用，测试不同输入尺寸（例如 (H, W) = (56, 56) 到 (7, 7)）下的性能。
  • 生产模型 (MLP):
    • 一个用于Facebook部分应用的、小尺寸、低延迟的多层感知机模型，用于推荐系统和点击率预测。
    • 批大小: 128。
    • 模型组件:
      • 2LUT: 两个并行的查找表嵌入层，输入大小 (E1, E2, D, L1, L2) = (10^7, 10^7, 64, 50, 50)。
      • C3: 一个转置矩阵乘法层，权重矩阵大小为 (1000, 1024)。
      • MLP1: 一个fcrelu层。
      • MLP3: 连续3个低延迟MLP层的序列。

A4 实验结果

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

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

• 转置矩阵乘法 (tmm)

  • 结果: 在大尺寸问题上，TC的最佳自动调优版本比Caffe2+CUBLAS慢（Maxwell上慢4.2倍，Pascal上慢3.4倍）。但在低延迟场景（小尺寸张量）下，TC生成的核函数（4.6µs）比CUBLAS的专用核函数（15.9µs）快得多。
  • 分析: CUBLAS的sgemm是高度手动优化的，达到了硬件峰值性能，而TC目前缺少寄存器级分块和SASS级优化。尽管如此，TC在延迟敏感的场景中显示出优势，并且生成的核函数性能受限于共享内存带宽，已达到该约束下的较高水平。

• 转置批处理矩阵乘法 (tbmm)

  • 结果: 对于因子分解机相关的尺寸，TC相比CUBLAS实现了高达3.5倍（Maxwell）和3.7倍（Pascal）的加速。
  • 分析: TC能够为这种特定尺寸组合生成比通用库更优化的核函数。

• 分组卷积 (gconv)

  • 结果: 尽管TC目前只支持时域卷积，但其性能仍优于基于CUDNN Winograd算法的Caffe2实现，加速比在1.4倍到3.6倍之间（Maxwell），在Pascal上更是高达1.9倍到8.8倍。
  • 分析: 这突显了TC的核心贡献之一：能够为研究社区中新出现但尚未被厂商库完全优化的算子快速生成高性能实现，填补了从算法发现到库实现之间的时间差。

生产模型 (图6)

• 查找表嵌入 (LUT)

  • 结果: 对于融合的2LUT操作，TC生成的核函数比Caffe2参考实现（依赖CUB库）快4倍（Maxwell）到4.1倍（Pascal）。
  • 分析: TC通过一种新颖的两阶段加载共享内存技术优化了数据访问，有效解决了长延迟依赖问题。

• 转置矩阵乘法 (C3)

  • 结果: 对于(1000, 1024)的权重矩阵，TC在Pascal上已具有竞争力。
  • 分析: 与tmm类似，性能受限于寄存器优化，但在特定尺寸下表现良好。

• 单层感知机 (MLP1)

  • 结果: 尽管缺少寄存器优化，TC在Pascal上比CUBLAS实现了高达1.5倍的加速。
  • 分析: 证明了TC在标准、中等大小的算子上也能提供性能优势。

• 融合的多层感知机 (MLP3)

  • 结果: MLP3是连续3个低延迟MLP层。Caffe2参考实现会产生多个核函数调用。TC通过JIT将这3层融合成一个单一的CUDA核函数，运行速度比Caffe2（混合使用CUDNN和CUBLAS）快高达3.6倍。
  • 分析: 在这种低延迟场景下，核函数启动开销非常显著。TC的算子融合能力极大地减少了这种开销，从而带来了巨大的性能提升。

总体性能加速比 (图7)

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

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 机器学习框架简介

• TensorFlow
  在核心层面，TensorFlow【索引1，OSDI 2016】通过算子（operators）来描述和执行n维张量上的计算。程序员将算子组合成一个静态的非循环数据流图（DAG），然后指定要运行的计算，最常见的是梯度下降的步骤。图8展示了使用该框架的示例代码。

• Caffe2
  Caffe2【索引37，CoRR 2017】在设计上与TensorFlow结构相似，但API更轻量，且更容易访问计算图中的中间结果。一个主要的设计差异是，Caffe2允许中间对象不仅是张量，还可以是其他不透明数据结构，如图9所示。

• PyTorch
  PyTorch【索引71，PyTorch Website 2017】与TensorFlow和Caffe2在许多抽象上相似，但一个主要区别是它没有静态计算图。PyTorch在指定计算的同时就运行它，而不是先构建图再执行。图10展示了其等效语法。

• MXNet
  MXNet【索引17，CoRR 2015】也共享相似的抽象，并提供声明式（如Caffe2/TensorFlow）和命令式（如PyTorch）两种范式。MXNet的核心工具链中加入了编译器抽象和遍（passes）。其基于图的中间表示（IR）NNVM允许执行图重写变换。NNVM还允许将图的部分切分并委托给TVM【索引16，TVM Annoucement 2017】【索引18，arXiv 2018】，这是一个基于Halide【索引72，PLDI 2013】的声明式核函数和变换规范。

• Lush/SN3
  现代ML框架中许多张量操作思想最早由Lush/SN3系统【索引13，Lush Manual 2002】在1990年代实现。idx-bloop构造迭代一组张量的首个索引。索引算术通过操作与每个索引关联的步幅的辅助函数来实现。

A.2 ML框架接口API

• 嵌入式语言API
  我们通过一个嵌入式语言API将张量推导式暴露给ML框架，类似于正则表达式、SQL查询或图形着色语言的库。代码作为字符串加载到一个管理编译和执行的执行引擎中，如图11所示。然后可以运行此代码，这将调用JIT编译器和自动调优器，如图12所示。

• 核心API
  这些针对Python和PyTorch的API只是一个核心框架无关API（如图13所示）的小型包装器。C++和其他ML框架（如Caffe2）也存在类似的API。

图 11: 构建执行引擎

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

A.3 多面体编译背景

• 迭代域（Iteration Domains）
  多面体框架操作于被称为语句实例的单个循环迭代。这些实例由多维向量空间中的点标识，其坐标对应于外围循环的归纳变量的值。受仿射边界约束，这些点的集合，即迭代域，形成一个凸多面体。在TC的背景下，一个语句实例由出现在该语句中的所有循环迭代变量的值标识。

• 调度（Schedules）
  迭代域的点按照其逻辑执行日期的字典序遍历。通常，这些日期通过一个封闭形式的调度来定义，该调度为每个点关联一个日期。本文主要使用调度树表示法。

• 代码生成（Code Generation）
  给定一个迭代域和一个调度，可以生成一系列循环嵌套来按调度规定的顺序访问所有迭代域点。存在高效的算法来生成FORTRAN、C + OpenMP、CUDA或OpenCL代码【索引6，PACT 2004】【索引14，PLDI 2012】【索引88，TACO 2013】。

• 表示和保持依赖（Dependences）
  多面体框架的一个显著优势在于其对依赖关系的抽象表示，描述为分段仿射关系的并集。这种分析确保调度树的修改保持程序语义。它基于访问关系，该关系将语句实例与其访问的数据（张量元素）相关联。

• 依赖分析（Dependence Analysis）
  当两个语句实例访问同一个张量元素，且至少其中一个是对其进行写操作时，它们之间就存在依赖关系。如果两个调度都以相同的相对顺序对依赖实例对进行排序，那么它们就能赋予程序相同的语义。依赖关系可以通过组合（逆）访问关系并根据调度进行约束来获得。例如，流依赖可以表示为 {(D1 → D2) | (D1 → D2) ⊂ Aread,D2 ◦ A−1write,D1 ∧ S(D1) ≺ S(D2)}。

• 依赖分析的精化
  上述是最简单的基于内存的依赖分析。通过执行精确的数据流分析【索引27，IJPP 1991】，可以将其进一步精化，只考虑实例级的定义-使用对。在isl中，依赖分析还区分了潜在（may）访问和确定（must）访问。我们的依赖关系在构建后可以被注解和进一步定制，例如，通过插入“kill”语句来向编译器表明某些张量元素在计算的不同部分之间不被重用。

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

• 数组瓦块计算
  内存提升需要特定的数组足迹计算和启发式方法来分组访问重叠区域的数组引用。描述块内访问的数组元素的关系是通过对访问关系的域应用分块和映射后的调度，然后投影掉未映射到块的调度维度来计算的。从这个关系中，我们寻找一个参数化偏移量，使得关系的像（image）是常数。如果可以找到这样的偏移量，该数组部分就会被提升到共享内存。

• 数组引用分组与同步
  在提升过程中，数组的一部分被临时复制到缓冲区；修改原始数组可能导致数据不一致。因此，对数组重叠部分的写访问应该被一起提升（或不提升）。我们基于块级访问关系进行数组引用分组。如果块内访问的元素重叠且至少有一个是写操作，则数组引用被分组。同步点被插入在从全局内存读取数据之后、写回全局内存之前，以及写引用和任何其他重叠引用之前。

A.4 TC示例的详细结果

• 转置矩阵乘法 (tmm)
  基线映射选项如下，自动调优器在 Maxwell 上为 (M, K, N) = (128, 32, 256) 找到了一个不同的、性能更好的配置。生成的CUDA核函数展示了数据从全局内存加载到共享内存，然后在私有寄存器中进行计算，最后写回全局内存的模式。

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);

• 转置批处理矩阵乘法 (tbmm)
  基线调优策略中，B 代表小批量大小，映射到CUDA处理器网格。对于 (B, N, M, K) = (500, 26, 72, 26)，自动调优器在 Maxwell 上找到了一个更优的配置，尽管生成的CUDA代码存在一些可以进一步优化的模式，但性能已优于CUBLAS。

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

• 分组卷积
  我们使用了一个显式的5-D TC版本来表示分组卷积，以避免非仿射索引和语义混淆。基线选项根据 W 的大小动态调整线程映射。在 Maxwell 上，对于 (W, H) = (7, 7)，自动调优器找到了一个不理想但性能远超CUDNN的解决方案，该方案通过过度配置线程来积极地将数据从全局内存加载到共享内存。

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)
    // ...

• 生产模型
  我们为生产模型中的每个部分（2LUT、C3、MLP1、MLP3）手动指定了基线映射选项。例如，对于MLP3，我们使用了最大融合策略，并将计算映射到128个线程和128个块。TC的灵活性使我们能够快速实验手动将语句分组到不同的TC中，从而轻松地将模型分解为更简单的函数进行优化。

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

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