作者｜Chip Huyen

翻译｜胡燕君、贾川、程浩源

我承认，在大学的编译器课上哭了，后来我挑选成为一名机器学习工程师，认为再也不用被编译器烦扰。

可是，当我逐步了解ML模型怎么投入生产运用，关于编译器的问题不断涌现。在许多用例中，尤其是用边际设备运转ML模型时，模型的成功与否仍然取决于运转它的硬件（hardwarelottery.github.io/ ）。因而，了解模型的编译和优化，以及模型在不同硬件加快器上的运转十分重要。

抱负状况下，咱们不需求特别关注编译器，一切都会“正常作业”。可是，要完结这个方针还有很长的路要走。

现在，越来越多公司期望用边际设备运转ML模型，运转ML模型的硬件越来越多，为了让ML模型能在硬件加快器上更好地运转，也诞生了越来越多编译器——例如MLIR dialects、TVM、XLA、 PyTorch Glow、cuDNN等。依据PyTorch创始人Soumith Chintala的说法，跟着ML的运用逐步成熟，公司之间的竞赛将转向谁能更好地编译和优化模型。

图1: ML的下一个战场是编译器（Soumith Chintala，Venture Beat 2020）

了解编译器的作业原理能够帮你选对编译器，让模型在所选硬件上的运转效果达到最佳，也能帮你诊断模型功用问题并加快模型运转速度。

本文对ML编译器进行了浅显介绍。ML编译器始于边际计算的兴起，这使编译器不再是体系工程师的专属，而是整体ML从业者关心的范畴。假如你已十分了解用边际设备运转ML的重要性，能够跳过下文榜首部分。

然后我将谈及在边际设备布置ML模型的两个主要问题：兼容性和功用，并阐明编译器怎么解决这些问题，以及它的作业原理。本文最终还将供给关于怎么通过几行代码来显著进步ML模型速度的参阅资料。

1

云核算 VS 边际核算

想象一下，你现已练习出一个超卓的ML模型，其精度远超你的期望，然后你刻不容缓地布置这个模型，来让用户运用。

最简略的办法是将模型打包并通过保管云服务（如AWS或GCP）进行布置，这也是许多公司在最初运用ML时的做法。云服务在这方面贡献巨大，使得企业能够轻松将ML模型投入生产。

可是，云布置也有许多缺陷。首要是本钱。ML模型需求进行很多核算，而核算的本钱很高。早在2018年，Pinterest、Infor、Intuit等大公司每年在云服务上的开销就已超数亿美元。而关于中小型公司来说，这个数字每年或许在5-200万美元之间。草创公司在云服务方面稍有不慎，或许就会导致破产。

图2: AWS运用量激增，企业云账单金额飙升（图源The Information, 2018)

跟着云核算费用不断攀升，越来越多公司试图将核算转移到消费设备（边际设备）。在边际设备上完结的核算越多，对云服务的需求就越少，他们要付出的费用就越少。

除了能下降本钱，边际核算还有许多优势。首要，边际核算的可运转范围更广。当模型位于公有云上时，向云端发送并接收数据需求依靠安稳的网络连接。但边际核算能够使模型在没有网络连接或连接不安稳的状况下持续运转，比如在农村地区或发展中国家。

其次，边际核算能够削减网络推迟的困扰。假如有必要运用网络传输数据（将数据发送到云端的模型进行猜测，然后将猜测发送回用户），那么某些用例或许就无法完结。在许多状况下，网络推迟比推理推迟更严峻。例如，你或许能够将ResNet50模型的推理推迟从30毫秒下降到20毫秒，但其网络推迟或许会高达数秒，具体状况取决于你的地点地网络。

