深度学习编译器后端和运行时

编译器前端将用户代码解析得到计算图 IR，并且做了一些和计算设备无关的通用优化。编译器后端做的优化就和具体的设备有关了（不同设备有不同的 allocator，不同的编程模型，比如英伟达的 CUDA），后端优化更加贴合硬件，会针对硬件特点为 IR 中的计算节点选择在硬件上的算子，然后为每个算子的输入输出分配硬件内存，最终生成一个可以在硬件上执行的任务序列。

如上图所示，编译器后端位于前端和硬件驱动层之间，主要做：计算图优化、算子选择、内存分配等工作。
首先要根据硬件特性对 IR 做等价图变化，以便在硬件上找到对应的执行算子。
前端 IR 是通过解析用户代码生成的，属于一个较高的抽象层次（high-level IR），隐藏一些底层运行的细节信息，此时无法直接对应硬件上的算子（算子是设备上的基本计算序列，例如 MatMul、Convolution、ReLU 等），需要将细节信息进行展开后，才能映射到目标硬件上的算子。
对于某些前端 IR 的子集来说，一个算子便能够执行对应的功能，此时可以将这些 IR 节点合并成为一个计算节点，该过程称之为算子融合，也就是把两个或更多的算子融合为一个底层算子，只需要发一次 kernel 而不用发多次 kernel；对于一些复杂计算，后端并没有直接与之对应的算子，但是可以通过几个基本运算的算子组合达到同样的计算效果，此时可以将前端 IR 节点拆分成多个小算子，也就是对于没有底层 kernel 的实现，通过拼几个 subgraph 的方式完成等效实现，缺点就是要多发几次 kernel。
在完成计算图优化之后，就要进行算子选择过程，为每个计算节点选择执行算子。算子选择是在得到优化的 IR 图后选取最合适的目标设备算子的过程。针对用户代码所产生的 IR 往往可以映射成多种不同的硬件算子，但是这些不同硬件算子的执行效率往往有很大差别，如何根据前端 IR 选择出最高效的算子，是算子选择的核心问题。为什么要选择算子呢？因为对于同一个算子比如 Conv2d，不同的输入尺寸下在 gpu（或者其他计算设备下）可能有多个实现，这是考虑到硬件设备的优化而做的，比如小 tensor 和大 tensor 应该调用不同的 kernel，所以如何选择一个合适的 kernel 去做真正的计算也是编译器后端需要做的工作。以 MegEngine 为例，选择底层算子的方式有 heuristic 和 fastrun 两种方式，它们都是在候选集中根据一定规则选择最适合当前算子和输入的一个 kernel。编译器一般都对每一个 IR 节点提供了多个候选的算子，算子选择目标就是从中选择最优的一个算子作为最终执行在设备上的算子。在机器学习系统中，对前端生成的 IR 图上的各个节点进行拆分和融合，让前端所表示的高层次 IR 逐步转换为可以在硬件设备上执行的低层次 IR（这个过程叫做 lowering）。得到了这种更加贴合硬件的 IR 后，对于每个单节点的IR可能仍然有很多种不同的选择，例如可以选择不同的输入输出格式和数据类型，需要对 IR 图上每个节点选择出最为合适的算子，算子选择过程可以认为是针对 IR 图的细粒度优化过程，最终生成完整的算子序列。最后，遍历算子序列，为每个算子分配相应的输入输出内存，然后将算子加载到设备上执行计算。

1. 计算图优化

计算图优化是在不影响模型的数值特性的基础上，通过图变换达到简化计算、减少资源开销、适配硬件的执行能力、提升执行性能的目的。
后端的计算图优化主要是针对硬件的优化（不针对硬件的优化一般放在前端做），根据优化适用于所有硬件还是只适合特定硬件，可以分为通用硬件优化和特定硬件优化，例如为了适配硬件指令限制而做的子图变换和与特定硬件无关的算子内存 IO 优化。

1.1 通用硬件优化

通用硬件优化主要指与特定硬件类型无关的计算图优化，优化的核心是子图的等价变换：在计算图中尝试匹配特定的子图结构，找到目标子图结构后，通过等价替换方式，将其替换成对硬件更友好的子图结构。
以优化内存 IO 为例。深度学习算子按其对资源的需求可以分为两类：
计算密集型算子：这些算子的时间绝大部分花在计算上，如卷积、全连接等；
访存密集型算子：这些算子的时间绝大部分花在访存上，他们大部分是 Element-Wise 算子，例如 ReLU、Element-Wise Sum等。
在典型的深度学习模型中，一般计算密集型和访存密集型算子是相伴出现的，最简单的例子是 “Conv + ReLU”。Conv 卷积算子是计算密集型，ReLU 算子是访存密集型算子，ReLU 算子可以直接取 Conv算子的计算结果进行计算，因此可以将二者融合成一个算子来进行计算，从而减少内存访问延时和带宽压力，提高执行效率。

