Facebook发布张量理解库，自动编译高性能机器学习核心

Facebook AI Research今天发布了张量理解（Tensor Comprehension），这是一个C ++库，也是一种数学语言，它能够自动、按需地及时编译出机器学习所需的高性能代码。

简单来说，就是将（部分）人类能读懂的高级编程语言代码，转换成CUDA等专为机器优化的代码。

Facebook在博客中说，Tensor Comprehension有助于弥合研究人员与工程师之间的沟通鸿沟。研究人员往往专注于数学运算，而工程师专注于在多种硬件后端上运行大规模模型的实际需求。

以前，构建高性能机器学习层通常需要几天到几周的工作，分为两个环节：

一、研究人员在numpy级别的抽象中写一个新层，与PyTorch等深度学习库链接起来，进行小规模测试。要把这些用来验证想法的代码用到大规模实验上，需要加快一个数量级。

二、工程师根据这个层为GPU和CPU写高效代码。

这个过程中需要的工程师，得是高性能计算专家，这类人才数量非常有限。他们需要了解情境、制定策略、编写并调试代码，同时，这个将代码移到后端的工作还涉及一些杂活，比如冗长的参数检查、添加样板集成代码等。

因此，过去几年里，深度学习界在生成GPU和CPU尚运行的高性能代码时，越来越依赖CuBLAS、MLK、CuDNN等代码库。

当想要试验的新想法偏离了这些代码库提供的基元时，会涉及到大量的工程工作，研究人员会被吓得止步不前。

Facebook认为，简化从想法到高性能代码的过程，把它从几天、几周缩短到几分钟有着重大的实践价值。他们希望研究人员从数学的角度描述自己的想法之后，通过Tensor Comprehension的自动编译和调整，就能生成性能良好的专用代码。

这一版本的Tensor Comprehension包含：

• 能用简单语法表达一系列机器学习想法的数学符号；
• 基于Halide IR的，表达这些数学符号的C++前端；
• 一个基于整数集库（ISL）的多面即时（JIT）编译器；
• 一个基于进化搜索的多线程、多GPU自动调节器。

Tensor Comprehension将高性能图像处理领域的流行语言Halide的编译器作为一个库，建立在Halide的中间表示（intermediate representation，简称IR）基础上，和多面编译（polyhedral compilation）技术结合起来。让用户可以用类似的高级语法编写层，但不用说明它如何运行。

通过Halide IR和polyhedral compilation的结合，Tensor Comprehension通过委托内存管理和同步功能自动合成CUDA核。这种翻译针对特定的运算符融合、快速本地内存、快速缩减和特定尺寸的JIT专门化进行优化。

Facebook在博客中表示，因为这一工作流不尝试拥有或优化内存管理，因此可以轻松高效地集成到任何ML框架和能调用C++函数的语言中。

不同于传统的编译器技术和库方法，多面编译让Tensor Comprehension能为每个新网络按需调度单个张量元素的计算。

在CUDA层面，它结合了仿射循环转换，融合/分裂和自动并行处理，同时确保数据在存储器层次结构中正确移动。

图中的数字表示张量元素最初计算的顺序，箭头表示它们之间的依赖关系，数字旋转对应着循环交换，让深层运算符融合成为可能。

Facebook还为推动搜索过程提供了一个集成的多线程、多GPU自动调整库，用进化搜索来生成和评估数千种实现方案，并选择性能最佳的方案。调用Tensor Comprehension上的调节（tune）功能就能实时看到性能的提升。

在性能方面，Tensor Comprehension在部分案例中可以匹敌甚至超越集成了手动调节代码库的机器学习框架。这主要依靠让代码生成策略适应特定问题大小的能力。

下图可以看出，Tensor Comprehensions自动生成核的性能，比Caffe2和ATen用供应商提供的CuDNN等库生成的核有所提升。

对于英伟达、英特尔等硬件厂商所提供的快速库来说，Tensor Comprehension是一种补充，也可以和CuDNN、MKL、NNPack等一起使用。

稍后，Facebook还将发布Tensor Comprehension的PyTorch集成。

Tensor Comprehensions是Facebook，法国国家信息与自动化研究所（Inria），苏黎世联邦理工学院和麻省理工学院的合作项目。

论文：

Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, Albert Cohen https://arxiv.org/abs/1802.04730

代码：

https://github.com/facebookresearch/TensorComprehensions

文档：

https://facebookresearch.github.io/TensorComprehensions/

博客：

https://research.fb.com/announcing-tensor-comprehensions/