AI 硬件公司接入框架与 AI 框架算子层级的思考

这个思考其实源于参加 WAIC 上组织的 AI 编译相关的闭门讨论的内容，观点有不少来源于现场讨论的朋友们，因为对这个主题感兴趣，我又结合自己的理解做了一些梳理。
之所以关注 AI 框架算子层级的问题，是因为自己最近在关注AI硬件行业，也就在关注为一款新硬件提供软件支持可能面临的问题。
AI 硬件公司，能成事非常不容易。首先是硬件要流片成功，然后是软件栈建设过得去，再之后还需要商业策略不掉链子（包括市场定位、切入策略、客户解决方案团队建设、量产方案等等）。
这非常不像典型的互联网创业团队的玩法，小快灵就可以 move 起来， fail fast 。硬件行业，有其很强的物理规律不能简单靠加人加班就超越，哪怕是天才型选手的加入，也只能确保这个过程犯错尽可能少，并不能确保大幅缩短执行周期。
而我的背景偏软件，至少希望确保软件这一环不会掉链子。本文将围绕一款 AI 新硬件上面，软件栈的建设，展开思考。

0. AI 硬件的软件栈 · 概览

首先来看看一款 AI 硬件的软件栈应该长什么样子，下图是我的一个理解。在这套技术栈里面。有几个细节可以稍微展开一点：

图作者理解的AI 硬件软件栈概览：底层硬件与上层框架通过中间插件、运行时、编程模型、设备编译器和驱动进行桥接

1. 硬件编译器与编程模型

新硬件的编译工具链与编程模型的建设往往需要时间。因为新硬件上软件系统的建设本质上是一个自底向上的工作，所以抽象层次越高的工作其实现复杂度往往更大，交付时间会靠后。

图从高层语言到汇编语言、机器码，汇编与机器码之间需要汇编器，机器码到硬件之间需要指令集

但上层 AI 应用的验证需求又是客观存在的，所以就会存在一个从框架插件或者算子库，甚至编程模型这个层级，直接接入到汇编器（assembler）这一层的快捷路径（NeuralTalk评：再怎么快捷，这中间应的编译工具链是少不了的）。

图 TPU ISA 主要涵盖数据操作（主内存读写）和计算操作（/乘法/卷积/激活计算）

因为 assembler 通常做的是一个字面汇编描述到二进制码的翻译，外加一些可轻可重的优化，更容易迭代式交付给上层使用。当然了，如果有可能的话，把编译工具链和编程模型，也可拆解成可以阶段式交付的 milestone ，尽可能让上层开发可以提前借助于高级语言来解放生产力会更有帮助。

举一个非常小的例子，在汇编码里做 pass by value/pass by reference 的区分，和在 cpp 里，前者还是需要消耗一些额外的心智负担，也会影响一些开发调试的效率，这属于实际操作中的考虑了（NeuralTalk评：这中间往往有调试用的 Profiler 、模拟器等等工具，参考OpenVINO）。

2. 框架插件接入硬件

这里对接AI框架的层级我定义为 "Framework plugins"，背后的思考是我在分析一些硬件公司的方案的时候，能够看到几乎所有新硬件厂商都有一个倾向——就是假设或者希望AI框架能够导出一个完整的静态图（训练和推理），用作自己软件栈的输入，来完成训练和推理能力的支持。

图特定领域架构（DSAs）设计规范。可以看出对数据资源从切分粒度到编程语言，并行方式等

这种作法存在的原因其实比较容易理解，新硬件为了获得相较于NV更好的性能表现，通常会添加 DSA 的支持。

DSA（特定领域架构）也是体系结构和摩尔定律发展到如今的必然选择，比如AI芯片、自动驾驶芯片等都属于 DSA 的范畴。这也是未来一段时间商业硬件设计工作的重要方向就是 DSA 。
随着摩尔定律的发展，硬件已经支持了深度的存储层次、很宽的 SIMD 部件、级数很深的流水线、分支预测、乱序执行、投机预取、多线程和多处理；但是仍然有部分软件的行为是不能被硬件较好的加速的，所以要研究 DSA 。
DSA 的目标就是通过抽取软件的行为来发现不能被当前架构支持较好加速的部分，进行新的架构设计从而提高软件性能。

而 DSA 会加剧计算能力和访存以及计算 kernel dispatch 到加速器的差距，所以如果能全图都交到硬件加速器上处理，从充分发挥新硬件性能是最理想的了。

图硬件如何加入到TensorFlow中