第三，边际核算能够更好地保护灵敏的用户数据。云核算ML模型意味着或许要通过网络来发送用户数据，增加了数据被阻拦的危险。一起，云核算还将许多用户的数据存储在同一方位，这意味着一旦发生数据走漏就会影响许多人。据《安全》杂志2020年报导，近80%的公司在曩昔18个月内曾遭遇过云数据走漏(*
www.securitymagazine.com/articles/92… )。边际核算使得企业能够更安全地传输和存储用户数据，防止违反GDPR等数据保护条例。

2

编译：兼容性

相关于云核算，边际核算具有许多优势，因而许多企业正竞相开发针对不同ML用例优化的边际设备。谷歌、苹果、特斯拉等闻名巨子都宣布自研芯片。与此一起，不少ML硬件草创公司也融资数十亿美元来开发更好的AI芯片。

图3: AI硬件草创公司的不完全统计（融资信息来自CrunchBase）

运转ML模型的硬件产品变多了，于是就呈现一个问题：怎么让运用恣意结构构建的模型能够在恣意硬件上运转？

一个结构要在某种硬件上运转，它有必要得到硬件商的支撑。例如，虽然Google早在2018年2月就已公开发布TPU，但直到2020年9月，TPU才支撑PyTorch。在此之前，假如想运用TPU，就有必要要运用Google的TensorFlow或JAX。

要使一种硬件（或平台）支撑某一结构需求消耗很多时刻和作业量。将ML作业负载映射到硬件需求了解并运用硬件的根底架构。可是，有一个根底性问题有必要克服：不同的硬件类型有不同的内存布局和核算原语，如下图所示：

图4: 不同硬件后端的核算原语和内存布局（图源：arxiv.org/abs/1802.04…

例如，早年，CPU的核算原语是数字（标量），GPU的核算原语是一维向量，而TPU的核算原语是二维向量（张量）。可是，现在许多CPU具有向量指令，而一些GPU具有二维张量中心。给定一个256张图画 x 3 通道 x 224 W x 224 H的batch，假如要对这个batch履行卷积算子，一维向量核算原语和二维向量核算原语将有很大不同。相同，还要考虑不同的L1、L2和L3布局和缓冲区巨细，如此方能有效地运用内存。

正因如此，结构开发人员更倾向于只支撑少量服务器级别的硬件（如GPU），而硬件商也倾向于只向少量结构供给自己的内核库（例如，英特尔的OpenVino东西库仅支撑Caffe、TensorFlow、MXNet、Kaldi和ONNX。NVIDIA自己则有CUDA和cuDNN）。将ML模型布置到新硬件（例如手机、嵌入式设备、FPGA 和 ASIC）需求消耗很多的人力。

图5: 怎么在恣意硬件后端运转运用恣意结构构建的模型？图上每个空格都需求一种编译器？

中心表明 (IR)

与其针对每种新的硬件类型和设备装备新的编译器和库，何不创建一个中心前言来桥接结构和平台？结构开发人员将不再需求支撑每种类型的硬件，而是只需将他们的结构代码“翻译”成这种中心前言。这样，硬件商就只需求支撑一个中心结构。

图6: 中心表明（IR）作为中心前言

这种中心前言称为中心表明（IR）。IR是编译器作业的中心。编译器在模型原始代码的根底上生成一系列高档和初级中心表明，然后生成硬件原生代码以在特定平台上运转模型。

编译器一般运用代码生成器（codegen）依据IR生成机器原生代码。ML编译器中最常用的代码生成器是Vikram Adve和Chris Lattner开发的LLVM，LLVM改变了咱们对体系工程的了解。TensorFlow XLA、NVIDIA CUDA编译器 (NVCC)、MLIR（用于构建其他编译器的元编译器）和TVM都在运用LLVM。

生成代码的进程也称为“降级（lowering）”，由于是将高档的结构代码“下降”为初级的硬件原生代码。精确地说，这个进程不能称作“翻译”，由于两种代码之间不是1对1的映射联系。

高档IR一般是ML模型的核算图。关于了解TensorFlow的人来说，这儿的核算图相似TensorFlow 1.0中的核算图，那时TensorFlow还没有切换到Eager Execution方式。TensorFlow 1.0在运转模型之前会先构建模型核算图，核算图可让TensorFlow了解模型，然后优化运转时。

