Facebook发布Tensor Comprehensions，可自动按需生成机器学习代码

来源：Facebook Research

编译：Bot

近日，Facebook AI Research（FAIR）宣布发布Tensor Comprehensions，这是一个C++库，它用数学语言帮助弥合研究人员和工程师之间的沟通差距，能帮他们解决数学运算方面的一些问题，并按需求在各种硬件后端上运行大型模型。Tensor Comprehensions的主要特色在于它跟进了行业对Just-In-Time编译的独特需求，可以自动按需生成机器学习社区所需的高性能代码。

大幅提高生产力

以现有的情况看，如果要创建高性能的机器学习模型神经网络层，研究人员至少需要耗费数天或数周的时间，它主要可分为以下两块内容：

1.研究人员首先用numpy等语言编写新层，再调用PyTorch这样的深度学习库把代码和操作连接起来，最后进行小规模测试。如果要进行大规模测试，他们还需要通过数量级加速检测其代码实现。

2.把新层交到工程师手中，并由他根据GPU、CPU条件编写高效代码：

这些工程师必须是编写高效代码方面的专家，但是这样的人才十分稀缺；

这些工程师不仅要理解任务、制定策略，还要编写和调试代码；

把代码移植到后端会涉及一些普通操作，如冗长的参数检查和添加样板集成代码。

因此，在过去的几年中，深度学习社区已经发展到依靠CuBLAS、MKL和CuDNN等高性能库来获得GPU和CPU上的高性能代码。这在一定程度上提供了便利，但也成了研究规模扩大和水平提高的一个瓶颈。

为了解决这个问题，Facebook推出了Tensor Comprehensions，它能让研究人员用数学符号写出自己的想法，并把符号自动编译和调整成高性能代码，从而将编写机器学习模型的时间从几天或几周缩短为几分钟。

初代Tensor Comprehensions主要有以下几个亮点：

用简单的语法表达一系列ML概念的数学符号；

基于Halide IR的数学符号的C ++前端；

基于整数集库（ISL）的多面体即时（JIT）编译器；

一个基于进化搜索的多线程、多GPU自动调节器。

一些前期工作

Halide是最近在高性能图像处理领域颇受欢迎的一种语言，它用类似的高级函数语法来描述图像处理流程，然后在单独的代码块中明确地将其调度到硬件上，详细指定操作如何平铺、矢量化、并行化和融合。这对于有建筑背景知识的人来说是一种非常高效的语言，但对于大多数ML从业者来说却很难使用。那么，我们能不能把Halide的自动调度引申到编写在GPU上运行的ML代码上呢？

Tensor Comprehensions把Halide编译器作为库。它基于Halide的中间表征（IR）和分析工具，并将它与多面体编译技术进行匹配，也就是说，研究人员并不需要理解怎么写代码就能用高级语法编写图层。

细节

Tensor Comprehensions用Halide和Polyhedral Compilation技术，通过委托内存管理和同步功能自动合成CUDA内核。该编译针对特定尺寸对一般操作符融合、快速本地内存、快速缩减和JIT专业化进行优化。由于我们不尝试拥有或优化内存管理，因此我们的流程可以轻松高效地集成到任何ML框架和任何允许调用C ++函数的语言中。

与经典编译器技术和库方法相反，Polyhedral Compilation技术允许Tensor Comprehensions为每个新网络按需调度单个张量元素的计算。

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

图中的数字表示最初计算张量元素的顺序，箭头表示它们之间的依赖关系。在这个例子中，数字旋转对应于循环交换，这使得深层操作器融合成为可能。

为了推动搜索过程，Facebook还提供了一个集成的多线程、多GPU自动调谐库，它使用Evolutionary Search来生成和评估数千种实现方案，并选择性能最佳的方案。只需调用Tensor Comprehension上的tune功能，研究人员就可以根据表现进行筛选。最好的策略是通过protobuf序列化，并立即或在离线情况下可重用。

在性能方面，尽管开发团队做了不少改进，但在理想情况下，Tensor Comprehensions已经可以匹配，甚至击败当前ML框架与手动调整库集成的性能。这主要通过将代码生成策略适应特定问题来实现。下图是Tensor Comprehensions自动生成的内核与Caffe2和ATen（使用供应商库实现，如CuDNN）中的现有替代品的比较，我们能从中观察到它们各自的性能增益。

随着这一技术日益优化，Tensor Comprehensions将有望填补硬件制造商（如NVIDIA和Intel）编写的快速库，并将与CUDNN、MKL或NNPack等库一起被研究人员所使用。

Tensor Comprehensions地址：github.com/facebookresearch/TensorComprehensions

原文地址：research.fb.com/announcing-tensor-comprehensions/

论文地址：arxiv.org/abs/1802.04730