不过在 eager mode 广为用户接受的情况下，目前主流的 AI 框架想做到所有模型都导出一个干净的静态图我觉得并不容易（NeuralTalk评：这个角度主要是针对数据中心的 AI 加速器），所以我其实更希望的是新硬件把对接 AI 框架的模块做成一个插件，可以集成进不同框架的后：

一方面可以更 native 地配合 AI 框架的 eager mode ，在 python host 端和加速器端做灵活的交互；
另一方面，当用户在eager mode下完成交互开发，想追求更高性能的时候，可以再通过一些扩展的API来将模型可静态导出的部分尽可能导出成一张或几张静态图。相关的工作比较早也有人进行过尝试，比如MxNet Gluon，交给加速器来执行，至于一些不易导出静态图的部分，仍然通过python与加速器交互的方式进行fallback执行。这样可以在性能和通用性上获得更好的平衡。当然，这里有一个细节，如果引入了过多的 Python 侧与加速器的交互，很可能会把加速器的速度明显拖慢，甚至相比 NVIDIA 不再有优势。

所以，从加速器角度，会期望不断扩大用户模型能够交给加速器上的子图比例，以及通过类似 cuda-Graph 或 pipeline 执行的方式来隐藏这种开销，但通过这种插件的方式至少保证对接建模需求的可用性，不会因为用户模型一调整，就跑挂了，因为模型训练环节，对于交互式开发调试效率是非常看重的。

3. 兼容灵活性与性能

在框架插件这个层面，想在对接主流 AI 框架时做到既保证在经典代表性模型上的极致性能，比如MLPerf里的模型，同时兼顾对算法用户灵活开发的模型变种的适配程度，就需要考虑 AI 框架算子覆盖度的问题了。

图 TensorFlow 计算图

因为算法用户通常是使用 AI 框架提供的算子来完成模型搭建，以及少量的自定义算子。

TensorFlow现在有多少个算子呢，在最新的master code base里已经接近2000个了。
PyTorch也在相近的数量规模。

这里更麻烦的是算子的数量会随着建模需求的演化还在不断地扩散增加。再加上AI框架版本的快速演化，这就让新硬件对接适配的时候可能出现疲于奔命的现象。这里不讨论业务场景里 AI 框架版本碎片化的问题，那也是一个考虑框架插件化接入方案的原因。

4. 算子级别的接入？

对接 AI 框架的时候，一定需要对接到算子层级么？我觉得未必。原理上，存在两条对接 AI 框架的路径：

直接对接到 AI 框架算子粒度；
对接到 AI 框架算子粒度之下一个更细的原子IR粒度。比如HLO/Relay/ONNX/TorchScript，其中HLO只有200条左右指令)。从

硬件厂商角度来说，最希望的还是能够有一套具备“封闭且可完备描述模型行为的”原子IR粒度可供对接。从用户使用角度，则更习惯于在算子层级这个粗粒度来进行模型描述，而不是使用原子 IR 描述模型。通过原子 IR （Primitive IR）来装配组装出上层的 AI 框架算子，来桥接上层用户易用性和下层硬件后端对接性能 / 工作量的平衡，类似下图所示。

图通过原子IR来装配框架算子

理想很美好，现实很骨感。虽然从硬件厂商角度来说，能够通过直接对接一套低层级中间表示（Primitive IR）支持不同 AI 框架是符合硬件厂商需要的，但对于 AI 框架来说却未必。

5. 软硬件全栈设计与 IR 定义

作为对比，这里简述一下 XLA 的工作原理：

图 XLA 的编译过程：两层 HLO 中间有一层设备无关的优化分析PASS，最后是设备相关的优化分析以及代码生成

XLA 的输入语言称为“HLO IR”或仅为“HLO”（高级优化器）。运算语义页面中介绍了 HLO 的语义。可以将 HLO 简单理解为编译器 IR。XLA 接受在 HLO 中定义的计算图（“计算”）并将其编译为适用于各种架构的机器指令。XLA 采用模块化设计，可以轻松融入其他后端以针对某些新颖的硬件架构。TensorFlow 源代码树中包含适用于 x64 和 ARM64 架构的 CPU 后端，以及 NVIDIA GPU 后端。
XLA 提供了多种与目标无关的优化和分析过程（例如 CSE）、与目标无关的运算融合，以及用于为计算分配运行时内存的缓冲区分析。完成与目标无关的步骤之后，XLA 会将 HLO 计算发送到后端。后端可以执行进一步的 HLO 级优化，而此时将考虑目标特定的信息和需求。例如，XLA GPU 后端可以执行特别有利于 GPU 编程模型的运算融合，并确定如何将计算划分为计算流。在此阶段，后端还可能对某些运算或运算组合针对优化库调用执行模式匹配。
下一步是针对特定目标生成代码。XLA 所含的 CPU 和 GPU 后端使用 LLVM 进行低级 IR、优化和代码生成。这些后端发出有效表示 XLA HLO 计算所需的 LLVM IR，然后调用 LLVM 以从此 LLVM IR 中发出原生代码。目前，GPU 后端通过 LLVM NVPTX 后端支持 NVIDIA GPU。CPU 后端支持多个 CPU ISA。

