深度神经网络(DNN)编译器原理简介
目录
1 什么是DNN编译器
**编译器(compiler)**在计算机编译代码的时候往往指一种程序，会将某种编程语言的源代码（原始代码）转换成另一种编程语言（目标语言），转换的过程中进行的程序优化就是编译优化。
Fig. 1 LLVM编译过程（例如cpp编程IR中间表示，然后中端优化，最后给后端编译给机器执行）
随着xxxx，现在的model要部署在各类计算设备和不同的硬件架构上(x86, ARM, RISC-V)；实际部署对性能需要更多要求，性能的优化能够为算法提供更大的空间。这些需求在通用的计算框架中以及难以满足，深度学习的计算任务在现有的计算框架中以DSL(Domain Specific Language)的方式进行编程和表达，使得深度学习计算任务的优化和执行符合传统计算机语言的编译和优化过程。
因此，DNN编译就是将当前的深度学习计算任务通过一层或多层中间表达进行翻译和优化，最终转化成目标硬件上的可执行代码的过程。
Fig 2. DNN编译器流程，也由前端，后端，中间表达和优化过程构成
2 前端
DNN编译器的前端shared深度学习框架的前端表达，例如TF和torch，一般都是Python的DSL(Domain Specific Language)。当中的基本数据结构是Tensor(张量)，拿来描述基本数据类型(int, float, double, string等)构成的高维度数组。在Tensor上进行的基本计算操作称为算子(operator)，通常是基本的线性代数计算组成，如矩阵乘法，向量加减法。如下：
一些常用算子
Python作为胶水语言，把一个深度学习的计算模型描述成一系列算子的操作。
3 后端
DNN编译器的**后端指最终转换后的代码要执行的设备或神经网络加速器，目前常见的支持深度学习的计算设备有CPU、GPU、FPGA、TPU等其它专用加速器。**不同设备的不同芯片架构，对应的编程模型和优化方法都不相同。例如：GPU采用多个并行的SM(streaming multiprocessor)和共享内存的架构。（在GPU上的执行要符合SIMT的计算模型）；CPU采用多核架构以及多线程模型（如线程池）来实现高性能的计算任务。
4 中间表达（Intermediate Representation，后文用IR代替）
前端到后端的过程需要经过若干个IR。**目前DNN中常见的IR主要是计算图（CG）和算子表达式。计算图作为连接深度学习框架和前端语言的主要格式，也是标准化深度学习计算模型的常用格式，**如ONNX格式即为一种深度学习模型的标准可交换格式，目前主流框架如TensorFlow和PyTorch的大部分程序都可以被转换或导出成ONNX格式。
除了计算图外，算子继续向下转换到下层并生成设备代码的时候，算子表达式(Tensor Expression)作为另一类IR被用于DNN编译器中，如TVM，Ansor，Tensor Comprehension等。算子表达式等价于算子的计算逻辑，从而可以被编译器进一步变成Device的可执行代码。
C = t.compute((m, n),
lambda i, j: t.sum(A[i, k] * B[j * k], axis = k)
5 优化过程
DNN编译器最核心的就是优化过程(Optimization Pass)，定义在**每种IR上的函数，它的输入是某一种IR，经过一系列优化后输出一个被优化后的IR。**如计算图上的经典优化过程，**常数传播、公共子表达式消除等，这个过程对输入的计算图等价变换为新的计算图。**也有一些优化过程是将高层的IR转换为底层的IR，甚至直接变成device上的可执行代码。综上优化过程分为：设备相关 和 设备无关 的优化。
6 计算图优化
计算图是深度学习框架和前端语言的中间表达，它的节点是算子，边是张量，所有节点和边构成一个有向无环图（DAG），节点之间的依赖关系表示每个算子的执行顺序。
Fig 3. 计算图示例
计算图的优化是**通过一系列等价或者近似的优化操作将输入的计算图变换为一个新的计算图。其目标是通过这样的图变换来化简计算图，从而降低计算复杂度或内存开销。**常见的图优化方法：
6.1 表达式化简
计算图的子图部分所对应的算术表达式，在数学上化简为更简单的子图。下图展示了一个利用算术表达式化简计算图的例子，左边的子图包含了两个算法：Const算子（返回元素值为0的常量张量）和Mul算子（计算两个相同形状的算子的元素乘积），通过表达式化简，这个子图可以直接被化简成右边的只包括Const算子的子图。
6.2 公共子表达式消除
英文为Common Subexpression Elimination。其是通过找到程序中等价的计算表达式，然后通过复用结果的方式消除其他的冗余表达式的计算。一个简单的实现：按照图的拓扑顺序（保证前继节点已访问）遍历图中的节点，每个节点按照输入张量和节点类型组合作为键值进行缓存，后续节点如果由相同的键值就可以消除，并将其输入边连接到缓存的节点的输入上。
代码表示的话就是：
a = b * c + g;
d = b * c + f;
当中的b * c要计算两次，我们把b * c的结果保存就只要计算一次，减少开销。
tmp = b * c;
a = tmp + g;
d = tmp + f;
6.3 常数传播
常数传播（constant propagation）就叫常数折叠（constant folding），是通过编译期计算出常数表达式的值，然后计算出的值替换原来表达式。在计算图中，一个节点的所有输入张量都是常数张量的话，那么这个节点就可以在编译期计算出输入张量，并替换为新的常数张量。
值得注意的是，常数传播需要编译器具有计算的能力，甚至对于一些较大的算子还需要能够在加速硬件上（如GPU）上计算，否则优化的过程就会非常的慢。常数传播的优化在深度学习尤其是模型推理的时候非常有用，因为在推理时，模型中的参数张量全部固定为常数张量，大量计算可以在编译期计算好，极大的化简了推理运算时的计算开销。但是，在深度学习的场景中，常数传播有时候也会带来否优化，如增加计算内存甚至计算时间，一个典型的例子就是一个标量常数张量后面跟一个Broadcast的算子时，如果做了常数传播就会增加内存占用，如果后面是访存密集型的算子的话，也会增加内存压力，从而增加计算时间。
6.4 矩阵乘自动融合
在同一个网络里，经常会出现形状相同的矩阵乘法，根据一些矩阵的等价规则，如果把些矩阵乘算子融合成一个大的矩阵乘算子，可以更好的利用GPU的算力，从而加速模型计算。
如图，两个矩阵乘法共享同一个输入张量（图中方框内左侧），我们就可以自动把其中的两个输入张量拼接成一个大的矩阵乘算子（图中方框内右侧），其计算的结果刚好是原算子计算结果的拼接。
利用这种规则，图中最右侧的GRU网络中的两组矩阵乘算子可以分别融合成两个大的矩阵乘算子。类似的融合规则还有BatchMatMul，可以把两个相同形状的矩阵拼接成一个新的BatchMatMul算子。
6.5 算子融合
矩阵乘的自动融合是针对矩阵乘算子特有的，其他的小算子融合成大算子的效果也如此。可以提高GPU的利用率，减少内核启动开销、减少访存开销等好处。Element-wise的算子（如Add，Mul，Sigmoid，Relu等）其计算量非常小，主要计算瓶颈都在内存的读取和写出上，如果前后的算子能够融合起来，前面算子的计算结果就可以直接被后面算子在寄存器中使用，避免数据在内存的读写，从而提交整体计算效率。
在没有融合前，执行两个算子需要启动两个GPU内核，前一个计算的结果需要写出到主存中，下一个内核计算的时候需要再次读取到计算核上。然后，融合后的代码只需要启动一个内核，并且可以有效复用中间计算结果。
//融合前为两个单独内核函数
__global__ mul(float *x0, float *x1, float *y)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
y[idx] = x0[idx] * x1[idx];
}
__global__ add(float *x0, float *x1, float *y)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
y[idx] = x0[idx] + x1[idx];
}
//融合后为一个单独内核函数
__global__ fused_muladd(float *x0, float *x1, float *x2, float *y)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
y[idx] = x0[idx] * x1[idx] + x2[idx];
}
6.6 子图替换和随机子图替换
算子融合在深度学习计算中能够带来较好的性能优化，然而在实际的计算图中有太多算子无法做到自动的算子融合，主要原因包括算子的内核实现逻辑不透明、算子无法在特有加速器上融合等等。为了在这些的情况下还能进行优化，用户经常会实现一些手工融合的算子来提升性能。那么，编译器在计算图中识别出一个子图并替换成一个等价的新的算子或子图的过程就是子图替换优化。
基于规则的子图替换示例，需要在系统中注册系列替换规则，如Conv和Relu的子图可以替换为Conv+Relu融合后的算子。
7 内存优化
除了上述的计算性能之外，DNN编译器对深度学习的计算任务的内存占用优化也是一个重要的目标。
内存占用包括：输入数据，中间计算结果和模型的参数。推理的场景中，前面的算子计算完的中间结果所占用的内存，后面的算子都可以复用。但是训练场景，反向求导计算需要使用前向输出的中间结果。因此，前面计算的算子需要一直保留到对应的反向计算结束后才能释放，对整个计算任务的内存占用挑战比较大。
所幸的是，在计算图中，这些所有的数据都被统一建模成计算图中的张量，都可以表示成一些算子的输出。计算图可以精确的描述出所有张量之前的依赖关系以及每个张量的生命周期，因此，根据计算图对张量进行合理的分配，可以尽可能的优化计算内存的占用。
根据计算图优化内存分配的例子
在上图中，默认的执行会为每一个算子的输出张量都分配一块内存空间，假设每个张量的内存大小为N，则执行该图需要4N（x + a + b + c）的内存。但是通过分析计算图可知，其中的张量a可以复用张量x，张量c可以复用a，因此，总的内存分配可以降低到2N。
7.1 基于拓扑序的最小内存分配
计算图的张量内存分配包括：张量生命期的分析 和 **内存分配。**首先，给定计算图后，唯一决定张量生命期的就是节点（算子）的执行顺序。在计算框架中，执行顺序是由运行时决定的，所以内存也都是运行时分配的。但是在编译器中，我们可以通过生成固定顺序的代码来保证最终的节点以确定的顺序执行，因此在编译器就可以为所有张量决定内存分配的方案。一般只要以某种拓扑序遍历计算图就可以生成一个依赖正确的节点执行顺序，如BFS、Reverse DFS等，进而决定每个张量的生命期，即分配和释放的时间点。
然后，根据每个张量的分配和是否顺序分配对应的内存空间，使得总内存占用最小。一种常用的内存分配方式就是建立一个内存池，由一个块内存分配管理器（如BFC内存分配器）管理起来，然后按着每个张量的分配和释放顺序依次向内存池申请和释放对应大小的内存空间，并记录每个张量分配的地址偏移。**当一个张量被释放回内存池时，后续的张量分配就可以自动复用前面的空间。当所有张量分配完时，内存池使用到的最大内存空间即为执行该计算图所需要的最小内存。在真实的运行时，我们只需要在内存中申请一块该大小的内存空间，并按照之前的记录的地址偏移为每个张量分配内存即可。**这样即可以优化总内存的占用量，也可以避免运行时的内存分配维护开销。 值得注意的是，不同拓扑序的选择会同时影响模型的计算时间和最大内存占用，同时也强制了运行时算子的执行顺序，可能会带来一定的性能损失。
7.2 张量换入换出
上面的方法只考虑了张量放在GPU的内存中，当实际内存还是不够的时候，可以将一部分张量放到host端的内存中（如CPU），等需要的时候再移动回GPU的内存即可。但是需要考虑CPU和GPU直接内存拷贝的延迟和带宽，需要合理的将内存的读写与其他算子的计算重叠起来。
给定上述假设以及必要的数据（如每个内核的执行时间、算子的执行顺序等），关于每个张量在什么时间放在什么地方的问题就可以被形式化的描述成一个最优化问题。AutoTM就提出了一种把计算图中张量在异构内存环境中的问题建模成一个整数线性规划的问题并进行求解。
上图展示了一个利用整数线性规划优化计算图内存分配的优化空间示例，图中每一行表示一个张量，每一列表示算子的执行顺序。每一行中，黄色Source表示张量的生成时间，紫色的SINK表示张量被消费的时间，每个张量都可以选择是在内存中（DRAM）还是外存（PMM）中。那么问题优化目标为就是给定任意的计算图最小化其执行时间，约束为主存的占用空间，优化变量就是决定放在哪个存储中，在有限的节点规模下，这个问题可以通过整数线性规划模型求解。
7.3 张量重计算
深度学习计算图的大多算子都是确定性的，即给定相同的输入其计算结果也是相同的。因此，我们可以进一步利用这个特点来优化内存的使用。当我们对连续的多个张量决定换入换出的方案时，如果产生这些张量的算子都具有计算确定性的话，我们可以选择只换出其中一个或一少部分张量，并把剩下的张量直接释放，当到了这些张量使用的时机，我们可以再换入这些少量的张量，并利用确定性的特点重新计算之前被释放的张量，这样就可以一定程序上缓解CPU和GPU之前的带宽压力，也为内存优化提供了更大的空间。如果考虑上换入换出，内存优化方案需要更加仔细的考虑每个算子的执行时间，从而保证重计算出的张量在需要的时候能及时的计算完成。
8 内核优化
在计算图上完成前面的编译优化之后，就需要继续向下编译。其中最主要问题是：如何对计算图中的每一个算子生成相应的代码。在计算框架中，每个算子都是预先实现并注册到框架中的，这样计算图在执行时只需要调用相应的代码即可。然而，计算框架的缺点是无法快速适配到一个新的硬件上，其需要为每一种硬件都实现一套算子代码，这不仅需要大量人力和时间成本，并且算子实现的性能也无法得到保证。因为，在对每个后端平台针对每个算子实现内核代码的时候都需要考虑不同的编程模型、数据排布、线程模型、缓存大小等等因素。
为了解决这个问题，就有了张量编译（或算子编译）的研究工作以及张量编译器。**算子编译的核心思想是：先找到能够描述通用算子与硬件无关的计算逻辑表示，然后编译器根据这个逻辑表示结合具体的硬件生成相应的内核代码。**相关的研究工作有TVM，Halide，TACO，Tensor Comprehension，FLEXTensor等。
8.1 算子表达式
大多数算子，其计算逻辑都可以描述为：针对输出张量中的每个元素的独立同构计算。以矩阵乘C= AB为例，矩阵C中的每个元素C[i][j]都是由A矩阵的 i 行和B矩阵的 j 列内积得到。所以，大多数算子的计算逻辑都要以描述其中的元素的计算逻辑来表示，这就是i算子表达式的作用。
一个算子表达式包括：
- 所有输入和输出的张量
- 输出张量的计算形状
- 输出张量中的每个元素的计算表达式（包括元素在张量中的位置参数，一般以lamda表达式的形式描述为坐标参数的匿名函数）
| 算子 | 算子表达式 |
|---|---|
| 矩阵乘 | C = t.compute((m, n), lambda i, j: t.sum(A[i, k] * B[k, j]), axis=k) |
| 仿射变换 | C = t.compute((m, n), lambda i, j: C[i, j] + bias[i]) |
| 卷积 | C = t.compute((c, h, w), lambda i, x, y: t.sum(data[kc, x+kx, y+ky] * w[i, kx, ky]), axis=[kx, ky, kc]) |
| ReLU | C = t.compute((m, n), lambda i, j: t.max(0, A[i, j]) |
8.2 算子表示与调度逻辑的分离
有了算子表达式之后，我们就得到了一个算子的计算逻辑。为了生成硬件上的最终代码，我们需要把算子表达式的逻辑计算变化成符合硬件编程模型的代码，并考虑硬件特性进行代码优化，这个过程就叫作表达式的调度（Schedule）。 通常来说，一个最简单的调度方案就是通过生成多重循环来遍历一个算子表达式中输出张量中的每一个元素，然后调用其提供的lambda函数，即可完成一个简单的内核代码的生成。下面的代码展示了一个简单的张量加算子的表达式，以及为其在TVM中创建一个默认调度的示例（上半部分），同时调度后产生出的内核代码（下半部分）。
# 在TVM中创建一个默认调度的示例
C = tvm.compute((n,), lambda i : A[i] + B[i])
S = tvm.create_schedule(C.op)
// 调度后生成的内核代码
for (int i = 0; i < n; i ++) {
C[i] = A[i] + B[i];
}
可以看到，上面生成的内核代码只是一个简单的循环，实际中这样的代码往往性能不好。我们希望对上述循环进行一系列的变化，如把一个循环拆分成两重循环、或者把两个循环合并一个循环、或者把两个循环的顺序颠倒等等。为了方便这些优化，算子编译器也提供了一些相应的调度操作接口，如下图中的split操作即可以上述循环按照32为因子进行拆分成一个两重循环。
# 在TVM中创建一个默认调度的示例
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
# 在TVM中按照32为因子进行拆分成一个两重循环
xo, xi = s[C].split(s[C].axis[0], factor = 32)
// 调度后产生出的内核代码
for (int xo = 0; xo < ceil(n /32); ++xo)
{
for (int xi = 0; xi < 32; ++xi)
{
int i = xo * 32 + xi;
if (i < n)
C[i] = A[i] + B[i];
}
}
除了优化，我们还希望一个算子表达式能生成特定硬件上符合其编程模型的代码。这就需要我们能针对这些硬件提供一些调度操作。例如，当我们想让上述代码能在CUDA GPU上执行，我们就需要把一些循环绑定到CUDA编程模型中的threadIdx或blockIdx上，同样，我们可以使用算子编译器中的bind接口来完成，如下代码所示，最终我们就可以得到一个简单的可以GPU执行的内核代码。
# 在TVM中创建一个默认调度的示例
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
# 在TVM中按照32为因子进行拆分成内个两重循环
xo, xi = s[C].split(s[C].axis[0], factor = 32)
# 使用bind接口来完成和threadIdx或blockIdx的绑定
S[C].reorder(xi, xo)
s[C].bind(xo, tvm.thread_axis("blockIdx.x"))
s[C].bind(xi, tvm.thread_axis("threadIdx.x"))
// 调度后产生出的内核代码
int i = threadIdx.x * 32 + blockIdx.x;
if (i < n)
{
C[i] = A[i] + B[i];
}
8.3 自动调度搜索与代码生成
有了算子表达式和对表达式的调度机制，我们就可以较容易的在一个新的硬件设备上生成一个算子的内核代码了。然而，我们可以看到，在调度的时候，有非常多种决定需要抉择，而且这些决定都会根据硬件的不同而产生不一样的性能影响，这些都需要经验非常丰富的专家才能知道一个较好的调度方案。为了进一步克复这个问题，一类利用机器学习进行自动调度搜索的方法被广泛应用。
给定一个算子表达式，我们首先需要针对该表达式自动生成出一个调度的代码模板，模板中可以预留出大量的可配置的参数。生成的模板需要能够尽可能包括各种代码的可能性，也就是保证足够大的搜索空间。给定了代码模板后，剩下的事情就是决定哪一个配置可以生成最优的代码，实际中，一个代码模板可能有成千上万种可选配置，**因此，一般的编译器会采用机器学习的方法通过不断尝试，生成代码、测量性能、反馈给机器学习模型、再生成下一个（一批）代码的方式不断迭代搜索，直到搜索到一定的步数后找到一个较优的代码配置，并生成最终代码。**通过机器学习的方法的好处是可以针对特别的问题输入和硬件，利用黑盒的方式找到一个较好的专用代码，但其缺点也很明显，在编译的过程中需要大量的编译和尝试，需要花费较长的编译时间和较多的算力。
9 跨算子的的全局调度优化
前面的优化和算子生成分别在计算图和算子表达式两个层次完成。这种分层的优化给编译器的设计和实现带来更清楚的模块化和可维护性，但是同时也由于上下层的分离损失了一些更进一步的优化机会，例如硬件利用率低，无法完全发挥硬件的计算性能。造成这些低效的主要原因包括：
- 单个Op的调度时间与计算时间相比不可忽略，造成较大的调度开销；
- OP的并行度不足以占满GPU的计算核心。
9.1 任意算子的融合
为了解决上述问题，一个很自然的想法就是：能不能对任意的算子进行融合，从而提高硬件利用率，降低算子的调度开销。一种最简单的方法就是实现更加激进的自动算子融合。如下图所示的为一个简单的计算图，与前面的算子融合不同的是，我们为了对任意算子做融合，引入了非element-wise算子，如矩阵乘法，这会给之前的融合方法增加难度。
为了实现任意算子融合，我们需要每一个算子的内核函数，如该例子中需要的Sigmoid、Relu和MatMul算子，下面的代码为这三个算子的内核函数：
__device__ float sigmoidf (float in) {
return 1.f / (1.f + expf(-in));
}
__device__ float reluf (float in) {
return fmaxf(0.f, in);
}
__device__ float Matmul (float *a, float *b, float *c, int m, int n, int k) {
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < m && col < k) {
float tmp = 0.f;
for (int i = 0; i < n; i++) {
tmp += a[row * n + i] * b[i * k + col];
}
c[row * k + col] = tmp;
}
}
为了按照前面讲到的的方法进行这三个算子的融合，我们需要将上述三个函数生成到同一个全局核函数内，如下图示例。值的注意的是，为了保证任意算子之前有正确的数据依赖，我们有时候需要在两个算子之间插入一个全局的数据同步。如下代码中的kernel_0是将三个融合后：
#include <cuda_runtime.h>
#include <iostream>
#include <chrono>
// Sigmoid 激活函数
__device__ float sigmoidf(float in) {
return 1.f / (1.f + expf(-in));
}
// ReLU 激活函数
__device__ float reluf(float in) {
return fmaxf(0.f, in);
}
// 矩阵乘法函数
__global__ void Matmul(float *a, float *b, float *c, int m, int n, int k) {
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < m && col < k) {
float tmp = 0.f;
for (int i = 0; i < n; i++) {
tmp += a[row * n + i] * b[i * k + col];
}
c[row * k + col] = tmp;
}
}
// Sigmoid 算子
__global__ void Sigmoid(float *A, float *C, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
C[idx] = sigmoidf(A[idx]);
}
}
// ReLU 算子
__global__ void ReLU(float *A, float *C, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
C[idx] = reluf(A[idx]);
}
}
// 融合三个算子后的全局内核函数
__global__ void kernel_0(float *A, float *B, float *C, int m, int n, int k) {
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
// 计算Sigmoid算子
if (row < m && col < n) {
float sigmoid_val = sigmoidf(A[row * n + col]);
A[row * n + col] = sigmoid_val; // 原地更新A
}
__syncthreads(); // 确保所有线程完成Sigmoid计算
// 计算矩阵乘法算子
if (row < m && col < k) {
float tmp = 0.f;
for (int i = 0; i < n; i++) {
tmp += A[row * n + i] * B[i * k + col];
}
C[row * k + col] = tmp;
}
__syncthreads(); // 确保所有线程完成矩阵乘法计算
// 计算ReLU算子
if (row < m && col < k) {
C[row * k + col] = reluf(C[row * k + col]);
}
}
int main() {
int m = 1024, n = 1024, k = 128;
size_t size_A = m * n * sizeof(float);
size_t size_B = n * k * sizeof(float);
size_t size_C = m * k * sizeof(float);
float *h_A = (float *)malloc(size_A);
float *h_B = (float *)malloc(size_B);
float *h_C = (float *)malloc(size_C);
// 初始化h_A和h_B
for (int i = 0; i < m * n; i++) {
h_A[i] = static_cast<float>(rand()) / RAND_MAX;
}
for (int i = 0; i < n * k; i++) {
h_B[i] = static_cast<float>(rand()) / RAND_MAX;
}
float *d_A, *d_B, *d_C;
cudaMalloc(&d_A, size_A);
cudaMalloc(&d_B, size_B);
cudaMalloc(&d_C, size_C);
cudaMemcpy(d_A, h_A, size_A, cudaMemcpyHostToDevice);
cudaMemcpy(d_B, h_B, size_B, cudaMemcpyHostToDevice);
dim3 blockDim(16, 16);
dim3 gridDim((m + blockDim.x - 1) / blockDim.x, (k + blockDim.y - 1) / blockDim.y);
// 测试分开执行三个算子的耗时
float *d_temp1, *d_temp2;
cudaMalloc(&d_temp1, size_A);
cudaMalloc(&d_temp2, size_C);
// 测试融合内核的耗时
cudaEvent_t start1, stop1;
cudaEventCreate(&start1);
cudaEventCreate(&stop1);
cudaEventRecord(start1);
Sigmoid<<<(m * n + 255) / 256, 256>>>(d_A, d_temp1, m * n);
Matmul<<<gridDim, blockDim>>>(d_temp1, d_B, d_temp2, m, n, k);
ReLU<<<(m * k + 255) / 256, 256>>>(d_temp2, d_C, m * k);
cudaEventRecord(stop1);
cudaEventSynchronize(stop1);
float separate_time = 0;
cudaEventElapsedTime(&separate_time, start1, stop1);
std::cout << "Separate kernels time: " << separate_time << " ms" << std::endl;
// 测试融合内核的耗时
cudaEvent_t start2, stop2;
cudaEventCreate(&start2);
cudaEventCreate(&stop2);
cudaEventRecord(start2);
size_t sharedMemSize = (m * n + m * k) * sizeof(float);
kernel_0<<<gridDim, blockDim, sharedMemSize>>>(d_A, d_B, d_C, m, n, k);
cudaEventRecord(stop2);
cudaEventSynchronize(stop2);
float fused_time = 0;
cudaEventElapsedTime(&fused_time, start2, stop2);
std::cout << "Fused kernel time: " << fused_time << " ms" << std::endl;
// 释放内存
cudaFree(d_A);
cudaFree(d_B);
cudaFree(d_C);
cudaFree(d_temp1);
cudaFree(d_temp2);
free(h_A);
free(h_B);
free(h_C);
return 0;
}
上面的实现可以看出有很大的局限性和问题，例如这种方法打破了现有的模块化设计，内核融合过程需要对每个算子的内核函数有一定要求，并需进行二次修改，还需要获取其一些额外的隐式参数：如threadBlock的个数，大小等。更进一步，这种方法也引入了一些“非标准”的GPU用法，如在kernel内部做全局同步可能会引入**死锁（指程序中某些线程由于等待某些条件的满足而陷入永久等待的状态，导致程序无法继续执行）**的问题 。 尽管有学术界可以使用持久化线程（Persistent threads）的方法来实现同步，但这种方法和GPU有较强的绑定，无法把优化过程通用化到其它硬件上，有大量GPU相关的实现细节混在其中。
9.2 编译时全局算子调度
为了更好的解决这种问题，我们就是需要一种能根据计算流图中的并行度以及算子内部的并行度的整体信息来进行一个全局的任务调度方法。本章中以Rammer的技术为例，介绍一种全局算子调度的优化方法。
首先，在计算表达层，为了打开现有算子的黑盒实现，Rammer引入rOperator来代替原有算子抽象，暴露出每一个算子中的所有并行任务（rTask）。 在硬件层，引入虚拟设备（vDevice）的抽象，并提供计算单元（vEU）级别的的粒度调度接口，中以允许将一个rTask调度到任意指定的vEU上。然而，rTask粒度的调度可能带来更严重的调度开销，Rammer利用DNN计算性能有较强的确定性，即算子的计算时间在数据流图运行前就可以通过测量得到。因此，在编译时可以将整个数据流图的所有rTask静态的编排成一个确定性执行方案，通过vDevice映射到物理Device的执行单元进行执行。
这种全局调度的抽象解耦了调度机制与优化策略，通过暴露出rTask粒度的调度接口，从而可以基于该接口设计任意的编排方案来优化整体性能。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
