一、深度学习框架
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_branch = relay.multiply(x, relay.const(3, "float32")) result = relay.If(condition, then_branch, else_branch) # 创建 Relay 函数 func = relay.Function([x], result) # 编译 Relay 函数 mod = tvm.IRModule.from_expr(func) target = "llvm" compiled_func = relay.create_executor(mod = mod) # 输入数据 input_data = tvm.nd.array(np.array(5., dtype=np.float32)) # 执行函数 output = compiled_func.evaluate()(input_data) print(output) ## # 10.0
-
导数(Derivative):支持自动微分,用于模型训练。
算子Op的导数算子将Op的输出梯度和输入数据作为输入,然后计算Op的梯度。
尽管一些DL编译器(例如TVM)支持自动微分,但当应用链式规则时,它们需要高级IR种所有运算符的导数。
TVM 致力于提供代数算子和神经网络算子的导数算子,程序员可以使用这些导数算子来构建定制算子的导数。
from tvm import relay import numpy as np import tvm mod = tvm.IRModule() x = relay.var('x', shape=()) y = relay.multiply(x, x) mod['main'] = relay.Function([x], y) mod = relay.transform.InferType()(mod) grad_ir = relay.transform.gradient(mod['main'], mode = 'first_order') b_mod = tvm.IRModule.from_expr(grad_ir) print(mod['main']) ''' fn (%x: float32 /* ty=float32 */) -> float32 { multiply(%x, %x) /* ty=float32 */ } /* ty=fn (float32) -> float32 */ ''' lib =relay.build(b_mod,'llvm') m = tvm.contrib.graph_executor.GraphModule(lib['default'](tvm.device('llvm',0))) input_data = tvm.nd.array(np.array(5.,dtype="float32")) m.set_input('x',input_data) m.run() res = m.get_output(0,tvm.nd.empty(input_data.shape)).numpy() grad = m.get_output(1,tvm.nd.empty(input_data.shape)).numpy() print(f"result is : {res}, gradient is : {grad}") ''' result is : 25.0, gradient is : 10.0 '''
-
自定义操作符:允许用户定义新的操作符,增强编译器的扩展性。
它允许程序员为特定目的定义其运算符。提供对自定义运算符的支持提高了 DL 编译器的可扩展性。
例如,在Glow中定义新的算子时,程序员需要实现逻辑和节点的封装。此外,如果需要的话,还需要额外的努力,例如降低步骤、操作IR生成和指令生成。
而 TVM 和 TC 除了描述计算实现之外,需要较少的编程工作。具体来说,TVM的用户只需要描述计算和时间表并声明输入/输出张量的形状。而且定制的算子通过hook的方式集成了Python函数,进一步减轻了程序员的负担。
-
5.2 低级IR
5.2.1 与High-Level的区别
-
High-Level IR:通常更接近于神经网络模型的原始描述。它负责高层次的优化和模型表示,允许进行更抽象的操作。在这个层次上,针对模型的整体结构进行优化、融合、替代等操作。
-
Low-Level IR: 更接近底层硬件的特性和操作。它负责将高级表示翻译成更接近硬件的表达形式,以便进行低层次的优化和代码生成。
5.2.2 Low-Level IR的实现
Low-Level IR 比 High-Level IR更细粒度的表示形式描述 DL 模型的计算,从而通过提供调整计算和内存访问的接口来实现目标相关的优化。
在本节中,我们将低级 IR 的常见实现分为三类:基于Halide-based IR、polyhedral-based IR 和其他的 IR。
Halide-based IR
Halide最初被提出用于并行化图像处理,并且在DL编译器(例如TVM)中被证明具有可扩展性和高效性。
Halide的基本理念是计算和调度的分离。
-
计算:定义算法的数学逻辑(如卷积、池化等)。
-
调度:定义如何高效地执行计算(如循环展开、并行化、内存布局)。
核心思想:
分离计算和调度。
计算部分:定义了算法的数学逻辑,例如一个简单的图像卷积可以表示为:
Func blur_x, blur_y; Var x, y; blur_x(x, y) = input(x, y) + input(x+1, y) + input(x+2, y); blur_y(x, y) = (blur_x(x, y) + blur_x(x, y+1) + blur_x(x, y+2)) / 9;
调度部分:定义了如何高效地执行计算,例如可以通过一下调度策略优化卷积操作:
blur_x.compute_root().parallel(y).vectorize(x, 8); blur_y.compute_root().parallel(y).vectorize(x, 8);
parallel(y):在y维度上并行化
vectorize(x,8):在x维度上使用SIMD指令(每次处理8个像素)
举个例子:
下面是一个简单的自定义图像处理函数,该函数计算了一个图像的梯度,并在生成的图像上检查了计算结果的正确性。其中计算的定义(像素计算)与调度策略(reorder和parallel)是分开的
#include <iostream>
#include <Halide.h>
int main() {#定义计算逻辑,gradient接收两个变量并计算他们的和,输出是一个二维数组,每个位置(x,y)的值为x+yHalide::Func gradient;Halide::Var x, y;Halide::Expr e = x + y;gradient(x, y) = e;
#定义调度策略#重新排序循环的顺序,y放在内循环,x放在外循环。可以优化内存访问模式,减少缓存未命中。gradient.reorder(y,x);#将内存循环y并行化,利用多核CPU的并行计算能力。gradient.parallel(y);
#用于存储计算结果Halide::Buffer<int32_t> output = gradient.realize({800, 600});
#验证结果for (int j = 0; j < output.height(); j++) {for (int i = 0; i < output.width(); i++) {if (output(i, j) != i + j) {printf("Something went wrong!\n""Pixel %d, %d was supposed to be %d, but instead it's %d\n",i, j, i + j, output(i, j));return -1;}}}printf("Success!\n");return 0;
}#include <iostream>
#include <Halide.h>
int main() {#定义计算逻辑,gradient接收两个变量并计算他们的和,输出是一个二维数组,每个位置(x,y)的值为x+yHalide::Func gradient;Halide::Var x, y;Halide::Expr e = x + y;gradient(x, y) = e;
#定义调度策略#重新排序循环的顺序,y放在内循环,x放在外循环。可以优化内存访问模式,减少缓存未命中。gradient.reorder(y,x);#将内存循环y并行化,利用多核CPU的并行计算能力。gradient.parallel(y);
#用于存储计算结果Halide::Buffer<int32_t> output = gradient.realize({800, 600});
#验证结果for (int j = 0; j < output.height(); j++) {for (int i = 0; i < output.width(); i++) {if (output(i, j) != i + j) {printf("Something went wrong!\n""Pixel %d, %d was supposed to be %d, but instead it's %d\n",i, j, i + j, output(i, j));return -1;}}}printf("Success!\n");return 0;
}
编译器会自动尝试多种不同的计算任务排序、并行化策略等,然后评估每个方案的性能,最终选择效果最好的一种来生成高效的底层代码,从而加速深度学习模型的执行。 这种自动化的过程有助于减轻开发者的工作负担,同时确保生成的代码尽可能地优化。
Halide 中内存引用和循环嵌套的边界仅限于与轴对齐的有界框。
polyhedral-based IR
它使用线性规划、仿射变换和其他数学方法来优化具有边界和分支的静态控制流的基于循环的代码。
与 Halide 不同的是,内存引用和循环嵌套的边界可以是多面体模型中任意形状的多面体。
多面体模型(polyhedral model)是在编译器优化中应用的一种高级技术,旨在优化循环结构的并行性和局部性。 它将程序的循环依赖关系转化为多维几何空间中的多面体,从而提供了一种抽象的方式来分析、优化和重组循环代码。
for (int i = 1; i < N; ++i) for (int j = 1; j < N; ++j) A[i][j] = A[i-1][j] + A[i][j-1];
其中数据的依赖关系图如下所示:
其中循环迭代顺序图如下所示:
如果经过倾斜变换,可得到
for (int d = 2; d <= 2 * 5; ++d) { for (int j = max(1, d - 5); j <= min(5, d - 1); ++j) {int i = d - j;A[i][j] = A[i-1][j] + A[i][j-1]; } }
其中循环迭代顺序图如下所示:
因为在j维度上,没有数据依赖,我们将可以在j维度上实现并行,优化循环结构的并行性。
这种灵活性使得多面体模型广泛应用于通用编译器中。然而,这种灵活性也阻碍了与调整机制的集成。
然而,由于能够处理深度嵌套循环,许多深度学习编译器,例如 TC 和 PlaidML(作为 nGraph 的后端),都采用多面体模型作为其low-level IR。
基于多面体的 IR 可以轻松应用各种多面体变换(例如fusion、tiling、sinking和mapping),包括设备相关和设备无关的优化。基于多面体的编译器借用了许多工具链,例如 isl 、Omega 、PIP 、Polylib 和 PPL 。
TC在低强度IR方面有其独特的设计,它结合了Halide和多面体模型。它使用基于Halide的IR来表示计算,并采用基于多面体的IR来表示循环结构。 TC通过抽象实例呈现详细的表达,并引入具体的节点类型。
简而言之,TC使用域节点来指定索引变量的范围,并使用上下文节点来描述与硬件相关的新迭代变量。它使用带节点来确定迭代的顺序。过滤器节点表示与语句实例组合的迭代器。 Set和sequence是关键字,用于指定过滤器的执行类型(并行和串行执行)。此外,TC使用扩展节点来描述代码生成的其他必要指令,例如内存移动。
PlaidML 使用基于多面体的 IR(称为 Stripe)来表示张量运算。它通过将并行多面体块的嵌套扩展到多个级别来创建可并行代码的层次结构。此外,它允许将嵌套多面体分配给嵌套内存单元,提供了一种将计算与内存层次结构相匹配的方法。在Stripe中,硬件配置独立于内核代码。 Stripe 中的标签(在其他编译器中称为 pass)不会更改内核结构,但提供有关优化遍的硬件目标的附加信息。 Stripe 将 DL 运算符拆分为适合本地硬件资源的图块。
Other unique IR
有些 DL 编译器可以在不使用 Halide 和多面体模型的情况下实现定制的low-level IR。
在定制的low-level IR 上,他们应用特定于硬件的优化并lowers to LLVM IR。
Glow 中的低级 IR 是基于指令的表达式,对地址引用的张量进行操作。
Glow(Graph Lowering)是一个开源的深度学习编译器项目,由Facebook AI Research(FAIR)开发
Glow 低级 IR 中有两种基于指令的函数:声明和编程。
-
第一个声明了在程序的整个生命周期中存在的恒定内存区域的数量(例如输入、权重、偏差)。
-
第二个是本地分配区域的列表,包括函数(例如 conv 和 pool)和临时变量。
指令可以在全局内存区域或本地分配的区域上运行。
此外,每个操作数都用限定符之一进行注释:
-
@in 表示操作数从缓冲区读取;
-
@out表示操作数写入缓冲区;
-
@inout 表示操作数读取和写入缓冲区。
这些指令和操作数限定符帮助 Glow 确定何时可以执行某些内存优化。
MLIR 深受 LLVM 的影响,它是比 LLVM 更纯粹的编译器基础设施。 MLIR 重用了 LLVM 中的许多思想和接口,位于模型表示和代码生成之间。
MLIR 具有灵活的类型系统并允许多个抽象级别,并且它引入了dialect来表示这些多个抽象级别。每种dialect都包含一组定义的不可变操作。
MLIR 当前的dialect包括 TensorFlow IR、XLA HLO IR、实验性多面体 IR、LLVM IR 和 TensorFlow Lite。还支持dialect之间的灵活转换。
此外,MLIR 可以创建新的dialect来连接到新的低级编译器,这为硬件开发人员和编译器研究人员铺平了道路。
XLA 的 HLO IR 可以被视为高级 IR 和低级 IR,因为 HLO 的粒度足够细,可以表示特定于硬件的信息。此外,HLO 支持特定于硬件的优化,并可用于发出 LLVM IR。
Code Generation based on Low-Level IR
大多数DL编译器采用的低级IR最终可以lowers to LLVM IR,并受益于LLVM成熟的优化器和代码生成器。
此外,LLVM 可以从头开始显式地为专用加速器设计自定义指令集。
然而,传统编译器在直接传递给 LLVM IR 时可能会生成质量较差的代码。
为了避免这种情况,DL 编译器应用两种方法来实现硬件相关的优化:
-
在 LLVM 的上部 IR 中执行特定于目标的循环转换(例如,基于 Halide 的 IR 和基于多面体的 IR),
-
提供有关优化过程的硬件目标的附加信息。
大多数深度学习编译器都应用这两种方法,但侧重点不同。
一般来说,更喜欢前端用户的 DL 编译器(例如 TC、TVM、XLA 和 nGraph)可能会关注 1。
而更倾向于后端开发人员的 DL 编译器(例如 Glow、PlaidML 和 MLIR)可能会关注 1,重点关注2。
DL编译器中的编译方案主要分为两类:即时(JIT)和提前(AOT)。
-
对于 JIT 编译器来说,它可以即时生成可执行代码,并且可以通过更好的运行时知识来优化代码。
-
AOT 编译器首先生成所有可执行二进制文件,然后执行它们。因此,它们在静态分析方面比 JIT 编译具有更大的范围。
-
此外,AOT 方法可以应用于嵌入式平台的交叉编译器(例如 C-GOOD ),并支持在远程计算机(TVM RPC)和定制加速器上执行。
5.3 前端优化
将深度模型作为输入,然后将模型转换为计算图的表示形式。即将深度学习模型从高级深度学习框架转换为可执行的机器码或中间表示之前,对模型进行一系列的优化操作。
需要实现不同的格式变换来支持不同深度学习框架中的格式。
构建计算图后,前端应用图级优化,许多优化更容易在图级别上识别和执行,因为图提供了计算的全局视图。
计算图优化包含通用编译器里的优化技术以及 DL 领域特定的优化,这些优化在计算图的基础上减少了冗余并且提高了了性能。
前端优化通常由passes定义,通过遍历计算图的节点并执行图转换实现。
在编译器领域,"pass"(通常称为编译器通行证或优化通行证)是指编译器的一个阶段或一个模块,它对输入程序或代码执行一系列特定的转换、优化或分析操作。
每个pass通常执行一组相关的任务,以改变程序的形式、提高代码质量、优化性能或进行其他编译任务。
pass的目的是将源代码或中间表示(IR)从一个状态转换到另一个状态,以便后续pass可以继续执行更高级别的优化或代码生成任务。通行证之间的顺序和数量通常取决于编译器的设计和优化策略。
以下是一些pass可能执行的任务示例:
-
词法分析和语法分析:这是编译器的前端阶段,它将源代码解析成语法树或中间表示,以便后续pass能够理解代码的结构。
-
语义分析:pass可以执行类型检查、作用域分析、错误检测等任务,以确保源代码的语义正确性。
-
优化pass:这些pass执行各种代码优化操作,例如常量折叠、死代码消除、循环展开、操作融合等,以提高程序性能。
-
中间表示生成:pass可以将优化后的代码转换为中间表示(IR),这是编译器用来生成目标代码的内部表示形式。
-
代码生成:pass将中间表示转换为目标平台的机器代码或汇编代码。
-
链接:如果编译器处理多个源文件或库,pass会将它们组合成一个可执行程序或库。
pass之间的顺序和数量可以因编译器的设计而异。
优化pass通常在前端pass后,但在代码生成之前。
编译器开发者可以根据编译器的目标和优化策略来确定pass的组织和执行顺序。
这些pass是编译器内部的模块,它们协同工作以实现源代码到目标代码的转换过程。
5.3.1 节点级优化
关注的是计算图中的单个操作节点,即神经网络中的单个计算步骤。旨在改进每个节点的计算效率和资源利用。
-
Nop消除:消除无用的操作节点。
-
零维张量消除:消除输入为零维张量的操作。
-
节点替换(用其他成本较低的节点替换节点)
在通用编译器中,Nop Elimination(空操作消除)删除占用少量空间但不指定任何操作的无操作指令。
在深度学习编译器中,Nop Elimination负责消除缺乏足够输入的操作。
例如:可以消除只有一个输入张量的 sum 节点
可以消除填充宽度为零的 padding 节点。
假设A是零维张量,B是常量张量,那么A和B的求和运算节点可以替换为已经存在的常量节点B,而不影响正确性。
假设C是3维张量,但一维形状为零,如{0,2,3},因此C没有元素,可以消除argmin/argmax运算节点。
我们使用tvm实现一个节点消除的例子:
import tvm
import numpy as np
from tvm import te
import tvm.relay as relay #tvm中定义和优化计算图的模块
def add_example(shape):a = relay.var("a", relay.TensorType(shape, "float32"))y = relay.multiply(a, relay.const(2, "float32")) # 先定义 y 为 a 的两倍y = relay.add(y, relay.const(0, "float32")) # 然后加上 0return relay.Function([a], y) # 只返回 a,因为 b 没有定义
shape1 = (1, 256)
f = add_example(shape1)
mod = tvm.IRModule.from_expr(f)
print(mod)
'''
def @main(%a: Tensor[(1, 256), float32]) {
%0 = multiply(%a, 2f);
add(%0, 0f)
}
'''
mod1 = relay.transform.SimplifyExpr()(mod)
print(mod1)
'''
def @main(%a: Tensor[(1, 256), float32] /* ty=Tensor[(1, 256), float32] */) -> Tensor[(1, 256), float32] {
multiply(%a, 2f /* ty=float32 */) /* ty=Tensor[(1, 256), float32] */
}
'''
可以看到,当add操作的y变量为0时候,add节点被消除了
5.3.2 块级优化
-
代数简化:通过代数恒等式简化计算,如常量折叠、强度折减等。
一个深度学习模型在前端表示一般都是一个DAG图,各种算子可以利用等价计算的方式进行计算量优化,其中代数简化主要包括:
1.代数识别:这一优化技术旨在识别和简化计算图中的代数表达式。通过识别节点之间的代数关系,编译器可以将一系列节点替换为更简单的等效形式,从而减少计算的复杂性。
合并相同权重的操作:如果在计算图中多次使用相同的权重参数进行卷积操作,编译器可以将这些操作合并为一个,以减少计算的复杂性。 消除冗余操作:如果计算图中包含冗余的操作,例如相同的激活函数应用多次,编译器可以消除其中一些操作,以减少计算量。 GEMM优化例子: 1.有两个输入矩阵 A 和 B。 2.对两个输入矩阵进行转置操作,分别得到 AT 和 BT。 3.然后将 AT 和 BT 相乘,得到结果矩阵 C。
这种方法在数学上是正确的,但它涉及两次矩阵的转置操作,这可能会导致性能下降,尤其是在大规模矩阵计算时,因为矩阵的转置需要额外的计算和内存访问。
优化的思路是改变操作的顺序,如下: 1.有两个输入矩阵 A 和 B。 2.将矩阵 B 与矩阵 A 直接相乘,得到结果矩阵 C。 3.如果需要,再对矩阵 C 进行转置操作,得到 CT。 这种优化的关键在于,通过改变操作顺序,我们只需要在最后一步才进行一次转置操作,而不是在两个输入矩阵上都进行转置操作,从而减少了计算和内存访问的开销、
2.强度降级(Strength Reduction):强度降级是一种优化策略,通过将高成本的运算操作替换为相对廉价的操作来减少计算的成本。
用移位操作代替乘法:将乘法操作替换为位移操作,特别是在权重是2的幂次方时,可以显著提高计算速度。例如,将
x * 8
替换为x << 3
。 用累加操作代替多次加法:如果计算图中包含多次相同的加法操作,编译器可以将它们替换为累加操作,从而减少加法的次数。3.常量折叠(Constant Folding):常量折叠是一种将常量表达式替换为其计算结果的优化技术。如果计算图中包含了大量的常量节点,编译器可以在编译时计算这些常量表达式的值,并将其替换为结果,从而减少计算的复杂性和运行时开销。
计算常量表达式:如果计算图中包含常量操作,编译器可以在编译时计算这些常量表达式的值。例如,将
3 * 4
替换为12
。 移除无用的常量:如果计算图中包含未使用的常量节点,编译器可以将其移除以减少计算图的复杂性。这些代数简化优化技术考虑了节点序列,并利用不同类型节点之间的可交换性、可结合性和可分配性等代数性质,以简化计算。
-
算子融合:将多个操作融合为一个操作,减少中间结果的存储和计算开销。
算子融合(Operator Fusion)是一项至关重要的优化技术。它的目标是将多个神经网络操作或算子合并成一个更大的算子,以提高计算效率和减少资源消耗。
-
更好的计算共享:通过将多个操作融合成一个,可以减少计算之间的数据传输和中间结果的存储需求。这提高了计算资源的利用率,特别是在GPU等硬件加速器上。
-
消除中间分配:运算符融合通常会减少或消除不必要的中间分配和内存操作。这有助于减少内存占用和提高内存带宽的效率。
-
进一步优化的便利性:融合后的运算符通常更容易进行进一步的优化。例如,可以更容易地对融合的运算符应用常量传播、代数化简、强度降级等优化技术,以减少计算开销。
-
减少启动和同步开销:在某些硬件上,启动和同步运算符的开销可以很显著。通过融合运算符,可以减少启动和同步操作的次数,从而提高计算效率。
算子融合通常在深度学习计算图的不同操作之间执行。 例如,卷积操作、激活函数操作和池化操作可以融合成一个单一的运算符,称为卷积层(Convolution Layer),这有助于减少计算和数据传输的复杂性。
以下是一个简单的示例,说明运算符融合如何提高计算效率: 考虑以下计算图片段:
Input -> Convolution -> ReLU -> Pooling -> Output
在算子融合之前,这个计算图中有四个独立的运算符。 但是,通过算子融合,可以将这些算子合并为一个算子,如下所示:
Input -> Conv-Relu-Pool -> Output
这个融合后的运算符执行相同的功能,但在计算上更高效,减少了计算和数据传输的开销。 总之,深度学习编译器前端中的块级优化的运算符融合是一项关键的优化技术,可以显著提高深度学习模型的执行效率,减少资源消耗,并为后续的优化提供更好的基础。 它是实现高性能深度学习推断和训练的重要组成部分。
-
-
算子下沉:将某些操作(如转置)移动到计算图的较低层次,以便进一步优化。
通过这种优化,许多类似的运算彼此更加接近,为代数简化创造了更多的机会。
我们使用tvm实现constant floding 和operator fusion的例子:
import numpy as np
import tvm
from tvm import te
import tvm.relay as relay
def example():
shape = (1, 64, 54, 54)
c_data = np.empty(shape).astype("float32")
c = relay.const(c_data)
weight = relay.var("weight", shape=(64, 64, 3, 3))
x = relay.var("x", relay.TensorType((1, 64, 56, 56), "float32"))
conv = relay.nn.conv2d(x, weight)
y = relay.add(c, c)
y = relay.multiply(y, relay.const(2, "float32"))
y = relay.add(conv, y)
z = relay.add(y, c)
z1 = relay.add(y, c)
z2 = relay.add(z, z1)
return relay.Function([x, weight], z2)
f = example()
mod = tvm.IRModule.from_expr(f)
print(mod)
'''
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) {
%0 = add(meta[relay.Constant][0], meta[relay.Constant][0]);
%1 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]);
%2 = multiply(%0, 2f);
%3 = add(%1, %2);
%4 = add(%3, meta[relay.Constant][0]);
%5 = add(%3, meta[relay.Constant][0]);
add(%4, %5)
}
'''
mod1 = relay.transform.FoldConstant()(mod)
print(mod1)
'''
def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */;
%1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */;
%2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */;
%3 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */;
add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */
}
'''
mod2 = relay.transform.FuseOps()(mod1)
print(mod2)
'''
def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] {
%4 = fn (%p0: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %p1: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %p2: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, %p3: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */;
%1 = add(%0, %p2) /* ty=Tensor[(1, 64, 54, 54), float32] */;
%2 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */;
%3 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */;
add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */
} /* ty=fn (Tensor[(1, 64, 56, 56), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(1, 64, 54, 54), float32], Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */;
%4(%x, %weight, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */
}
'''
可以看出经过constant floding和operator fusion最后只剩下一个算子%4=fn
5.3.3 数据流级优化
数据流级别优化关注的是计算图中不同块之间的数据流和依赖关系。它旨在优化数据的传输和处理流程以提高整个计算图的效率
-
公共子表达式消除(CSE):消除重复计算的子表达式。
共同子表达式消除是一种优化技术,它旨在消除计算图中重复计算相同表达式的情况。
当多个操作需要计算相同的中间结果时,只需计算一次,并在需要时重复使用这个结果,而不必重复计算
考虑以下计算图片段:
A = B + C D = B + C
这里的表达式
B + C
在两个地方都计算了两次。通过共同子表达式消除,可以将其计算一次,然后重用结果:A = B + C D = B + C
-
死代码消除(DCE):消除无用的代码。
死代码消除是一种用于移除计算图中不会影响最终输出结果的无效操作或节点的技术。
这些节点通常是由于模型重构或其他原因而变得无效的。
-
静态内存规划:优化内存分配,减少内存占用。
静态内存规划是指在模型编译期间为模型的中间结果分配内存空间,以减少在运行时的内存分配和释放开销。这有助于提高执行效率。
In-Place Memory Sharing(原地内存共享):
-
解释:不同的操作或层次可以共享相同的内存空间,以减少内存占用。这意味着在模型的计算过程中,相同的内存区域可以用于不同操作的输入和输出,而不必每次都为它们分配新的内存。
-
优点:原地内存共享可以减少内存占用和内存分配开销,提高内存使用效率。它特别适用于内存有限的设备,如边缘设备。
-
示例:在卷积神经网络中,输入特征图和输出特征图的内存可以被多个卷积操作共享,因为它们的尺寸和数据类型相同。这可以通过指定输入和输出张量的内存布局来实现。
Standard Memory Sharing(标准内存共享):
-
解释:不同操作或层次可以共享相同的内存空间,但在计算过程中需要确保不会互相干扰。这通常需要使用额外的同步和管理机制来确保数据的正确性。
-
优点:标准内存共享可以减少内存占用和内存分配开销,但与原地内存共享不同,它更注重数据的安全性和正确性。
-
示例:在多线程或多设备环境中,标准内存共享可能会用于确保多个操作之间的数据共享和同步。例如,在分布式深度学习训练中,不同的计算节点可以共享模型参数,但需要使用同步机制确保参数的一致性。
-
-
布局转换:在计算图中不同操作之间进行数据布局的变换,以适应不同硬件的需求。例如将数据从行优先布局转换为列优先布局以提高内存访问效率。
import numpy as np
import tvm
from tvm import te
import tvm.relay as relay
def add_example(shape):
a = relay.var("a", relay.TensorType(shape, "float32"))
b = relay.add(a, relay.const(1,"float32"))
d = relay.add(a, relay.const(100,"float32"))
y1 = relay.multiply(b,relay.const(2,"float32"))
c = relay.add(a, relay.const(1,"float32"))
y2 = relay.multiply(c,relay.const(3,"float32"))
y = relay.add(y1,y2)
return relay.Function([a],y)
shape1 = (1, 256)
f = add_example(shape1)
mod = tvm.IRModule.from_expr(f)
print(mod)
'''
def @main(%a: Tensor[(1, 256), float32]) {
%0 = add(%a, 1f);
%1 = add(%a, 1f);
%2 = multiply(%0, 2f);
%3 = multiply(%1, 3f);
add(%2, %3)
}
'''
mod1 = relay.transform.EliminateCommonSubexpr()(mod)
print(mod1)
'''
def @main(%a: Tensor[(1, 256), float32] /* ty=Tensor[(1, 256), float32] */) -> Tensor[(1, 256), float32] {
%0 = add(%a, 1f /* ty=float32 */) /* ty=Tensor[(1, 256), float32] */;
%1 = multiply(%0, 2f /* ty=float32 */) /* ty=Tensor[(1, 256), float32] */;
%2 = multiply(%0, 3f /* ty=float32 */) /* ty=Tensor[(1, 256), float32] */;
add(%1, %2) /* ty=Tensor[(1, 256), float32] */
}
'''
可以看出:d=relay.add(a,relay.const(100,"float32"))
作为一个死代码被消除了,其中b,c都是a+1的操作进行CSE优化。
5.4 后端
将 high-level IR 变换到 low-level IR,然后进行硬件有关的优化。
在一方面,它可以直接将 high-level IR 变换到 third-party 的 tool-chains 例如 LLVM IR,来复用现有的在通用编译器上的优化和代码生成的一些基建。
另一方面,它可以利用在 DL models 和硬件特性上的先验知识来实现更高效的代码生成,通过一些 customized 的编译 passes。
4.4.1 硬件特定优化
特定于硬件的优化可以为不通用的硬件高效地生成代码。
后端优化的方法有两种方式:
(1)将低级IR转化为LLVM IR,以利用LLVM的基础设施生成优化的CPU/GPU代码。
(2)利用深度学习领域知识设定定值优化,更有效的利用目标硬件。
深度学习编译器后端优化时,为什么不直接使用LLVM呢? 深度学习编译器的后端优化时不直接使用LLVM的一个原因是,深度学习计算图(DNN模型)的特殊性和对性能的要求可能与通用编程语言的编译不完全匹配,需要定制化的优化策略和技术。
以下是一些原因:
特定硬件的优化:深度学习编译器通常需要为特定类型的硬件进行优化,如图形处理单元(GPU)、张量处理单元(TPU)或定制化的深度学习加速器。 这些硬件对于深度学习任务具有特定的计算需求,因此需要定制化的优化策略,而不仅仅是通用的LLVM优化。
高级抽象:深度学习模型的计算图通常包含高级抽象,例如卷积、池化、批归一化等操作,这些操作在底层需要进行高效的优化,但这种优化可能在通用编程语言中不太容易实现。
自动微调:深度学习编译器通常需要进行自动微调(Auto-Tuning)来找到最佳的参数配置,以满足模型和硬件的性能需求。 这种自动微调的过程通常需要与深度学习框架和硬件特性紧密结合,而不仅仅是LLVM的标准优化。
低级细节:深度学习编译器需要处理低级细节,如内存管理、数据布局和量化等,以最大程度地减少计算和内存开销。这些细节需要特定于深度学习的优化技巧。
尽管深度学习编译器的后端可以受益于LLVM的一些通用优化技术,但深度学习编译器通常需要更多的领域特定优化,以满足深度学习模型的性能需求。
因此,通常会在LLVM之上构建定制的深度学习编译器后端,以更好地适应深度学习任务的特殊性。 这些深度学习编译器后端会结合LLVM的一些优化技术,但也包括许多领域特定的优化策略,以提高深度学习模型的性能、减少计算资源的占用,并支持不同类型的硬件加速器。
由于特定于硬件的优化是针对特定硬件量身定制的,因此我们提出了现有深度学习中广泛采用的五种方法。
-
硬件内在映射:将低级IR指令映射到硬件特定的优化内核。
在TVM中硬件指令映射是通过extensible tensorization实现的,他可以声明硬件指令的行为以及指令映射的lower。
这种方法使编译器后端能够将硬件实现以及高度优化的手工micro-kernels应用于特定的patterns,从而显着提高性能。
此外,Halide/TVM 将特定的 IR 模式映射到每个架构上的 SIMD opcodes,以避免 LLVM IR 映射在遇到vector patterns时效率低下。
以下例子采用
TVM
利用x86
中的AVX512_DOT_16x4_INTRIN
指令实现矩乘加速,并且对比baseline
的执行效率。 关于TVM详细用法,在此处不作详细展开和解释。import tvm from tvm import te import numpy as np from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN, AVX512_DOT_16x4_INTRIN # baseline m, n, k = 128, 128, 128 lhs_dtype = 'uint8' rhs_dtype = 'int8' X = te.placeholder((m, k), name="X", dtype=lhs_dtype) W = te.placeholder((n, k), name="W", dtype=rhs_dtype) ak = te.reduce_axis((0, k), name="k") matmul = te.compute( (m, n), lambda i, j: te.sum(X[i, ak].astype("int32") * W[j, ak].astype("int32"),axis=ak, ), name="compute", ) func = te.create_prim_func([X,W,matmul]) sch_baseline = tir.Schedule(func, debug_mask="all") print(sch_baseline.mod.script()) '''output: # from tvm.script import ir as I # from tvm.script import tir as T @I.ir_module class Module: @T.prim_func def main(X: T.Buffer((128, 128), "uint8"), W: T.Buffer((128, 128), "int8"), compute: T.Buffer((128, 128), "int32")):T.func_attr({"tir.noalias": T.bool(True)})# with T.block("root"):for i, j, k in T.grid(128, 128, 128):with T.block("compute"):v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])T.reads(X[v_i, v_k], W[v_j, v_k])T.writes(compute[v_i, v_j])with T.init():compute[v_i, v_j] = 0compute[v_i, v_j] = compute[v_i, v_j] + T.Cast("int32", X[v_i, v_k]) * T.Cast("int32", W[v_j, v_k]) ''' ctx = tvm.cpu() mod = tvm.build(sch_baseline.mod, target="llvm -mcpu=skylake-avx512") a = tvm.nd.array(np.ones((128,128)).astype("uint8")) b = tvm.nd.array(np.ones((128,128)).astype("int8")) res = tvm.nd.array(np.zeros((128,128)).astype("int32")) mod(a,b,res) evaluator = mod.time_evaluator(mod.entry_name, ctx, number=50) print("Baseline: %f" % evaluator(a, b, res).mean) ''' Baseline: 0.000937 '''
可以看出,正常的baseline需要0.000937
m, n, k = 128, 128, 128 lhs_dtype = 'uint8' rhs_dtype = 'int8' X = te.placeholder((m, k), name="X", dtype=lhs_dtype) W = te.placeholder((n, k), name="W", dtype=rhs_dtype) ak = te.reduce_axis((0, k), name="k") matmul = te.compute( (m, n), lambda i, j: te.sum(X[i, ak].astype("int32") * W[j, ak].astype("int32"),axis=ak, ), name="compute", ) func = te.create_prim_func([X,W,matmul]) sch = tir.Schedule(func, debug_mask="all") block = sch.get_block("compute") sch.transform_layout(block, "W", lambda i, j: [i//16, j//4, i%16, j%4]) _, j, k = sch.get_loops(block) _, ji = sch.split(j, factors=[None, 16]) ko, ki = sch.split(k, factors=[None, 4]) sch.reorder(ko, ji, ki) sch.decompose_reduction(block, ko) sch.tensorize(ji, AVX512_DOT_16x4_INTRIN) print(sch.mod.script()) '''output: # from tvm.script import ir as I # from tvm.script import tir as T @I.ir_module class Module: @T.prim_func def main(X: T.Buffer((128, 128), "uint8"), W: T.Buffer((8, 32, 16, 4), "int8"), compute: T.Buffer((128, 128), "int32")):T.func_attr({"tir.noalias": T.bool(True)})# with T.block("root"):for i, j_0 in T.grid(128, 8):for j_1_init in range(16):with T.block("compute_init"):v_i = T.axis.spatial(128, i)v_j = T.axis.spatial(128, j_0 * 16 + j_1_init)T.reads()T.writes(compute[v_i, v_j])compute[v_i, v_j] = 0for k_0 in range(32):with T.block("compute_update_o"):v_i_o, v_j_o, v_k_o = T.axis.remap("SSR", [i, j_0, k_0])T.reads(compute[v_i_o, v_j_o * 16:v_j_o * 16 + 16], X[v_i_o, v_k_o * 4:v_k_o * 4 + 4], W[v_j_o, v_k_o, 0:16, 0:4])T.writes(compute[v_i_o, v_j_o * 16:v_j_o * 16 + 16])A = T.match_buffer(X[v_i_o, v_k_o * 4:v_k_o * 4 + 4], (4,), "uint8", offset_factor=1)B = T.match_buffer(W[v_j_o, v_k_o, 0:16, 0:4], (16, 4), "int8", offset_factor=1)C = T.match_buffer(compute[v_i_o, v_j_o * 16:v_j_o * 16 + 16], (16,), "int32", offset_factor=1)A_u8x4: T.uint8x4 = A[0:4]A_i32: T.int32 = T.reinterpret("int32", A_u8x4)A_brdcst: T.int32x16 = T.Broadcast(A_i32, 16)A_u8x64: T.uint8x64 = T.reinterpret("uint8x64", A_brdcst)B_i8x64: T.int8x64 = B[0, 0:64]Red: T.int16x32 = T.call_llvm_pure_intrin("int16x32", T.uint32(6900), T.uint32(2), A_u8x64, B_i8x64)C[0:16] = C[0:16] + T.call_llvm_pure_intrin("int32x16", T.uint32(6901), T.uint32(2), Red, T.Broadcast(T.int16(1), 32)) ''' ctx = tvm.cpu() AVX512_DOT_16x4_INTRIN_mod = tvm.build(sch.mod, target="llvm -mcpu=skylake-avx512") a = tvm.nd.array(np.ones((128,128)).astype("uint8")) b = tvm.nd.array(np.ones((8, 32, 16, 4)).astype("int8")) res = tvm.nd.array(np.zeros((128,128)).astype("int32")) AVX512_DOT_16x4_INTRIN_mod(a,b,res) evaluator = AVX512_DOT_16x4_INTRIN_mod.time_evaluator(AVX512_DOT_16x4_INTRIN_mod.entry_name, ctx, number=50) print("AVX512_DOT_16x4_INTRIN: %f" % evaluator(a, b, res).mean) ''' AVX512_DOT_16x4_INTRIN: 0.000026 '''
经过AVX512_DOT_16x4_INTRIN的硬件指令映射,使得计算加速36x+
-
内存分配和获取:优化内存分配和数据获取,减少内存访问延迟。
内存分配是代码生成中的另一个挑战,特别是对于 GPU 和定制加速器而言。
例如,GPU主要包含
shared memory
(内存大小有限,latency较低)和local memory
(容量大,latency较高)。这种内存层次结构需要有效的
Memory allocation
和fetching
(获取)技术来提高数据局部性。为了实现这种优化,TVM引入了内存范围的调度概念。
内存范围调度原语可以将计算
stage
记为shared
或者thread-local
。对于标记为
shared
的计算stage
,TVM 生成具有shared memory allocation和协作数据fetching的代码,这会在正确的代码位置插入memory barrier
以保证正确性。此外,TC还通过扩展PPCG编译器提供类似的功能(称为内存提升)。
然而,TC仅支持有限的预定义规则。特别的,TVM 通过内存范围调度原语在加速器中启用特殊缓冲。
-
内存延迟隐藏:通过重新排序执行流水线来隐藏内存访问延迟。
通过对执行流水线重新排序,
Memory latency hiding
(内存延迟隐藏)也是后端使用的一项重要技术。 -
循环优化:包括循环融合、滑动窗口、分块、循环重排序和循环展开。
面向循环的优化也应用于后端,以便为目标硬件生成高效代码。
由于 Halide 和 LLVM (集成了多面体方法)已经集成了此类优化技术,一些 DL 编译器在其后端利用了 Halide 和 LLVM。
应用于面向循环优化的关键技术包括:
-
loop fusion, 循环融合
循环融合是一种循环优化技术,可以将具有相同边界的循环融合在一起,以实现更好的数据重用。
对于 PlaidML、TVM、TC 和 XLA 等编译器来说,这种优化是通过 Halide 计划或多面体方法来实现的
而 Glow 则通过算子堆叠来实现循环融合。
-
sliding windows, 滑动窗口
滑动窗口是 Halide 采用的一种循环优化技术。
它的核心理念是在需要时计算数值,并在不再需要时将其存储起来以重复使用数据。
由于滑动窗口将两个循环的计算交错进行,并使其串行化,因此需要在并行性和数据重用之间做出权衡。
-
tiling, 分块
分块法将循环拆分为多个块,从而将循环分为通过分块迭代的外循环和在分块内迭代的内循环。
这种转换通过将一个分块放入硬件缓存,使分块内部的数据具有更好的定位性。
由于分块的大小取决于硬件,因此许多 DL 编译器通过自动调整来确定分块模式和大小。
-
loop reorder, 循环重排
循环重新排序(又称循环置换)改变嵌套循环中的迭代顺序,可以优化内存访问,从而提高空间局部性。
它与数据布局和硬件特性有关。不过,当迭代顺序存在依赖关系时,执行循环重排序并不安全。
-
loop unrolling, 循环展开
循环展开可以将特定循环解卷为固定数量的循环体副本,从而允许编译器应用积极的指令级并行性。
通常,loop split 与loop unroll结合使用,首先将循环拆分为两个嵌套循环,然后完全解卷内部循环。
-
-
并行化:利用硬件的多线程和SIMD并行性。
由于现代处理器普遍支持多线程和SIMD并行,编译器后端需要利用并行性来最大限度地提高硬件利用率,从而实现高性能。
Halide使用名为parallel的调度原语来指定线程级并行化的循环并行化维度,并通过映射为并行的循环维度与块和线程注释来支持GPU并行化。
此外,它还用一个 n-width vector语句取代了大小为 n 的循环,并可通过硬件内在映射将其映射到特定于硬件的 SIMD 操作码。
Stripe 开发了多面体模型的一种变体,称为嵌套多面体模型,它引入了并行多面体块作为迭代的基本执行元素。
经过这种扩展,嵌套多面体模型可以检测平铺和分层之间的层次并行化。
此外,一些 DL 编译器还依赖于 Glow 等手工库或硬件供应商提供的优化数学库。
同时,Glow 会将矢量化工作lower to LLVM,因为当提供张量维度和循环次数信息时,LLVM 自动矢量化器会工作得很好。
然而,完全由编译器后端利用并行性可以应用更多 DL 模型的特定领域知识,从而在牺牲更多工程努力的情况下获得更高的性能。
4.4.2 自动调优
深度学习模型在不同硬件平台上运行时,性能表现会有所不同。 为了充分发挥硬件潜力,开发者需要手动调整模型的参数和编译器选项,这需要大量的时间和专业知识。 Auto-Tuning 的目标是将这些优化自动化,提高深度学习模型在不同平台上的性能。
在编译器后端至关重要,可以减轻手动获取最佳参数配置的工作量。
工作原理: Auto-Tuning 的工作原理通常包括以下步骤:
a. 搜索空间定义:首先,需要定义一个搜索空间,其中包含了各种编译器选项、硬件配置和算法变体的组合。 b. 性能评估:针对搜索空间中的每个组合,使用一组性能指标来评估深度学习模型的性能。这些指标可以包括模型的训练速度、推理速度、内存消耗等。 c. 搜索算法:选择合适的搜索算法,通常采用启发式搜索方法,例如遗传算法、粒子群优化或贝叶斯优化,来探索搜索空间。 d. 自动化决策:根据性能评估的结果,自动选择最佳的编译器选项、硬件配置和算法变体组合。
由于硬件特定优化中参数调整的搜索空间巨大,因此有必要利用自动调整来确定最佳参数配置。
在本次Survey中的 DL 编译器中,TVM、TC 和 XLA 支持自动调整。
-
参数化:定义调优参数,如数据形状和硬件特性。
-
成本模型:使用机器学习模型预测不同配置的性能。
-
搜索技术:使用遗传算法、模拟退火和强化学习来搜索最优配置。
-
加速:通过并行化和配置重用来加速调优过程。
4.4.3 优化的内核库
也广泛应用于通用处理器和其他定制的深度学习加速器上,例如cuDNN,DNNL,MIOpen
-
使用高度优化的内核库(如cuDNN、MKL-DNN)来加速计算密集型操作。
up主入驻深度编译器方向了,只能说太抽象,太难了,太难了。
参考文献:
The Deep Learning Compiler: A Comprehensive Survey
<七> 深度学习编译器综述:Backend Optimizations(2) - 知乎 (zhihu.com)
感觉知乎这个大佬总结的特别好,但是我的基本功不扎实所以第一次只能到这样了,好多都是ds给我的,希望后面我能有所进步吧