TensorFlow 算是一个对硬件厂商比较友好的框架，因为 Google 基于 TPU 做软硬全栈设计，所以 XLA 以及 XLA 所基于的 HLO IR 是比较适合硬件厂商进行对接的。

XLA 是 JAX 使用的编译器，目前主要用于为 TPU ，未来不久将支持更多的硬件设备， XLA 的计算过程表示为 HLO IR 中的计算图，之后再进一步表示为硬件绑定的如 CPU/GPU/TPU 的 LLO 表示。JAX 是一个自动求导工具如使用原生的 Python 语法与 Numpy 模块，与 XLA 结合起来可以实现高性能机器学习。

至少，原则上，硬件厂商对接 XLA/HLO 可以获得跟 Google 自身将 TPU 对接到TensorFlow 非常接近的便利性。而TPU的未来规划从一开始同时考虑训练和推理也使得 XLA/HLO 这套技术体系的完备性是符合新硬件的扩展性需要的。包括 Graphcore 以及一些其他硬件厂商选择基于 XLA 对接 TF 就有这方面的考虑。

然而，随着 PyTorch 的 adoption ratio 日渐增加，故事变得有些不一样了。我的认识，FB 在推出 PyTorch 的时候，并没有考虑过软硬全栈设计。事后诸葛亮式的来看，这一方面解放了 FB 的约束，不像 Google 在 TF 1.X 极度倾向于静态图，回头来看，被吐槽很厉害的静态图设计也许当时会有考虑 TPU 软硬全栈设计的 trade-off 考虑？

6. eager mode 下的便利性与 DSA 不友好

而 PyTorch 是可以更多从用户易用性出发，以 eager mode 为主来设计 PyTorch 的整体 API 体系以及执行流程；另一方面，这也为 PyTorch 对接 DSA 架构潜在的挑战埋下了伏笔。

图 eager mode 下带来的 Python 便利性以及 DSA 不友好（红框）

参见上图里的红框部分，这种 eager mode 下很容易出现 Python 操作加速器的逻辑，但实际中一方面为模型开发的可调试性带来便利，另一方面也对 DSA 架构不够友好。

DSA 在获得单位晶体管更强 AI 算力的同时，也对主机端与加速器交互带来的 kernel launch 开销和访存同步开销更为敏感。一个小例子，可以在商汤开源的 MMDetection 的 git repo 里看到类似的代码。实际上，哪怕是在 NV GPU 架构上，随着 V100 引入 TensorCore 开始，这个对 DSA 不友好的问题也变得更加明显了，到了 A100 时代，简直是不可忽略了。

于是，就存在了现实的一些局限：

原子IR（Primitive IR）定义上，存在 HLO （注：这里的 HLO 更多是指 HLO 表示层，可以与 XLA 独立开，比如 MLIR 里的 HLO Dialect ）、ONNX、TorchScript、Relay 几种候选。目前来看 HLO 对训练和推理同时支持，并且在生产环境里经受过检验。ONNX、TorchScript、Relay 对推理的支持尚可，对训练的支持还需要打磨和完善。
即便假设 HLO 这层 IR 能够被大家接受，但 eager mode 带来的 python 侧和加速器侧交互对性能的影响，也会相当大程度上影响到新硬件的效力发挥。所以在IR定义存在可能共识的前提之外，还需要一套有效的机制能够将 eager mode 的建模代码高效映射到加速器上执行。在 WAIC 林达华老师材料里提到的非规则计算和动态计算过程都算是相关的例子了。

在上述两个局限都被克服的情况下，在 AI 框架算子层面之下，找到一个更有利于桥接上层用户和底层硬件后端的途径才存在可行性。这也许是后续硬件厂商可以一起努力推进的方向。

7. eager mode 下的便利性与 DSA 不友好

MLPerf 目前仍然关注模型全网层级的benchmark。未来是否有可能在更原子层面，提供出具备代表性刻划能力的IR描述片段？

图 MLPerf