算子融合是 mlsys 中一个非常常见的优化计算图的手段。

1.2 特定硬件优化

特定硬件优化是指该计算图的优化是在特定硬件上才能做的优化（不同的硬件有不同的指令集、不同的编程模型以及其他限制），常见的基于硬件的优化包括由于硬件指令的限制而做的优化，特定硬件存储格式导致的优化等。

1.2.1 硬件指令限制

在一些特定的硬件上，IR 中计算节点没有直接对应的硬件算子，只能通过子图的变换来达到子图中所有算子在对应的硬件上的存在。也就是把一个没有 kernel 实现的 op 拆成几个有 kernel 的子 op，起到等价替换的效果。
举个例子，假设我们 Concat op 不支持输入个数是 100，最大支持 63，那就拆成 63 + 37 两个 op：

1.2.2 数据排布格式的限制

针对不同特点的计算平台和不同的算子，为了追求最好的性能，一般都需要选择不同的数据排布格式（Format），而这些排布格式可能跟框架缺省的排布格式是不一样的。比如有些硬件上只能跑 NCHW 格式的，有些硬件上跑 NHWC 的，这时候考虑到通用性以及执行性能就需要在 format 上做一些文章。
一般的做法是算子在执行完成后对输出插入一个格式转换操作，把排布格式转换回框架的缺省排布格式，这就引入了额外的内存操作。
举个例子i：在华为昇腾平台上 Conv 算子在输入和输出的内存排布为 5HD 时是性能最优的，所以可以看到 Conv 算子输出结果的格式是 5HD，然后通过一个转换操作转回了框架缺省的 NCHW，紧接着，后面又是一个 Conv 算子，它需要 5HD 的输入，所以又做了一个 NCHW 到 5HD 的转换。我们很容易看出，虚线框内的两个转换操作互为逆操作，可以相互抵消。通过对计算图的模式匹配，可以将该类型的操作消除，这样可以减少不必要的内存分配和计算操作。

2. 算子选择

算子选择是将 IR 图上的每个计算节点映射到设备上可执行算子的过程，一个 IR 图上的计算节点往往可以对应多个设备上的算子，这个过程中需要考虑算子的规格，算子的执行效率等问题，算子选择目标就是从中选择最优的一个算子。

2.1 算子选择的基本概念

经历了后端的图优化后，IR 图中的每一个节点都有一组算子与之对应。此时的 IR 图中的每一个节点可以认为是用户可见的最小硬件执行单元，代表了用户代码的一个操作（比如一个 Conv2d 这样的操作就没法再拆了），对于这个操作还没有具体生成有关设备信息的细节描述。
这些信息是算子选择所选择的内容信息，称之为算子信息。算子信息主要包括以下内容：
针对不同特点的计算平台和不同的算子，为了追求最好的性能，一般都需要选择不同的数据排布格式（format）。机器学习系统常见的数据排布格式有 NCHW 和 NHWC 等。
对于不同的硬件支持不同的计算精度，例如 float32、float16 和 int32 等（如果是模型量化还有更低精度）。算子选择需要在所支持各种数据类型的算子中选择出用户所设定的数据类型最为相符的算子。

2.1.1 数据排布格式（Format）

机器学习系统中很多运算都会转换成为矩阵的乘法（matmul），例如卷积运算（输入和卷积核频繁做 matmul 操作）。我们知道矩阵乘法 A × B = C 是以 A 的一行乘以 B 的一列求和后得到 C 的一个元素。

上图的上方，矩阵数据的存储是按照行优先来进行存储，虽然 B 在存储时是按照行存储，但是读取数据时却按照列进行读取，假如我们能把 B 的格式进行转换转换为列存储，例如上图的下方所示，这样就可以通过访问连续内存的方式加快数据访问速度进而提升运算速度。
由此可见不同的数据排布方式对性能有很大影响。
在机器学习系统中常见的数据格式一般有两种，分别为 NCHW 类型和 NHWC 类型。其中 N 代表了数据输入的批大小 (batch size)，C (channels)代表了图像的通道数，H (height)和 W (width) 分别代表图像输入的高和宽。
下图是 batch size 为 2、通道数为 N，长宽为 5 × 4 的 tensor 的逻辑示意图：

