深度学习编译器（整理某survey）

最新推荐文章于 2025-05-12 18:00:45 发布

原创

最新推荐文章于 2025-05-12 18:00:45 发布 · 1.1k 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能

一、深度学习框架

TensorFlow

PyTorch

MXNet

ONNX:定义了一个统一的表示，DL models的格式方便不同框架之间的转换模型

二、深度学习硬件

通用硬件（CPU、GPU）：通过硬件和软件优化支持深度学习工作负载

GPU:通过多核架构实现高并行性
专用硬件（TPU）：专门为深度学习计算设计，以提高性能和能效
神经形态硬件（TrueNorth）：模拟生物大脑的电子技术

三、硬件特定的代码生成器

FPGA（现场可编程门阵列）在深度学习中具有重要作用。FPGA的代码生成器可以分为处理器架构和流架构两类。

四、通用设计架构

深度学习编译器的通用设计架构包括前端、中间表示（IR）和后端。

前端：将深度学习模型从框架中导入，并转换为计算图表示（如Graph IR）。
中间表示（IR）：DL models 在 DL 编译器中被翻译成多级 IRs，其中 high-level IR 在前端部分，low-level IR 在后端部分。基于 high-level IR，编译器前端需要做一些硬件无关的变换和优化。基于 low-level IR，编译器后端需要做硬件相关的优化，代码生成以及编译。
后端：将高级IR转换为低级IR，并进行硬件特定的优化和代码生成。

五、关键组件

5.1 高级IR

高级IR（Graph IR）用于表示计算和控制流,能够捕捉并且表示多种多样的深度学习模型。

设计目的：建立operators和data之间的控制流和依赖，同时为graph level的优化提供接口

常见的表示方法包括DAG（有向无环图）和Let-binding。高级IR还支持张量计算的不同表示方法，如函数式、Lambda表达式和爱因斯坦表示法。

DAG-based IR（基于有向无环图的IR）

DAG是传统编译器中最常用的表示方式。

在深度学习编译器中：
- 节点表示原子操作（如卷积、池化等）
- 边表示张量或依赖关系（数据流）
- 无环：图中没有循环依赖，数据只能单向流动
有着 DAG 计算图的帮助，DL 编译器可以分析不同算子之间的关系和依赖，并且使用它们来指导后续优化。

下图是一个典型的DAG-based IR,通过节点和边表示深度学习模型的计算图。

缺点：
- 变量作用范围不明确：DAG 无法明确表示 x 和 y 的作用范围。
- 依赖关系隐含：DAG 通过边表示数据流，但无法明确表示变量的生命周期。
- 控制流支持不足：如果计算图中包含条件语句或循环，DAG 无法直接表示。
Let-binding-based IR（基于Let绑定的IR)
- Let-binding是一种编程语言中的概念，用于引入一个新的变量，并将其绑定到一个特定的值或表达式。允许创建一个变量将其初始化为一个值，并在特定的范围内使用这个变量，提高代码的可读性和可维护性。
```
let x = a + b in
let y = x * c in
let z = y + d in
z
```
  解释：
  - 第一个 let 绑定 x = a + b，x 的作用范围是后续的 let 表达式。
  - 第二个 let 绑定 y = x * c，y 的作用范围是后续的 let 表达式。
  - 第三个 let 绑定 z = y + d，z 的作用范围是最后的 z。
    
    最后整个表达式的返回值是z
- 当使用let关键字定义一个表达式时，一个let Node生成，然后它指向表达式中的operator和variable。
  
  let节点包含变量绑定部分和作用域部分：
  
  以上面代码为例，编译器会构建如下的let节点结构：
  第一层 Let 节点：
  - 绑定变量：x
  - 绑定表达式：a + b
  - 作用域：let y = x * c in let z = y + d in z
  第二层 Let 节点：
  - 绑定变量：y
  - 绑定表达式：x * c
  - 作用域：let z = y + d in z
  第三层 Let 节点：
  - 绑定变量：z
  - 绑定表达式：y + d
  - 作用域：z（即 z 的值）