图7: 高档IR和初级IR

高档IR一般与硬件无关（不管在什么硬件上运转），而初级IR一般与结构无关（不管模型是用什么结构构建的）。

3

优化：功用

完结代码“降级”后，在所选硬件上运转模型时或许还会遇到功用问题。Codegen十分拿手将IR降级为机器代码，但受方针硬件后端的影响，机器代码的运转效果或许不够好。生成的代码或许无法充分运用数据局部性和硬件缓存，或许无法运用能够加快代码的高档功用，例如向量操作或并行操作。

典型的ML作业流程需求用到许多结构和库。例如，用 Pandas/Dask/Ray 从数据中提取特征；用NumPy履行向量化；用LightGBM等树模型来生成特征，然后用由不同结构（如sklearn、TensorFlow或Transformers）构建的各种模型来进行猜测。

虽然这些结构内部的单个功用或许得到优化，但跨结构的优化几乎不存在。一种简略的办法是把数据搬运到不同的功用上分别核算，但这会导致整个作业流的速度下降一个数量级。斯坦福 DAWN试验室的研讨人员的一项研讨发现：与手动优化代码相比，运用NumPy、Pandas和TensorFlow的典型ML作业负载在一个线程中的运转速度要慢23倍(www.vldb.org/pvldb/vol11… )。

在生产中，数据科学家/ML工程师常用pip来装置他们作业所需的包，在开发环境中运转状况良好，于是他们将模型布置到生产环境。当他们在生产中遇到功用问题时，他们的地点公司一般会聘请优化工程师来依据运转的硬件优化模型。

图8: Cruise优化工程师的作业描绘

图9: Mythic优化工程师的职位描绘

可是优化工程师十分稀缺，薪资要求也高，由于这份作业既要懂机器学习，也要懂硬件架构。另一个办法是选用优化编译器（Optimizing Compiler），即能够优化代码的编译器，也能够主动优化模型。在将ML模型代码降级为机器代码的进程中，编译器能够剖析模型核算图以及其中包含的算子（如卷积算子、循环算子、穿插熵等），然后设法加快核算速度。

前文总结：编译器能够把ML模型和它赖以运转的硬件桥接起来。优化编译器包含两部分：降级与优化。这两部分有时也能够整合在一起。优化或许发生在从高档IR到初级IR的任何阶段。

• 降级：编译器为模型生成硬件原生代码，让模型能够在特定硬件上运转。

• 优化：编译器依据硬件环境优化模型。

4

怎么优化ML模型

优化ML模型的办法有两种：局部优化和大局优化。局部优化是指优化模型的某一个或某一组算子；大局优化是指端到端优化整个核算图。

现在已有一些规范的局部优化办法，大部分是通过进步并行度或削减芯片内存访问来加快模型。以下是四种常见的办法。

• 向量化：给定一个循环或嵌套循环，但单次不止履行一个元素，而是运用硬件原语履行在内存中接连的多个元素。
• 并行：给定一个输入数组（或许n维数组），将其分割为多个不同的独立作业块，分别对每个作业块履行操作。
• 循环分块（Loop Tiling） ：修正循环中的数据访问次序，然后更好地运用硬件的内存布局和缓存。这种优化办法对硬件的依靠极高。合适CPU的访问方式未必合适GPU。详情可参阅Colfax Research供给的可视化解说(colfaxresearch.com/how-series/… )。
• 算子交融：将多个算子交融为一个算子，能够防止冗余的内存访问。例如，对同一数组履行两个操作需求两次循环，但通过交融后就只需求一次循环。详情可参阅Matthias Boehm供给的比如。(mboehm7.github.io/teaching/ss…)

图10: 来自Colfax Research的可视化解说

图11: 来自Matthias Boehm的算子交融比如

