文@小P家的 001314
前言
大家好，本次分享的主题是深度学习编译器。
我们在 AI 框架前沿技术分享 中提到，将灵活的模型表达和高效的模型执行联系在一起的深度学习编译技术是 AI 框架演进的一个主要方向，在本次分享中，我们将从以下三个方面展开详细介绍：
- 从 AI 框架演进看深度学习编译器
- 深度学习编译技术面临的挑战
- 我们商汤 SenseParrots 的深度学习编译方案
从 AI 框架演进看深度学习编译器
笔者认为 AI 框架的核心任务有两个，第一个是定义用户友好的模型表达方式，第二个是能够高效地执行模型。
从第一代 AI 框架 caffe、Theano、Torch，到第二代框架 TensorFlow、PyTorch，包括近几年不断涌现出的新的 AI 框架 TVM、PaddlePaddle、MindSpore、MegEngine、OneFlow，其中涉及的绝大多数深度学习框架技术，都是为了这两个任务服务的。其中，
- 基于 Python 构建深度学习模型，在框架内部实现自动求导，动态图机制等技术可以认为是在为第一个任务服务；
- 基于计算图进行模型表示与执行，计算与通信并行，量化，混合精度等技术可以认为是在为第二个任务服务。
编译是把用源语言编写的源程序变成目标程序的过程，所以 AI 框架中把用户模型表达变成可执行代码这个过程，本身就是一个编译过程，AI 框架可以看做深度学习编译器+执行引擎+其他组件的复合体。
基于这个理论，有了 TensorFlow 的 XLA，PyTorch 的 TorchScript、TVM 中的两层 IR(Intermediate Representation)等设计，深度学习编译器已经被深深地嵌入到了 AI 框架设计中。
深度学习编译技术面临的挑战
我们认为深度学习编译技术面临的挑战主要存在于以下几个方面：
- 用户模型表达解析层面的挑战
- 深度学习编译器系统设计层面的挑战
- 高性能代码生成层面的挑战
2.1 用户模型表达解析层面的挑战
现有 AI 框架在前端模型表达方面主要分为两个流派，一部分框架使用限定的算子集合搭建模型，如 TensorFlow、PyTorch，另一部分框架基于自定义 DSL(Domain Specified Language)搭建模型，如 TVM。
两种做法各有利弊，使用限定的算子集合搭建模型，更方便进行模型表达到框架中间表示的转化过程，但是表达能力有限，尤其在表达控制流上会有一些困难；使用自定义 DSL 搭建模型，往往具备更完善的表达能力，但是增加了模型表达解析和后续基于中间表示优化的难度。
无论使用哪种前端表达，基于 Python 构建深度学习模型已经形成了大家的共识，进而会遇到共同的挑战，我们将这些挑战总结为四个方面：
1. 多种数据结构：
在表达深度学习模型，尤其是添加自定义算子时，算法研究员往往会使用 AI 框架中定义的张量数据结构（PyTorch 中的 Tensor）作为基础数据结构，并且使用 Python native 数据结构（list、tuple 和 map）作为补充。因此，用户给出的模型表达中，总是混合使用多种数据结构，如何在编译过程中表示各种数据结构类型是一个挑战。
另外，Python 是一种动态语言，无法提前预测数据类型，这使上述问题变得更加复杂。
2. 多种运算操作：
第二个挑战是多种运算操作带来的挑战。
用户给出的模型表达中，可能用到三类运算：
- 以 Python 运算符形式体现的基本运算，如加减乘除等；
- 第三方软件提供的一些 API，例如 PyTorch 提供的一系列 API 接口；
- 索引运算。
特别的，索引运算的实现往往会引入 view 的概念，即输出与输入数据共享相同的物理内存，处理输出时，输入也会更改。在编译过程中，需要正确地记录并处理这种内存共享机制。
3. 辅助语句处理：
用户在表达深度学习模型时，出于代码安全性和易用性的考虑，还可能会用到一些辅助语句。因此，除了必需的计算和内存访问语句外，在用户模型表达解析过程中还需要处理装饰器、断言、函数调用等辅助语句。这些语句往往不会直接映射到深度学习编译器的IR中，但是需要保证它们不能带来副作用。
4. 不确定性行为：
最后，除了上述挑战之外，在用户模型表示语句中还存在一些不确定性行为。我们这里总结出两种不确定性行为：
value dependent问题，即输入数据不同会导致输出结果的形状或者控制流走向变化；- 非纯函数现象，指即使输入数据完全相同，函数也会产生不同的结果。
深度学习编译器系统设计层面的挑战
从系统设计层面来看，深度学习编译器设计需要考虑以下挑战：
1. 编译时机与开销：
首先，通过一次编译优化生成代码，然后重复使用这些编译好的代码是深度学习编译技术取得性能提升的主要来源。在模型表达中，一个函数可能被多次调用，每次调用输入数据都可能不一样，如何处理多次调用函数的编译开销是深度学习编译器设计的一个挑战。
另一方面，随着PyTorch 用户的持续增长，一些基于静态图的框架，如 TensorFlow、PaddlePaddle 也把动态图模式作为默认执行模式，如何在基于动态图的框架中选取合适的时机，完成动态图静态化和编译优化也是一个关键问题。
2. 与反向过程兼容：
几乎所有 AI 框架中都有自动求导机制，这些自动求导机制的实现往往要求框架中的算子同时具备 forward 和 backward 两个过程的实现。
对于用户自定义算子，如何获取反向过程 IR 表示并进行编译优化则是另一个挑战。
高性能代码生成层面的挑战
生成代码的性能一直是评价深度学习编译器优劣的关键指标。从最开始的基于模板的代码生成方法，到 Halide 中提出的 compute 和 schedule 分离的优化策略，再到 TVM 中的 auto-schedule，直至今天的 Ansor 、多面体优化 等，都在尝试提高深度学习编译器生成代码的性能。
商汤 SenseParrots 的深度学习编译方案
我们在自研 AI 框架 SenseParrots 中，内置了自研深度学习编译器 Elena，并通过和 AI 框架执行引擎协同设计，重新设计了一套深度学习编译方案。该方案主要由三个部分组成：
- 用户模型表达解析模块；
- 基于两阶段 IR 的代码生成模块；
- AI 框架执行引擎。
在这里，用户模型表达解析模块负责将用户基于 Python 的模型表达翻译到 AI 框架自己的 IR 上。我们会在后续的分享中介绍如何解决多种数据结构、多种运算操作、辅助语句处理、不确定性行为等问题。
基于两阶段 IR 的代码生成模块负责生成高性能的可执行代码，在后续的系列分享中也会介绍一些代码生成技术实现。
AI 框架执行引擎在这里起到一个连接作用，将用户模型表达解析模块和基于两阶段 IR 的代码生成模块组合在一起，串通整个深度学习编译流程。
结语
感谢阅读，关于深度学习编译相关问题，欢迎大家在评论区给我们留言，一起讨论哦~
更期待有相同兴趣的同学加入我们，一同探索解决深度学习编译器的问题和挑战！
关注 公众号「SenseParrots」，获取人工智能框架前沿业界动态与技术思考：