- "Let-binding"是解决语义歧义的一种方法，当使用let关键字定义表达式时，会生成一个let节点，然后它指向表达式中的运算符和变量，而不仅仅是像DAG一样构建变量之间的计算关系。
- 在基于DAG的编译器中，当计算需要获取某个表达式的返回值时，它首先访问相应的节点并搜索相关节点，也称为递归下降技术。
  
  相反，基于Let-binding的编译器计算出let表达式中变量的所有结果并构建变量映射。当需要特定结果时，编译器会查找此映射来决定表达式的结果。
- 在DL编译器中，TVM的Relay IR同时采用了DAG-based IR和Let-binding-based IR，以获得两者的优点。
Representing Tensor Computation(张量计算的表示)
- Functon-based(函数式表示)：
  - 核心思想
    
    是一种基于函数的表示方式，它将复杂的计算任务分解为一系列封装好的函数（算子）。这些函数没有副作用，即他们的输出只依赖于输入，不会影响其他部分的状态。这种方式使得计算过程更加模块化，易于优化和并行化。
  - XLA的HLO例子
    
    XLA是一个用于加速深度学习模型的编译器框架，它通过HLO IR(中间表示)来优化计算任务。
    
    由三个层级组成：
    
    HIoModule:表示整个程序
    
    HIoComutation:表示一个函数
    
    Hlilnstruction:表示一个具体的计算操作
    
    XLA 使用 HLO IR 来同时表示图IR 和操作IR，因此 HLO 的操作能从数据流级别覆盖到算子级别。
- Lambda表达式：
  
  Lambda 表达式是一种基于 index 的形式化表达式，它通过变量绑定和替换描述了计算。
  
  使用 lambda表达式，程序员可以迅速定义一个计算而不用去实现一个新函数。TVM 使用基于 lambda 表达式的 tensor expression（TE）来表示这种 tensor 计算。在 TVM 中，算子被 output tensor 的 shape 和用于计算的 lambda 表达式共同定义。
  - TVM中
    
    张量表达式中的计算运算符由输出张量的形状和计算规则的lambda表达式定义。
    
    举例：将矩阵A和B相加，并将结果存储在输出矩阵C中。
```
import tvm
from tvm import te

# 定义输入矩阵维度
M, N = 2, 2

# 创建TVM计算图上的符号变量
A = te.placeholder((M, N), name='A')
B = te.placeholder((M, N), name='B')

# 使用Lambda表达式定义相加操作
C = te.compute((M, N), lambda i, j: A[i, j] + B[i, j], name='C')

# 创建TVM的调度器
s = te.create_schedule(C.op)

# 编译计算图并执行
func = tvm.build(s, [A, B, C], "llvm")
ctx = tvm.Device("llvm", 0)
a = tvm.nd.array([[1, 2], [3, 4]], ctx)
b = tvm.nd.array([[5, 6], [7, 8]], ctx)
c = tvm.nd.empty((M, N), ctx)
func(a, b, c)
print(c.asnumpy())
```
- Einstein notation
  
  被称作求和约定，是一种用来表示求和的记号约定。它要比 lambda 表达式更容易编程。
  
  以 TC 为例，临时变量的索引不用被特地去定义。
  
  IR 可以基于 Einstein 记号，通过为定义变量的出现，来自动推断出真实的表达式。在 Einstein 记号中，operators 需要是既可以结合又可以交换的。这个限制保证了 reduction operator 可以以任意顺序被执行，使得进一步的并行成为可能。