逻辑上看这里有多个 tensor，但是计算机的存储并不能够直接将这样的矩阵放到内存中，需要将其展平成 1 维后存储，这样就涉及逻辑上的索引如何映射成为内存中的索引，即如何根据逻辑数据索引来映射到内存中的 1 维数据索引。
对于 NCHW 的数据是先取 W 轴方向数据，再取 H 轴方向数据，再取 C 轴方向，最后取 N 轴方向（也就是倒着来的，先确定了高维度再访问元素）。其中物理存储与逻辑存储的之间的映射关系为：offsetnchw(n,c,h,w)=nCHW+cHW+h⋅W+w​
这种格式中，是按照最低维度 W 轴方向进行展开，W 轴相邻的元素在内存排布中同样是相邻的。如果需要取下一个图片上的相同位置的元素，就必须跳过整个图像的尺寸 C⋅H⋅W。
比如有 8 张 32 * 32 的 RGB 图像，此时 N=8, C=3, H=32, W=32。在内存中存储它们需要先按照 W 轴方向进行展开，然后按照 H 轴排列，这样之后便完成了一个通道的处理，之后按照同样的方式处理下一个通道。处理完全部通道后，处理下一张图片。PyTorch 和 MegEngine 框架默认使用 NCHW格式。

可以看到相当于各个通道各处理各的。
类似的 NHWC 数据格式是先取 C 方向数据，再取 W 方向，然后是 H 方向，最后取 N 方向。NHWC 是 Tensorflow 默认的数据格式。这种格式在 PyTorch 中称为 Channel-Last：
offsetnchw(n,h,w,c)=n⋅HWC+h⋅WC+w⋅C+c​
很多框架都采用上述的两种格式作为默认的数据排布格式。但是在硬件上对数据操作时，此时的数据排布可能还不是最优的。在机器学习系统中，用户输入的数据往往会远远大于计算部件一次性计算所能容纳的最大范围，所以此时必须将输入的数据进行切片分批送到运算部件中进行运算。为了加速运算很多框架又引入了一些块布局格式来进行进一步的优化，这种优化可以使用一些硬件的加速指令，对数据进行搬移和运算。这种特殊的数据格式与硬件更为贴合，可以快速的将矩阵向量化，并且极大的利用片内缓存。

2.1.2 数据精度

通常深度学习的系统，使用的是单精度(float32)表示（也叫 fp32）。这种数据类型占用 32 位内存。还有一种精度较低的数据类型为半精度(float16，也叫 fp16)，其内部占用了 16 位的内存。由于很多硬件会对半精度数据类型进行优化，半精度的计算吞吐量可以是单精度的 2∼8 倍，且半精度占用的内存更小，这样可以输入更大的批大小 (BatchSize)，进而减少总体训练时间。

Sig 代表符号位，占 1 位，表示了机器数的正负，Exponent 表示指数位，Mantissa 为尾数位。
其中 float16 类型的数据采用二进制的科学计数法转换为十进制的计算方式为：

其中如果指数位全为 0 时，且尾数位全为 0 时表示数字 0。 如果指数位全为 0，尾数位不全为 0 则表示一个非常小的数值。 当指数全为 1，尾数位全为 0 表示根据符号位正无穷大，或者负无穷大。 若指数全为 1，但是尾数位不为 0，则表示 NAN。bfloat16 并不属于一个通用的数据类型，是 Google 提出的一种特殊的类型，现在一般只在一些 TPU 上训练使用，其指数位数与 float32 的位数保持一致，可以较快的与 float32 进行数据转换。由于 bfloat16 并不是一种通用类型，IEEE 中也并没有提出该类型的标准。

2.2 算子选择的过程

基于数据格式和数据精度这两个概念，在不同硬件下会有不同的算子支持，一个硬件上支持的所有算子的集合定义为该硬件的算子信息库（也就是当前算子的一个 kernel 候选集）。算子选择过程就是从算子信息库中选择最合适算子的过程。
下图展示了一个算子选择过程：