Weld编译优化器的缔造者Shoumik Palkar表明，在必定的条件下，这些规范局部优化办法能够将模型运转速度提高3倍左右。(www.youtube.com/watch?v=JbT… )

图12: 卷积神经网络核算图的笔直交融&水平交融（图源TensorRT）

完结更大的速度提高需求运用到核算图中更高层次的结构。例如，给定一个卷积神经网络及其核算图，可进行笔直或水平方向的交融，然后削减内存访问，加快模型运转速度。详情可参阅NVIDIA的TensorRT 团队供给的可视化解说(developer.nvidia.com/tensorrt )。

5

人工规划 VS 依据ML的编译器

人工规划优化规矩

正如上文的笔直/水平交融比如所示，履行核算图的办法许多。例如，给定3个算子A、B、C，能够交融A与B，B与C，也可交融A、B、C。

一般，结构和硬件商的优化工程师能依据经历探究出履行模型核算图的最佳办法。例如，NVIDIA或许会安排一名工程师甚至一个工程师团队来专门研讨怎么让ResNet-50模型在NVIDIA的DGX A100服务器上运转得更快（所以，不该太垂青MLPerf基准测试结果(mlcommons.org/en/inferenc… )。由于即使一种常见模型在某种硬件上能够运转得相当快，也并不代表任何模型运用该种硬件都能取得相同的高速度。这种模型或许只是通过过度优化罢了）。

人工规划优化规矩有一些不足之处。首要，得出的优化规矩或许不是最优的。谁也不能确保工程师想出来的优化办法便是最佳方案。

其次，人工规划出的优化规矩不具备自适应性。假如要针对新的结构或新的硬件架构优化模型，还得重新消耗很多人力。

何况，模型优化取决于核算图中的算子，这又增加了杂乱性。优化卷积神经网络不同于优化递归神经网络，后者又不同于优化Transformer模型。NVIDIA和Google专心于针对自家硬件优化ResNet和BERT等常见模型。但假如ML研讨人员又创造出新的模型架构呢？那他们就得先自行优化这种新模型，证明它具备高功用，才能让硬件商选用并持续优化这种模型。

用ML来加快ML模型

咱们的方针是找到履行核算图的最快办法。那么能不能将一切或许的办法都测验一遍，记录每种办法的运转时刻，然后找到时刻最短的一种呢？

能够。但问题是，潜在的办法及其组合实在太多，难以一一尽头。但假如借助ML呢？

• 用ML能够缩小查找空间（即一切或许办法的调集），不用测验一切办法。

• 用ML还能够猜测每种办法的所需用时，不用消耗时刻等候核算图完结履行。

可是，要预估某种办法履行核算图的所需用时十分困难，由于这需求对核算图作出很多假设。现在的技能只能做到关注核算图的一小部分。

假如你曾在GPU上运转PyTorch，你应该见过下列设置：

torch.backends.cudnn.benchmark=True

当设置为True时，就启用了cuDNN autotune。cuDNN autotune会在一组预先设定的履行卷积算子的办法中探究最快的一种。假如每次迭代的卷积神经网络形状都共同，启用cuDNN autotune能够大大进步效率。除了初次运转卷积算子时速度较慢以外（由于cuDNN autotune需求花时刻探究最快的办法），在后续运转中，cuDNN都能够运用autotune的缓存结果直接挑选最快的装备。

虽然cuDNN autotune能够进步效率，但它只适用于卷积算子，而且应该只在PyTorch和MXNet上能用。更通用的解决方案是autoTVM，它是开源编译器堆栈TVM的一部分。autoTVM可针对子图寻觅最佳履行办法，而不只是针对一个算子，所以它的查找空间要杂乱得多。autoTVM的作业原理也很杂乱，但能够总结如下：

1. 将核算图分割为多个子图。

2. 猜测每个子图的巨细。

3. 分配时刻为每个子图寻觅最佳履行办法。