```
def matmul(A: float[N, K], B: float[K, M]) -> float[N, M]:
    C(n, m) +=! A(n, k) * B(k, m)
```
数据表示 （管理）
- 占位符（Placeholder）
  - 广泛用于符号编程。只是一个具有显式形状信息（例如每个维度的大小）的变量，并且将在计算的后期阶段用值填充。
  - 用于描述张量的形状信息，允许在计算图中定义操作而不需要具体数据。即允许程序员描述操作并构建计算图，而无需关系确切的数据元素。
  - 可以通过占位符来改变输入/输出和其他相应中间数据的形状，而无需改变计算定义。
    
    内存指针直接表示： 当DL编译器使用内存指针直接表示张量数据时，它会将张量数据的实际值存储在内存中，并使用指针来引用这些内存位置。这种方式效率高，适用于已知形状和数据值的情况，但可能不够灵活，无法处理动态形状或未知数据的情况。
    
    Placeholder表示： Placeholder是一种更灵活的数据表示方式。在这种方式中，编译器并不直接存储张量的实际值，而是创建一个Placeholder，表示这个张量的数据将在运行时动态地提供。这对于模型的输入、输出以及未知形状的数据非常有用。Placeholder允许在运行时灵活地传入实际的张量数据，使得编译器能够适应不同的输入和情境。
  - Placeholder的作用
    
    预留输入位置
    
    定义计算的输入
    
    灵活的编程方式
- 动态形状表示：声明placeholder支持未知维度的张量，如TVM的Any和XLA的None。
  
  未知的形状表示对于支持动态模型是必要的。然而，为了完全支持动态模型，应该放宽约束推理和维度检查。此外，还应该实现额外的机制来保证内存的有效性。
- 数据布局（Data Layout）：描述张量在内存中的组织方式，通常是从逻辑索引到内存索引的映射。通常包括维度序列：如NCHW和NHWC格式,padding,striding。
  
  TVM和Glow将数据布局表示为算子参数，并需要此类信息进行计算和优化。
  
  在TVM中，数据布局信息通常作为操作符的参数之一来表示。每个操作符都有一个或多个输入张量，每个张量都有自己的形状(shape)，数据类型(dtype)和数据布局(layout)等信息。

操作符支持

深度学习编译器支持的算子负责表示深度学习工作流，他们是计算图的节点

算子通常包括：

代数算子（+, ×, exp and topK）

神经网络算子（convolution and pooling）

张量算子（reshape, resize and copy）

广播和归约算子（例如，min and argmin）

控制流运算符（conditional and loop）

在这里我们选择在不同的深度学习编译器中经常使用的三个代表性算子进行说明。

广播（Broadcast）：可以负责数据并生成具有兼容形状的新数据。

例如：对于加法运算符，输入张量应具有相同形状。一些编译器通过提供Broadcast来放宽机制。

import tvm
from tvm import relay
import numpy as np
# 创建输入变量
x = relay.var("x", shape=(3, 1), dtype="float32")
y = relay.var("y", shape=(3, 4), dtype="float32")

# 进行广播操作
broadcasted_x = relay.broadcast_to(x, shape=(3, 4))
result = relay.add(broadcasted_x, y)

# 创建 Relay 函数
func = relay.Function([x, y], result)

# 编译 Relay 函数
mod = tvm.IRModule.from_expr(func)
target = "llvm"
compiled_func = relay.create_executor(mod = mod)

# 输入数据
input_x = tvm.nd.array(np.array([[1], [2], [3]],dtype=np.float32))
input_y = tvm.nd.array(np.array([[4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]],dtype=np.float32))

# 执行函数
output = compiled_func.evaluate()(input_x, input_y)
print(output)
#[[ 5.  6.  7.  8.]
#[10. 11. 12. 13.]
#[15. 16. 17. 18.]]

控制流（Control Flow）：支持条件语句和循环，用于表示复杂的模型（如RNN）。RNN和强化学习(RL)等模型依赖于循环关系和数据依赖的条件执行，这需要控制流。

import tvm
from tvm import relay

'''
x = input()
if x < 10:
    x * = 2
else 
    x * = 3

'''

# 创建输入变量
x = relay.var("x", shape=(), dtype="float32")

# 创建条件语句
condition = relay.less(x, relay.const(10, "float32"))
then_branch = relay.multiply(x, relay.const(2, "float32"))
else_br