首先，选择算子执行的硬件设备。不同的硬件设备上，算子的实现、支持数据类型、执行效率通常会有所差别。这一步往往是用户自己指定的，若用户未指定，则编译器后端会为用户匹配一个默认的设备。 然后，后端会根据 IR 图中推导出的数据类型和内存排布格式选择对应的算子。
理想情况下算子选择所选择出的算子类型，应该与用户预期的类型保持一致。但是由于软硬件的限制，很可能算子的数据类型不能满足用户所期待的数据类型，此时需要对该节点进行升精度或者降精度处理才能匹配到合适的算子。
算子的数据排布格式转换是一个比较耗时的操作，为了避免频繁的格式转换所带来的内存搬运开销，数据应该尽可能地以同样的格式在算子之间传递，算子和算子的衔接要尽可能少的出现数据排布格式不一致的现象（所以在计算之前有时候会做一些 dtype promotion 的操作）。另外，数据类型不同导致的降精度可能会使得误差变大，收敛速度变慢甚至不收敛，所以数据类型的选择也要结合具体算子分析。
总的来说，一个好的算子选择算法应该尽可能的保持数据类型与用户设置的数据类型一致，且尽可能少的出现数据格式转换。

3. 内存分配

经过计算图优化和算子选择之后，我们可以得到 IR 图中每个算子的输入输出的形状（Shape）、数据类型、存储格式。根据这些信息，计算输入输出数据的大小，并为输入输出分配设备上的内存，然后将算子加载到设备上才能真正执行计算。此外，为了更充分地例用设备内存资源，可以对内存进行复用，提高内存利用率。
内存在传统计算机存储器层次结构中有着重要的地位，它是连接高速缓存和磁盘之间的桥梁，有着比高速缓存更大的空间，比磁盘更快的访问速度。随着深度神经网络的模型越来越复杂（网络结构越来越大），AI 芯片上的内存很可能无法容纳一个大型网络模型。因此，对内存进行复用是一个重要的优化手段。此外，通过连续内存分配和 In-Place 内存分配还可以提高某些算子的执行效率。

3.1 Device 内存概念

在深度学习体系结构中，通常将与硬件加速器（主要是 GPU）相邻的内存称之为设备（Device）内存，而与 CPU 相邻的内存称之为主机（Host）内存。CPU 可以合法地访问主机上的内存，而无法直接访问设备上的内存；同理，AI 芯片可以访问设备上的内存，却无法访问主机上的内存。也就是说 Host Memory 和 Device Memory 之间是比较独立的。
在网络训练过程中，往往需要从磁盘加载数据到主机内存中，然后在主机内存中做数据处理（现在 PyTorch 和 MegEngine 也在支持用 gpu 来做数据处理），再从主机内存拷贝到设备内存中，最后设备才能合法地访问数据。算子全部计算完成后，用户要获取训练结果，又需要把数据从设备内存拷贝到主机内存中。也就是说，Host 和 Device memory 之间还是有较为频繁的通信操作的（拷贝 tensor）。

3.2 内存分配

内存分配模块主要负责给图中算子的输入、输出分配 Device 内存。用户的前端脚本经过编译器前端处理后得到 IR，后端根据 IR 进行算子选择和相关优化，可以得到算子最终的输入输出张量的形状（Shape）、数据类型（DType）、格式（Format）等信息，根据这些信息可以计算出算子输入、输出张量的尺寸大小。基本的计算方法如下：

得到张量的尺寸大小后，往往还需要对内存大小进行对齐操作。内存通常以 4 字节、8 字节或 16 字节为一组进行访问，如果被搬运的内存大小不是这些值的倍数，内存后面会填充相应数量的空数据以使得内存长度达到这些值的倍数（这个操作叫做 padding），因为访问非对齐的内存可能会更加耗时。

上图中，首先给输入张量、Conv2D 的权重和 Conv2D 的输出分配内存地址。然后为 BatchNorm 的输入分配地址时，发现 BatchNorm 的输入就是 Conv2D 算子的输出，而该张量的地址已经在之前分配过了，因此只需要将 Conv2D 算子的输出地址共享给 BatchNorm 的输入，就可以避免内存的重复申请以及内存的冗余拷贝。以此类推，可以发现整个过程中可以将待分配的内存分成三种类型：

• 一是整张图的输入张量；
• 二是算子的权重或者属性；
• 三是算子的输出张量。

三种类型在训练过程中的生命周期有所不同。
在 CPU 上常常使用 malloc 函数直接申请内存，这种方式申请内存好处是随时申请随时释放，简单易用。然而在许多对性能要求严苛的计算场景中，由于所申请内存块的大小不定，频繁申请释放会降低性能。通常会使用内存池的方式去管理内存，先申请一定数量的内存块留作备用，当程序有内存申请需求时，直接从内存池中的内存块中申请。当程序释放该内存块时，内存池会进行回收并用作后续程序内存申请时使用。 在深度学习框架中，设备内存的申请也是非常频繁的，往往也是通过内存池的方式去管理设备内存，并让设备内存的生命周期与张量的生命周期保持一致。
在 MegEngine 中，不同的训练模式会有不同的 memory 分配方式，比如 sublinear 和 DTR 的分配方式就不同。