4. 将每个子图的最佳履行办法组合起来，履行整个核算图。

autoTVM能核算每种办法的实践运转时刻，然后收集实在数据用以练习本钱模型，令本钱模型猜测未来某种办法的所需时刻。这种做法的优点是，由于本钱模型是依据运转时发生的数据练习的，所以它能够适应任何硬件，坏处是等候本钱模型练习完善需求更长时刻。

图13: 依据ML的TVM带来的提速；依据ML的TVM大约需求70次迭代才能逾越cuDNN的功用；试验数据来自陈天奇等人。

TVM之类的编译器具有自适应性、灵活性，能协助用户方便地测验新硬件，例如苹果公司在2020年11月发布的M1芯片。

M1是一个依据ARM的片上体系（SoC），ARM架构咱们或多或少都有所了解，可是M1的ARM完结仍然有许多新颖的当地，需求很多优化才能使各种ML模型在M1芯片上快速运转。

在M1芯片发布一个月后，OctoML（译注：TVM团队2019年成立的公司，专心于ML模型的布置）表明，autoTVM 进行的优化比苹果公司的Core ML团队人工规划的优化快了近 30%(venturebeat.com/2020/12/16/… )。当然，跟着M1的成熟和人工优化的深入发展，主动优化将很难逾越人工优化，但体系工程师仍能够运用autoTVM等主动优化东西来加快优化。

虽然主动优化的效果杰出，但仍然存在一个问题：TVM的速度有或许很慢。需求试遍一切的或许办法，才能确认最佳优化途径，这一进程或许需求消耗数小时，杂乱的ML模型甚至或许需求数天。但这是一劳永逸的进程，优化查找的结果能够被缓存，然后用于优化现有模型，还可认为未来的优化供给根底。

当你针对一种硬件后端优化了一次模型之后，该模型就能够在运用同一硬件后端的多种设备上运转。当你准备好用于生产的模型，并且选定方针硬件来运转推理时，这种优化办法十分适用。

6

不同的编译器

现在最广泛运用的编译器类型，是由主流结构和硬件商开发的、针对特定结构和硬件组合的特定范畴编译器。果然如此，最常用的编译器是由最大的结构和硬件商开发的。

• NVCC（NVIDIA CUDA编译器）：闭源；仅适用于 CUDA。

• XLA（加快线性代数编译器，由Google开发）：已作为TensorFlow库的一部分开源；XLA原本计划用于加快TensorFlow模型，但已被运用到JAX。

• PyTorch Glow（Facebook）：已作为PyTorch库的一部分开源；虽然在TPU上启用了XLA+PyTorch，但关于其他硬件还是依靠PyTorch Glow。

第三方编译器一般雄心勃勃（假如某编译器开发商认为自己针对GPU进行的模型优化比NVIDIA做得还要好，这必然需求足够的自傲）。第三方编译器的存在很重要，在闻名巨子已针对自家产品深度调整自家编译器的状况下，第三方编译器能够协助小公司推出的新结构、新硬件和新模型时下降开销、进步功用，让小公司也有机会与闻名巨子竞赛。

我心目中最好的第三方编译器是Apache TVM，它适用于多种结构（包含TensorFlow、MXNet、PyTorch、Keras、CNTK）和多种硬件后端（包含CPU、服务器GPU、ARM、x86、 移动GPU和依据FPGA的加快器）。

另一个我看好的项目是MLIR。它最初也是由LLVM的创建者Chris Lattner在Google发起，但现在，该项目隶属LLVM。MLIR并不是真正意义上的编译器，而是一种元编译器，一种能够让用户构建自己的编译器的根底架构。MLIR能够运转多种IR，包含TVM的IR、LLVM的IR和TensorFlow核算图。

WebAssembly (WASM)

WASM是曩昔几年中最令我兴奋的技能趋势之一。它拥有高功用、易于运用，而且它的生态发展飞快。到2021年9月，全球93%的设备都支撑WASM。