这些IR定义及描述片段能够完整的描述任何一种AI模型（可以认为是 AI 建模层面的一种“图灵完备”），由硬件厂商在这些描述片段上验证其性能表现，再配合上模型全网层级的 benchmark 来展现新硬件的性能及通用性。这样的一个潜在好处是，除了刷榜的 benchmark 以外，可以提升硬件厂商 benchmark 结果在真实业务场景的可用性。

因为真实场景里使用的模型往往和 MLPerf 的标准模型还是存在差异的（推理稍好一些，训练的话，实际上在真实场景是很少去重跑 MLPerf 的 training benchmark 的，因为一但有了这些经典模型的pre-trained model，需要做的要么是在新的数据集上finetune，这对性能可能就不那么敏感，要么是修改模型结构做新的建模尝试，这就可能未必可以把 MLPerf 的性能结果直接transfer过来了）。

另，在整理这篇文章的时候，自己也梳理了一些之前思考的内容，有些认识随着自己所处软件背景的变化，以及最近吸收到的信息增加，也有所调整：

8. 华为自研框架的必要与硬件能力的制约

对华为当初选择自研框架的动机，现在多了一些理解。MindSpore这个项目的起源一直有不同的说法，这里不再搬各种江湖传闻了。但有一点我现在会更加认同，那就是想推广软硬全栈生态时，如果对上层框架没有掌控力，很容易会因为上层框架的定义导致底层硬件的能力被制约。

图 Mindspore架构图

MindSpore框架架构总体分为MindSpore前端表示层、MindSpore计算图引擎和MindSpore后端运行时三层。
MindSpore前端表示层（Mind Expression，简称ME）：该部分包含Python API、MindSpore IR（Intermediate representation，简称IR）、计算图高级别优化（Graph High Level Optimization，简称GHLO）三部分。
Python API向用户提供统一的模型训练、推理、导出接口，以及统一的数据处理、增强、格式转换接口。
GHLO包含硬件无关的优化（如死代码消除等）、自动并行和自动微分等功能。
MindSpore IR提供统一的中间表示，MindSpore基于此IR进行pass优化。
MindSpore计算图引擎（Graph Engine，简称GE）：该部分包含计算图低级别优化（Graph Low Level Optimization，简称GLLO）、图执行。
GLLO包含硬件相关的优化，以及算子融合、Buffer融合等软硬件结合相关的深度优化。
图执行提供离线图执行、分布式训练所需要的通信接口等功能。
MindSpore后端运行时：该部分包含云、边、端上不同环境中的高效运行环境。

另外商汤的同学曾经提到过内部自研的 Parrot 框架可以兼容 PyTorch，我对这个工作的细节也比较好奇，我想大概率是因为商汤还是关注在内部的垂直 CV 类业务，并且商汤的建模能力已经是国际水准，在引领这个领域的建模趋势。在一个封闭建模场景下，并且可以引领上层建模趋势的情况下，做到兼容 PyTorch API 我想是存在很大技术可行性的。而对于一家硬件厂商，或是一个面向通用场景的云厂商，就是另一个故事了。

9. TF 与 PyTorch以及硬件厂商对 ONNX 的不同态度

ONNX 这个中间 IR 表示，自己有相当长一段时间是持看低的认识的，因为 PyTorch 团队明显在减少投入在上面的精力， TensorFlow 团队在 ONNX 上也没有投入多少精力。随着了解更多硬件行业的细节，对 ONNX 的意义和价值自己又有了一些新的理解。在我看来只要还存在至少超过两家AI框架，ONNX 这种中间 IR 对于硬件公司就有着非常重要的价值。存在且合理，大量的硬件公司会选择对接 ONNX ，甚至优先对接 ONNX ，其实是有其权衡考量的。

图 ONNX High-level system architecture

The ONNX Flow is quit simple:
Starting from an ONNX model, ONNX Runtime first converts the model graph into its in-memory graph representation.
It then applies a number of graph transformations that a) perform a set of provider independent optimizations such cast transformations between float16 and float32, and b) partition the graph into a set of subgraphs based on the available execution providers.
Each subgraph is assigned to an execution provider. We ensure that a subgraph can be executed by an execution provider by querying the capability of the execution provider using the GetCapability() API.

当然，目前对于AI软硬生态来说，还缺乏一个真正的被认可的标准，ONNX虽然有一定认可度，但目前其对于训练场景的描述能力不足，以及ONNX标准已经呈现出碎片化分裂的一些趋势，使得硬件厂家在对接 AI 框架的时候，还是不得不做一些重复的、碎片化的工作。如果未来能够形成一套 bridge AI 硬件和软件生态之间的事实标准，对整个行业的发展也会是善莫大焉。