3.3 内存复用

内存复用是指分析张量的生命周期，将生命周期结束的张量的设备内存释放回内存池并用于后续张量的内存分配。说白了就是用之后不用的 tensor 的内存（而不是去抢更多的内存）。内存复用的目的是提高内存的利用率，让有限的设备内存容纳更大的模型。
比如下图中当 BatchNorm 算子计算结束后，输出 1 不再被任何算子使用，则该张量的设备内存可以被回收，并且如果输出 1 的内存尺寸大于等于输出 3 的内存尺寸，则从输出 1 回收的地址可以用于输出 3 的内存分配，从而达到复用输出 1 地址的目的。

3.3.1 内存复用的大致流程

上图中横坐标表示张量的生命周期，图中纵坐标表示内存大小。在生命周期内，某一个张量将一直占用某块设备内存，直至生命周期结束才会释放相应内存块。通过张量生命周期和内存大小可以构造出矩形块，而内存分配要求解的目标是在内存生命周期图中容纳更多的矩形块，问题的约束是矩形块之间无碰撞。
图的左边是在未使用任何内存复用策略的情况下的内存生命周期图，此时内存同时只能容纳 T0、T1、T2、T3 四个张量。内存复用策略的求解是一个 NP 完全的问题。许多深度学习框架通常采用贪心的策略去分配内存，例如采用 BestFit 算法，每次直接从内存池中选取可以满足条件的最小内存块。

3.4 内存分配优化手段 1 —— 内存融合

上面的内存分配的方式，都是以单个张量的维度去分配的，每个张量分配到的设备地址往往是离散的。但是对于某些特殊的算子，如 AllReduce 通信算子，需要为它们分配连续的内存。通信算子的执行包含通信等待、数据搬移、计算等步骤，而在大规模分布式集群的场景下，通信的耗时往往是性能瓶颈。
针对这种场景，可以将多个通信算子融合成一个，为通信算子的输入分配连续的内存，从而减少通信的次数：

又比如分布式训练中的神经网络权重初始化，通常将一个训练进程中的权重初始化，然后将该权重广播到其他进程中。当一个网络有较多权重的时候，需要多次进行广播。通常可以为所有权重分配连续的内存地址，然后广播一次，节省大量通信的耗时。

3.5 内存分配优化手段 2 —— In-Place 算子

在内存分配流程中，会为每个算子的输入和输出都分配不同的内存。然而对很多算子而言，为其分配不同的输入和输出地址，会浪费内存并且影响计算性能。
例如优化器算子，其计算的目的就是更新神经网络的权重；例如 Python 语法中的 += 和 *= 操作符，将计算结果更新到符号左边的变量中；例如 a[0]=b 语法，将 a[0] 的值更新为 b。诸如此类计算有一个特点，都是为了更新输入的值。
下面以张量的 a[0]=b 操作为例介绍In-Place的优点。下图左边是非 In-Place 操作的实现，step1 将张量 a 拷贝到张量 a’，step2 将张量 b 赋值给张量 a’，step3 将张量 a’ 拷贝到张量 a。图右边是算子 In-Place 操作的实现，仅用一个步骤将张量 b 拷贝到张量 a 对应的位置上。对比两种实现，可以发现 In-Place 操作节省了两次拷贝的耗时（zero-copy，零拷贝），并且省去了张量 a’ 内存的申请。

4. 计算调度与执行

经过算子选择与内存分配之后，计算任务可以通过运行时完成计算的调度与在硬件上的执行。根据是否将算子编译为计算图，计算的调度可以分为单算子调度与计算图调度两种方式。而根据硬件提供的能力差异，计算图的执行方式又可以分为逐算子下发执行的交互式执行以及将整个计算图或者部分子图一次性下发到硬件的下沉式执行两种模式。

4.1 单算子调度（动态图）

