MIT Taco 项目：自动生成张量计算的优化代码，深度学习加速效果提高 100 倍

2017 年 11 月 3 日 AI研习社 岑大师

我们生活在大数据的时代，但在实际应用中，大多数数据是 “稀疏的”。例如，如果用一个庞大的表格表示亚马逊所有客户与其所有产品的对应映射关系，购买某个产品以 “1” 表示，未购买以 “0” 表示，这张表的大部分将会是 0。

使用稀疏数据进行分析的算法最终做了大量的加法和乘法，而这大部分计算是无效的。通常，程序员通过编写自定义代码来优化和避免零条目，但这种代码通常编写起来复杂，而且通常适用范围狭窄。

AI研习社发现，在 ACM 的系统、程序、语言和应用会议（SPLASH）上，麻省理工学院、法国替代能源和原子能委员会的研究人员和 Adobe 研究院的研究者们最近提出了一个针对稀疏数据自动生成优化代码的系统。相比起现有的非优化软件包，该系统可实现 100 倍的加速，其性能与特定稀疏数据操作的精细手工优化代码相当，但程序员的工作要少得多。

该系统的名字叫 Taco，即 “张量代数编译器” 的缩写。在计算机科学上，类似亚马逊的表格这样的数据结构被称为 “矩阵”，张量只是一个矩阵的更高维的模拟。如果亚马逊表格还将客户和产品与亚马逊网站上的客户的产品评级以及产品评论中使用的字词进行了对比，那么这可以用四维张量来表示。

Taco 的杀手锏：内核融合

对张量的数学运算对于大数据分析和机器学习都是至关重要的问题。自从爱因斯坦的时代以来，它一直是科学研究的主要内容。传统上，为了处理张量计算，数学软件将张量运算分解为不同组成部分进行计算，例如如果需要计算两个张量相乘加第三个向量，则软件将在前两个张量上运行其标准张量乘法程序，存储结果，然后再运行其标准张量加法。

举例来说，目前我们常用的深度学习框架，如 TensorFlow、PyTorch 等都会将一个深度学习的模型转换为一个由基本计算符组成的数据流图，再有下层的计算引擎一次调度执行这些节点对应的内核函数（对于数据图的动态优化参见AI研习社此前文章《紧跟未来深度学习框架需求，TensorFlow 推出 Eager Execution》）。好的深度学习框架中会定义成百上千个 Operator，这些 Operator 定义了张量的加、减、乘、除和矩阵乘法等等，因此，在深度学习训练中，这些节点在 GPU 上的执行会转变成数千次 GPU 上的内核执行，从而使得张量的计算更加灵活。而对于这些内核函数的频繁调用也成为了影响深度学习框架的重要因素。

然而，在需要进行频繁计算的大数据时代，像目前深度学习框架虽然能同时进行数千次运算，大仍然耗时太多。对深度学习框架进一步优化、解决这些性能问题的一个直接方法就是 Taco 所采用的内核融合方法 (Kernel Fusion)，即将一个计算图中的节点所对应的内核函数融合成一个函数，这样整个数据流图纸需要通过一次函数调用就可以完成所有计算，从而将对平台的调度和内核启动的开销降到最低。

该论文的第一作者、MIT 博士生 Fredrik Kjolstad 解释说，Taco 采用对于大规模数据集的高效运行，张量运算的每一个序列都需要自己的 “内核” 或计算模板。如果在一个内核中执行操作，用户可以一次完成所有操作而不需将输出存放在内存中，从而极大加快计算速度。

此外，通过合理地设计不同内核函数的输入输出数据的放置（AI研习社注：例如使用 GPU 上的共享内存或寄存器），可以极大地提高数据传输效率，从而提升整体计算性能。许多张量运算涉及从一个张量与另一个张量相乘的条目，用于处理大型稀疏矩阵的程序可能浪费大量的时间来添加和乘以零。通过手工优化代码可以识别稀疏张量中的零条目，在运算中做到只对非零条目进行计算或者省略对零条目的计算，可以简化其操作从而加速张量的计算，但这需要程序员做更多的编程工作。

举例来说，有两个矩阵（即只有两个维度的张量）需要进行相乘，如果矩阵所有位置都是有效信息，那么对应的代码只有 12 行，但如果矩阵是稀疏的（包含大量 0），相同的操作相同的操作可能需要 100 行甚至更多的代码来跟踪遗漏和错误。

而在 Taco，系统会自动添加所有额外的代码。程序员只需要简单地指定张量的大小，以及张量的类型（完整的或者稀疏的），以及文件的位置，从它应该导入它的值。对于两个张量上的任何给定操作，Taco 会首先建立一个分层映射，指出来自两个张量的哪些配对条目是非零的，然后将每个张量的条目与零配对，并在这过程中丢弃所有的零对。

此外，Taco 还使用有效的索引方案来存储稀疏张量的非零值。以前文提到的 Amazon 所有客户和库存对应的表格为例，其原始图表需要的存储容量是目前 Google 服务器存储容量的 10 倍，而使用 Taco 的压缩方案只需要 13GB，也就是说，一个智能手机就可以存储这一信息。

俄亥俄州立大学计算机科学和工程教授 Saday Sadayappan 称，过去二十年来，许多研究小组试图解决稀疏矩阵计算的编译器优化和代码生成问题，而 Fredrik 和其导师 Saman 的研究是这个对这个长期存在的开放性问题的一个重要突破。” 他们的编译器可以自动生成非常有效的代码，让开发者可以以非常简单和方便的高级符号来指定非常复杂的稀疏矩阵或张量计算。“ Saday 说。 “从展示的结果看，由编译器自动生成的代码与手动编写的代码效果相当甚至更好，这是最近在编译器优化领域最令人振奋的进步之一。“

关于 Taco 编译器的更多信息可阅读相关论文《The Tensor Algebra Compiler》。

关注AI研习社，回复【20171104】获取论文原文PDF。