咱们一直在评论编译器怎么为模型生成机器原生代码，让模型能够在特定的硬件后端上运转。那能不能生成一种能够在任何硬件后端上运转的代码？

浏览器有一个优点：假如你的模型能够在浏览器上运转，那么这个模型就能够在任何支撑浏览器的设备上运转，包含Macbook、Chromebook、iPhone、Android手机等，且无需关心这些设备运用什么芯片。假如苹果决定从英特尔芯片改用ARM芯片，对你也没有影响。

WASM是一个敞开规范，它让你得以在浏览器上运转可履行程序。运用sklearn、PyTorch、TensorFlow等结构构建模型后，你不用针对特定硬件编译模型，而是能够将模型编译为WASM。然后你会得到一个可履行文件，能够用JavaScript来履行该文件。

由于WASM是在浏览器中运转，所以它的缺陷是速度很慢。虽然WASM现已比JavaScript快得多，但与在设备上运转原生代码相比（例如iOS或Android运用程序），它仍然很慢。Jangda等人的研讨显示，编译为WASM的运用程序，其运转速度比原生运用程序平均慢45%（火狐浏览器）到55%（谷歌浏览器）(www.usenix.org/conference/… )。

有许多编译器能够将代码编译到WASM运转时。例如最盛行的Emscripten（它运用的也是LLVM Codegen），但它只能将C语言和C++编译成WASM；Scailable能够将scikit-learn模型编译为WASM，但它在GitHub上只有十几个Star，并且现已有几个月没更新了（开发方是不是现已不再保护它了？）；TVM应该是现在还能用的仅有一个可将ML模型编译到WASM的编译器(tvm.apache.org/2020/05/14/… )。

温馨提示：假如你计划运用TVM的话，用非官方conda/pip命令能够快速装置（tlcpack.ai/ ），不要看Apache站点上的阐明。由于假如你按照后者操作遇到问题，只能去他们的Discord上找人帮助(discord.gg/8jNs8MkayG )。

7

编译器的未来

考虑模型在不同的硬件后端上的运转细节很有必要，这样有助于进步模型功用。Austin Huang(www.linkedin.com/in/austin-h… )在MLOps Discord上发帖称，只需运用简略的现成东西，例如量化东西、Torchscript、ONNX、TVM等，往往就能够取得两倍的模型加快。

在此我向我们推荐一篇文章（efficientdl.com/faster-deep… ），它列出了多个在GPU上加快PyTorch模型的有用技巧，甚至无需运用编译器。

在模型布置阶段，有必要测验不同的编译器，比较哪一个能够带来最佳的功用提高。你能够并行进行试验。假如一个推理请求取得小幅提速，那么数百万或数十亿推理请求就能够累积成巨大回报。

虽然用于机器学习的编译器现已取得巨大进步，但咱们还需求许多尽力，才能把编译器笼统出来，让广大ML从业者不用再为编译器烦恼。抱负的状况是，ML编译器能够像GCC传统编译器一样。GCC会主动将C语言或C++代码降级为机器代码，让大多数C语言程序员不用关心GCC会生成什么中心表明。

未来，相信ML编译器也能够做到这样，当开发者运用结构创建出核算图方式的ML模型，然后ML编译器能够依据任何方针硬件为模型生成机器原生代码，开发者也不用关心编译器生成什么中心表明。TVM等东西能够协助咱们完结这一未来。

8

称谢

感谢Luke Metz、Chris Hoge、Denise Kutnick、Parimarjan Negi、Ben Schreiber、Tom Gall、Nikhil Thorat、Daniel Smilkov、Jason Knight和Luis Ceze，他们耐心肠回答了我的问题，协助我写出了这篇文章。

（本文经授权后由OneFlow编译发布，原文：huyenchip.com/2021/09/07/…

欢迎下载体验 OneFlow v0.8.0 最新版本：

github.com/Oneflow-Inc…