单算子调度是相对于计算图而言，算法或者模型中包含的算子通过 Python 语言的运行时被逐个调度执行。例如 PyTorch 和 MegEngine 的默认执行方式 imperative（是一种动态图模式，单个 op 立即执行计算并返回结果），TensorFlow 的 eager 模式。
单算子执行的调用链路如下图所示，算子在 Python 侧被触发执行后，会经过机器学习框架初始化，其中需要确定包括算子的精度，输入与输出的类型和大小以及对应的硬件设备等信息，接着框架会为该算子分配计算所需的内存，最后交给具体的硬件计算设备完成计算的执行。

单算子调度方式的好处在于其灵活性，由于算子直接通过 Python 运行时调度，一方面可以表达任意复杂的计算逻辑，尤其是在需要复杂控制流以及需要 Python 原生数据结构支持来实现复杂算法的场景；另一方面单算子调度对于程序正确性的调试非常便利，开发人员可以在代码执行过程中打印任意需要调试的变量；最后一点是通过 Python 运行时驱动算子的方式，可以在计算中与 Python 庞大而丰富的生态库协同完成计算任务。可以发现这里其实就是动态图的优势，所以可以认为单算子调度=动态图模式。

4.2 计算图调度（静态图）

虽然单算子调度具有如上所述的优点，其缺点也很明显。一方面是难于进行计算性能的优化，原因是由于缺乏计算图的全局信息，单算子执行时无法根据上下文完成算子融合，代数化简等优化；另一方面由于缺乏计算的拓扑关系，整个计算只能串行调度执行，即无法通过运行时完成并行计算。
例如上述示例代码的计算逻辑可以表达为下图所示。由该计算图可以看出，其中乘法和减法之间并没有依赖关系，因此这两个计算可以并行执行，而这样的并行执行信息只有将计算表达为计算图后才能完成分析，这也是计算图调度相对于单算子调度的优势之一（静态图相比动态图的优势）。

在一个典型的异构计算环境中，主要存在 CPU、GPU、TPU 等多种计算设备，因此一张计算图可以由运行在不同设备上的算子组成为异构计算图。下图展示了一个典型的由异构硬件共同参与的计算图：

主流框架均提供了指定算子所在运行设备的能力，MegEngine 用户可以参考这个文档。

5. 算子编译器

作为 AI 编译器中一个重要组成部分，算子编译器把单个简单或复杂的算子经过表达和优化后编译为一个单独的可执行文件。
算子编译器，顾名思义，即对算子进行编译优化的工具。这里所谓的”算子”可以来自于整个神经网络中的一部分，也可以来自于通过领域特定语言（Domain Specific Language, DSL）实现的代码。而所谓编译，通俗来说起到的是针对目标语言进行表达和转换。
从目的上来说，算子编译器致力于提高算子的执行性能。从工程实现上来说，算子编译器的输入一般为 Python 等动态语言描述的张量计算，而输出一般为特定 AI 芯片上的可执行文件。MegEngine 目前正在实现 Functional 上的 Bytecode 虚拟机，来提高 functional 下的执行效率。

5.1 算子调度策略

算子编译器为了实现较好地优化加速，会根据现代计算机体系结构特点，将程序运行中的每个细小操作抽象为”调度策略”。
如果不考虑优化和实际中芯片的体系结构特点，只需要按照算子表达式的计算逻辑，把输入进来的张量全部加载进计算核心里完成计算，之后再把计算结果从计算核心里面取出并保存下来即可。这里的计算逻辑指的就是基本数学运算（如加、减、乘、除）以及其他函数表达式（如卷积、转置、损失函数）等。
下图是现代计算机的存储结构，越靠近金字塔顶尖的存储器造价越高但是访问速度越快。

基于这一硬件设计的事实，我们有局部性（Locality）概念：
（1）时间局部性，相对较短时间内重复访问特定内存位置。如多次访问 L1 高速缓存的同一位置的效率会高于多次访问 L1 中不同位置的效率。
（2）空间局部性，在相对较近的存储位置进行访问。比如多次访问 L1 中相邻位置的效率会高于来回在 L1 和主存跳跃访问的效率。
满足这两者任一都会有较好的性能提升。基于局部性概念，我们希望尽量把需要重复处理的数据放在固定的内存位置，且这一内存位置离处理器越近越好，以通过提升访存速度而进行性能提升。
另外，把传统的串行计算任务按逻辑和数据依赖关系进行分割后，有机会得到多组互不相关的数据，并把他们同时计算，如下图所示：

以上种种在程序实际运行的时候针对数据做出的特殊操作，统称为调度（Schedule）。调度定义了：
（1）应该在何时何处计算函数中的每个值？
（2）数据应该储存在哪里？
（3）每个值在多个消费者（Consumer）之间访存需要花费多长时间？另外在何时由每个消费者独立重新计算？这里的消费者指使用前序结构进行计算的值。
通俗理解，调度策略指的是：在编译阶段根据目标硬件体系结构的特点而设计出的一整套通过提升局部性和并行性而使得编译出的可执行文件在运行时性能最优的算法。这些算法并不会影响计算结果，只是干预计算过程，以达到提升运算速度的效果。

5.2 子策略组合优化

算子编译器的一种优化思路是：将抽象出来的调度策略进行组合，拼接排布出一个复杂而高效的调度集合。子策略组合优化，本质上还是基于人工手动模板匹配的优化方式，依赖于开发人员对于硬件架构有较深的理解。这种方式较为直接，但组合出的优化策略无法调优，同时对各类算子精细化的优化也带来较多的人力耗费。
下面以 TVM 为例，通过在 CPU 上加速优化一段实际代码，简要介绍其中几种基本调度策略组成的优化算法。
我们以形式为乘累加计算的代码为例简要分析描述这一算法。该代码的核心计算逻辑为：首先对张量 C 进行初始化，之后将张量 A 与张量 B 相乘后，结果累加到张量 C 中：

for (m: int32, 0, 1024) {
  for (n: int32, 0, 1024) {
    C[((m*1024) + n)] = 0f32
      for (k: int32, 0, 1024) {
        let cse_var_2: int32 = (m*1024)
          let cse_var_1: int32 = (cse_var_2 + n)
            C[cse_var_1] = (C[cse_var_1] + (A[(cse_var_2 + k)]*B[((k*1024) + n)]))
      }
  }
}

​
假定数据类型为浮点型（Float），此时张量 A、B、C 的大小均为 1024 × 1024，三者占用的空间共为 1024 × 1024 × 3 × sizeof(float) = 12MB。这远远超出了常见缓存的大小（如 L1 Cache 为 32KB）。因此按照此代码形式，要将整块张量 A、B、C 一起计算，只能放入离计算核更远的内存进行计算。其访存效率远低于缓存。
为了提升性能，提出使用平铺（Tile），循环移序（Reorder）和切分（Split）的调度策略。由于 L1 缓存大小为 32KB，为了保证每次计算都能够放入缓存中，我们选取因子（Factor）为 32 进行平铺，使得平铺后的每次计算时只需要关注 m.inner × n.inner 构成的小块（Block）即可，而其他的外层循环不会影响最内层小块的访存。其占用内存大小为 32 × 32 × 3 × sizeof(float) = 12KB，足够放入缓存中。以下代码展示了经过该策略优化优化后的变化:

// 由for (m: int32, 0, 1024)以32为因子平铺得到外层循环
for (m.outer: int32, 0, 32) {
  // 由for (n: int32, 0, 1024)以32为因子平铺得到外层循环
  for (n.outer: int32, 0, 32) {
    // 由for (m: int32, 0, 1024)以32为因子平铺得到内层循环
    for (m.inner.init: int32, 0, 32) {
      // 由for (n: int32, 0, 1024)以32为因子平铺得到内层循环
      for (n.inner.init: int32, 0, 32) {
        // 对应地得到相应系数
        C[((((m.outer*32768) + (m.inner.init*1024)) + (n.outer*32)) + n.inner.init)] = 0f32
      }
    }
    // 由for (k: int32, 0, 1024)以4为因子切分得到外层循环，并进行了循环移序
    for (k.outer: int32, 0, 256) {
      // 由for (k: int32, 0, 1024)以4为因子切分得到外层循环，并进行了循环移序
      for (k.inner: int32, 0, 4) {
        // 由for (m: int32, 0, 1024)以32为因子平铺得到内层循环
        for (m.inner: int32, 0, 32) {
          // 由for (n: int32, 0, 1024)以32为因子平铺得到内层循环
          for (n.inner: int32, 0, 32) {
            // 由n轴平铺得到的外轴系数
            let cse_var_3: int32 = (n.outer*32)
            // 由m轴平铺得到的外轴和内轴系数
            let cse_var_2: int32 = ((m.outer*32768) + (m.inner*1024))
            // 由m轴和n轴得到的外轴和内轴系数
            let cse_var_1: int32 = ((cse_var_2 + cse_var_3) + n.inner)
            // 这里是核心计算逻辑，划分成不同层次使得每次循环计算的数据能够放入cache中
            C[cse_var_1] = (C[cse_var_1] + (A[((cse_var_2 + (k.outer*4)) + n.inner)] * B[((((k.outer*4096) + (k.inner*1024)) + cse_var_3) + n.inner)]))
          }
        }
      }
    }
  }
}

​

5.3 调度空间算法优化

算子编译器的另外一种优化思路是：通过对调度空间搜索/求解，自动生成对应算子调度。此类方案包括多面体模型编译（Polyhedral Compilation）（基于约束对调度空间求解）和 Ansor（调度空间搜索）等。这类方法的好处是提升了算子编译的泛化能力，缺点是搜索空间过程会导致编译时间过长。 以多面体模型编译技术将代码的多层循环抽象为多维空间，将每个计算实例抽象为空间中的点，实例间的依赖关系抽象为空间中的线，主要对循环进行优化。该算法的主要思想是针对输入代码的访存特点进行建模，调整循环语句中的每一个实例的执行顺序，使得新调度下的循环代码有更好的局部性和并行性。
举个例子：

for (int i = 0; i < N; i++)
  for (int j = 1; j < N; j++)
    a[i+1][j] = a[i][j+1] - a[i][j] + a[i][j-1];

​
通过多面体模型算法先对此代码的访存结构进行建模，然后分析实例（即下图中红色圈起来的节点）间的依赖关系（即下图中箭头）：

再进行复杂的依赖分析和调度变换之后得到一个符合内存模型的最优解。如下代码显示了经过多面体模型优化后得到的结果：

for (int i_new = 0; i_new < N; i_new++)
  for (int j_new = i+1; j_new < i+N; j_new++)
    a[i_new+1][j_new-i_new] = a[i_new][j_new-i_new+1] - a[i_new][j_new-i_new] + a[i_new][j_new-i_new-1];

​
观察得到的代码，发现优化后的代码较为复杂。但是仅凭肉眼很难发现其性能优势之处。仍需对此优化后的代码进行如算法描述那样建模，并分析依赖关系后得出结论，如下图所示：经过算法优化后解除了原代码中的循环间的依赖关系，从而提高了并行计算的机会。即沿着下图中虚线方向分割并以绿色块划分后，可以实现并行计算：

5.4 芯片指令集适配

当下的AI芯片中，常见的编程模型分为：
单指令多数据（Single Instruction, Multiple Data, SIMD），即单条指令一次性处理大量数据：

单指令多线程（Single Instruction, Multiple Threads, SIMT），即单条指令一次性处理多个线程的数据：

一般来说，算子编译器在具体的代码中会按照：前端、中端、后端，逐渐差异化的思路进行实现。即在前端设计中兼容多种不同芯片后端的指令集，以帮助编译器用户（即 AI 程序员）不需要在乎芯片差异，而只需要专注在 AI 算法逻辑上即可；在中间表示（IR）设计中对不同芯片的体系结构进行区分，从而可以实现对不同芯片进行不同的优化方法；在后端的目标代码生成部分对各个芯片的不同指令集详细区分，以保证生成出的目标代码能够顺利运行在目标芯片上。

5.5. 算子表达能力

算子表达能力指的是：算子编译器的前端识别输入代码，并在不损失语义信息的情况下转换为 IR 的能力。
算子编译器承接的前端输入往往是 PyTorch 等的 Python 形式的代码，而 Python 中各种灵活的表达方式（包括而不限于索引、View 语义等）对算子编译器的前端表达能力提出了较高要求。另外在检测网络中，输入算子往往还有大量的控制流语句。此外，还经常可以看到神经网络中存在许多的动态形状问题，即网络中的算子形状会受网络迭代次数和控制流等条件的影响。这些都对算子编译器前端的表达能力提出了很高的要求。
在实际工程实践中，发现大量的长尾分布般不常见但性能很差的算子往往是整体网络训练或推理的瓶颈点。而这些长尾算子大都是由于其出现频次低而不至于实现在计算库中。同时其语法过于灵活或存在大量的控制流语句以及动态形状问题而难以被目前的算子编译器前端充分表达出来，因此也难以通过算子编译器进行优化加速。于是，这些长尾算子只好以运行速度较慢的 Python 解释器或者虚拟机的方式执行，从而成为整个网络中的性能瓶颈。此时，提高算子编译器前端的表达能力就成为了重中之重。

6. 参考

• 机器学习系统：设计和实现
• Training Deep Nets with Sublinear Memory Cost
• Dynamic Tensor Rematerialization
• MegEngine 官方文档：Tensor 所在设备
• TVM 官方文档：How to optimize GEMM on CPU