oneDNN Graph Compiler: A Hybrid Approach for High-Performance Deep Learning Compilation

oneDNN Graph Compiler：一种用于高性能深度学习编译的混合方法

作者/机构
Jianhui Li, Zhennan Qin, Yijie Mei, Jingze Cui, Yunfei Song, Ciyong Chen, Yifei Zhang, Longsheng Du, Xianhang Cheng, Baihui Jin, Yan Zhang, Jason Ye, Eric Lin, Dan Lavery
(Intel, US & China)

A1 主要贡献

随着深度学习（DL）模型和硬件的快速发展，DL工作负载的特性已从少数计算密集型操作的热点，转变为分散在模型中的广泛操作。仅仅使用专家调优的算子库来加速少数计算密集型操作，已无法充分发挥AI硬件的性能潜力。为了应对这一挑战，本文提出了oneDNN Graph Compiler，这是一个采用混合方法的张量编译器，它结合了编译器优化技术和专家调优的内核，旨在为深度神经网络（DNN）图生成高性能代码。

核心问题与研究目标:
当前DL框架通常逐个操作地执行计算图，并依赖算子库来加速关键操作。然而，随着模型演进（如自然语言处理和推荐系统模型）和硬件发展，内存密集型操作的比例显著增加，这使得传统的基于算子库的方法性能受限。张量编译器（如XLA、Triton、TVM）应运而生，它们通过循环变换、并行化、向量化等技术来优化整个计算图，但要在性能上媲美专家手工调优的内核仍是一个巨大的挑战，通常需要复杂的自动调优过程。

oneDNN Graph Compiler的目标是，在自动编译过程中融入专家在开发高性能内核时积累的领域知识，从而为整个DNN计算图生成代码，使其性能不仅能与专家调优的算子相媲美，甚至在图级别上超越它们。

创新点与主要贡献:
本文的主要贡献如下：
1. 提出了一种包含两级中间表示（IR）的张量编译器设计。
* Graph IR：支持图级别的转换，如低精度计算转换、常量权重预处理、内存布局传播以及细粒度融合区域的形成。
* Tensor IR：支持为细粒度融合生成parallel-for循环，将多个融合操作（Fused OP）合并为一个并行循环以实现粗粒度融合，以及内存缓冲区优化。
2. 引入了一种基于模板的计算密集型操作降级（lowering）方法。
* 该方法实现了从业界专家调优内核中学到的最优算法。通过使用微内核（microkernel）和分块布局（blocked layout），它避免了冗长且困难的编译器变换过程，就能生成专家级别的算子性能。
3. 引入了融合操作模板（fused operation template）以生成高效的细粒度融合代码。
* 该模板为融合提供了多个“锚点”（anchor points），使得编译器可以选择最佳位置来组合一个与邻近可融合操作循环合并的parallel-for循环。这使得编译器能够支持包括归约（reduction）和重排序（reordering）操作在内的多种灵活融合，而无需在低级IR中使用传统的编译器技术来处理复杂的并行循环合并问题。

通过结合编译器技术和内核库技术，并专注于深度学习领域的特定优化问题（如低精度计算、激进的图操作融合、静态张量形状和内存布局优化、常量权重优化和内存缓冲区重用），oneDNN Graph Compiler在CPU上为性能关键的DNN计算图和端到端模型带来了显著的性能提升。

A2 方法细节

II. 高层设计

关键设计选择。我们做出了一些重要的设计选择，使得我们能够以可控的开发工作量复制专家调优内核的性能，并进一步获得卓越的图级性能。首先，我们不将图IR降级为通用的嵌套循环表示并应用高级循环变换技术，而是使用模板来机械地生成计算密集型内核及其与邻近内存密集型操作融合的代码。其次，我们不使用多级IR和渐进式降级，而是选择了两级IR：图IR和张量IR。图IR保留操作（OP）的语义并支持图优化，而张量IR则抽象了硬件目标并支持低级优化。第三，我们使用微内核来隐藏性能最佳的矩阵指令序列的实现细节。模板和微内核继承了专家调优内核的算法和实现，而图IR和张量IR则支持图级优化和在其他硬件目标上复用的可能性。

图1. oneDNN Graph Compiler的IR和优化流程
图1. oneDNN Graph Compiler的IR和优化流程

整体架构。图1展示了oneDNN Graph Compiler设计的高层视图。输入的DNN计算图在内部表示为图IR。图IR优化模块执行一系列转换，将计算图优化并分组为一系列融合操作（Fused operations）。图IR会进一步降级为张量IR。张量IR不保留DNN操作的语义，其语义接近于C程序。它操作的数据结构是多维数组，代表物理内存中的张量缓冲区。张量IR随后会进一步降级为LLVM IR和对微内核的内在函数调用。

模板和微内核简化设计。使用模板和微内核极大地简化了编译器的设计。就像使用C语言实现高性能算子一样，模板实现了已知的最佳算法，并且可以通过参数实例化，然后降级为基于张量IR的高性能算子。由于oneDNN Graph Compiler也使用模板来支持与邻近操作的融合，因此它不需要任何复杂的循环变换和支持性的循环IR，例如MLIR的Affine和Linalg方言。此外，通过使用微内核，oneDNN Graph Compiler避免了开发用于在指令级别生成高效代码的低级优化技术。

图IR（Graph IR）。图IR将优化范围从单个算子扩展到包含多个计算密集型操作的更大子图。由于图IR保留了DNN OP的语义，大多数领域特定的优化都在这一级别完成。DNN OP的语义由模板实现，这些模板直接指导并行任务分解、循环调度和分块、张量内存布局以及如何与邻近操作融合的决策。图IR使用图（graph）、逻辑张量（logical tensor）和操作（OP）来描述计算图。一个图包含一组OP和逻辑张量。每个OP代表计算图中的一个操作。逻辑张量代表张量的元数据，如元素的数据类型、形状和内存布局。OP具有类型（kind）、类别（category）、属性以及输入和输出的逻辑张量。

图IR优化模块。图IR优化模块首先将复杂的OP分解为基本的DNN OP。复杂的DNN OP是指那些具有复杂语义、可以由简单的基本操作（如加法和乘法）组合而成的OP。它们由DL框架引入，以支持高级DNN OP语义，方便编程，例如batchnorm、quantize、gelu和许多激活操作。基本的DNN OP被分为可调优OP（Tunable OP）或可融合OP（Fusible OP）。可调优OP描述了那些使用可调参数来实例化预定义模板以生成最佳性能代码的DNN操作，例子包括像矩阵乘法（matmul）这样的计算密集型操作。可融合OP指的是可以融合到可调优OP中的操作，例如逐元素操作、广播、归约和数据移动操作。

优化流程。复杂DNN操作的分解简化了图IR优化模块，使其只需处理基本的DNN操作。除了通用的编译器优化，如公共子表达式消除（CSE）、死代码消除和常量折叠外，它还包括领域特定的优化，如低精度转换、张量内存布局传播、常量权重预处理和融合。融合优化过程会判断融合两个操作是否有利可图，并持续融合OP以形成一个子图，该子图表示为一个融合OP（Fused OP）。最终，图IR被转换为一个由Fused OP组成的图，然后降级为张量IR。

张量IR（Tensor IR）。张量IR支持将Fused OP机械地降级为简单的循环，而无需进行复杂的嵌套循环分析和转换。张量IR的优化主要关注张量缓冲区的优化，并支持低级代码生成。与C程序类似，张量IR支持函数、语句、表达式和内在函数。从图IR图降级而来的张量IR模块包含多个函数，每个函数代表一个降级后的Fused OP。张量IR模块有一个入口函数，其中包含对从Fused OP降级而来的其他函数的一系列调用。一个张量IR函数包含多个基于表达式构建的语句，这些表达式操作于常量、变量和张量。常量和变量代表单个数据元素，用于表示标量数据，如循环索引、张量形状、地址和张量缓冲区的偏移量。张量代表由数据缓冲区支持的多维数组。内在函数用于表示微内核，这些微内核经过精心手动调优，在单个CPU核心上使用最快的缓存来完成DNN OP的子任务。

III. 用于可调优OP降级的基于微内核的模板

模板化代码生成方法。自动化地为可调优OP（Tunable OP）生成高性能代码是张量编译器的基础。oneDNN Graph Compiler采用了一种继承自高性能库开发的方法，即首先为给定的可调优OP创建代码模板，然后用启发式方法决定的参数来实例化它。这些参数是根据输入数据张量的形状和微架构的硬件尺寸来决定的。

图2. 用于可调优OP的基于微内核的模板
图2. 用于可调优OP的基于微内核的模板

Matmul模板示例。图2中展示的模板用于一个执行矩阵乘法 $A[M, K]$ 和 $B[K, N]$ 产生 $C[M, N]$ 的matmul操作。该模板适用于一个常见的深度学习用例，其中计算使用多个核心，并且输入和输出张量的大小适合缓存系统。外部的并行循环将内核划分为多个子任务，供多核处理。每个子任务被分配给一个单独的核心，称为单核内核（single-core kernel），由内部循环表示，其在最内层循环体中调用一个微内核。

张量切片与子矩阵。微内核和单核内核操作的是张量切片（tensor slice），它代表张量元素的一个子集。例如，原始张量表示为 $A[0:M, 0:N]$，其中下标表示每个维度的起始偏移和大小。张量切片表示为 $A[0:MB, 0:NB]$，其中MB和NB指的是张量切片沿m和n维度的分块大小。子矩阵是二维张量切片的一个特例。在上面的模板中，微内核产生一个小的子矩阵 $C[0:MB, 0:NB]$，而单核内核输出一个更大的子矩阵 $C[0:MSN, 0:NSN]$。

微内核的角色。微内核是oneDNN Graph Compiler实现与专家调优算子相当性能的重要元素。oneDNN Graph Compiler使用了名为batch-reduce GEMM的微内核【8，LIBXSMM: Accelerating small matrix multiplications by runtime code generation, 2016, SC ’16】【24，Tensor processing primitives: A programming abstraction for efficiency and portability in deep learning workloads, 2021, CoRR】。该微内核有两个输入，都代表一批二维矩阵。它首先对每个批次元素进行矩阵乘法，生成一批中间二维矩阵，然后将它们求和得到最终的二维矩阵输出。这个接口可以用于推理和训练用例中matmul和convolution操作的许多变体，并已被oneDNN算子库和oneDNN Graph Compiler采用。

微内核调优与数据布局。微内核经过精细调优，通过充分利用计算功能单元以及寄存器和L1缓存提供的高带宽来最大化计算效率。它抽象了ISA（指令集架构）的差异，因此oneDNN Graph Compiler无需处理不同CPU提供的不同向量或矩阵指令。然而，oneDNN Graph Compiler需要为微内核选择合适的输入子矩阵大小，使其通常是向量和矩阵功能单元所用寄存器大小的倍数。同时，它还需要选择微内核的批处理大小（batch size），以确保整个输入和输出子矩阵都能容纳在L1缓存中。为了进一步简化缓存访问，输入和输出张量是分块的。为简化实现，输入和输出张量使用子矩阵大小 [MB, NB, KB] 进行分块。因此，每个微内核访问的是一个连续的内存缓冲区。

模板参数的确定。降级一个matmul操作的参数指的是上述模板中的变量值：MPN、NPN、MB、NB、KB、BS，以及由msi、ksi和nsi索引的循环的顺序。其他参数可以从这些参数推导出来。oneDNN Graph Compiler使用一种专家调优的启发式方法来决定这些参数。对于给定的输出矩阵大小，它首先提出单核内核大小的选项，即一组[MPN, NPN]，这些选项可以在所有核心上实现良好的负载均衡。然后，它进一步提出微内核大小的选项，即一组[MB, NB, KB, BS]，这些选项能确保良好的微内核性能。接着，启发式方法会选择一对这样的尺寸组合，使得在整个系统所有核心上获得最佳的整体内核性能。它基于一个考虑了多核负载均衡和单核内核效率的成本模型，迭代搜索最佳参数。启发式方法还会报告在搜索过程中计算单核内核效率时所假设的内部循环顺序。

模板的多样性。oneDNN Graph Compiler为不同的用途开发了多种模板。一个可调优OP可以有多个模板，具体取决于使用场景。例如，在推理场景中，有时用例只处理一个数据样本但使用多个核心，因此模板可能需要应用“k-slicing”来从归约轴中提取额外的并行性。

伪代码符号说明。张量用张量名称后跟每个维度的索引和大小来描述。张量 $A[0:M, 0:K]$ 指的是从位置[0,0]开始、大小为[M, K]的二维张量。$A[0:MB, 0:KB]$ 指的是一个张量切片，包含张量A的一个子集，沿m维度从位置0到MB-1，沿n维度从0到NB-1。伪代码对A、B和C使用了分块布局。$C[0:MPSN, 0:NPSN, 0:MB, 0:NB]$ 表示完整的C张量 $C[0:M, 0:N]$ 经过分块布局重排后的形式。$C[mps:1, np:NSN, 0:MB, 0:NB]$ 表示一个张量切片，它在C张量的前两个维度上从位置“mps”和“np”开始，大小为“1”和“NSN”。A_addr[0..BS] 表示一个包含BS个元素的数组，从A_addr[0]到A_addr[BS-1]。$A[mps:1, ks:0..BS, 0:MB, 0:KB]$ 表示一个包含BS个张量切片的数组，从 $A[mps:1, ks:0, 0:MB, 0:KB]$ 到 $A[mps:1, ks:BS-1, 0:MB, 0:KB]$。

IV. 用于融合OP降级的带锚点的模板

带锚点的融合模板。oneDNN Graph Compiler将一个可调优操作（Tunable op）与多个相邻的可融合操作（Fusible ops）组合成一个融合操作（Fused op），并使用融合OP模板将其降级为嵌套循环。该模板在每个循环级别的开始和结束处为输入和输出张量包含了占位符，称为锚点（anchors）。图IR融合优化过程决定将一个可融合操作融合到可调优操作中是否有利可图，并为该可融合操作分配一个锚点。融合OP降级过程会检索可融合操作的锚点，并直接在锚点位置插入其对应的张量IR。

图3. 带有锚点和成本表的融合OP模板
图3. 带有锚点和成本表的融合OP模板

锚点及其成本分析。图3展示了模板内的锚点以及与每个锚点关联的张量切片。微内核之前的锚点被称为前置操作锚点（pre-op anchors），而微内核之后的锚点则被称为后置操作锚点（post-op anchors）。图3右侧的表格显示了每个锚点处张量切片的工作集大小，这描述了在单个核心上，融合操作在锚点处访问的内存大小。它还显示了计算融合操作在单个核心内核内被调用多少次，以及每个锚点总共需要多少张量元素内存访问的公式。当模板用可调优操作的参数实例化后，可以推导出具体的数值。

锚点选择启发式。融合优化使用启发式方法来决定选择哪个锚点。该启发式方法评估所有可能的锚点与不融合选项之间的单核内核成本，然后选择估计成本最低的那个。最内层循环内的提交锚点（commit anchors）处理最小的张量切片，由于数据位于最快的缓存中，因此提供了较低的单次访问成本。所以后置操作（post-op）通常会选择朝向最内层循环的第一个锚点作为最佳选择。然而，对于前置操作（pre-op），融合OP降级过程需要同时考虑前置操作融合引入的计算量和临时缓冲区大小。内层循环体中的锚点需要较小的临时缓冲区，但可能会产生冗余计算，这可以通过精心选择锚点来避免。

图4. 融合OP的伪代码
图4. 融合OP的伪代码

融合示例。图4展示了一个将数据布局重排序（reorder）和ReLU（修正线性单元）操作融合到一个实例化的GEMM操作的伪代码。第一个重排序操作作为前置操作融合插入到锚点#4，它将一个普通布局的张量A转换为一个分块布局的A'，分块因子为MB和KB。融合的重排序操作作用于A'的张量切片，表示为 $A'[mpsi:1, ksi:BS, 0:MB, 0:KB]$，它从位置 $A'[mpsi, ksi, 0, 0]$ 开始，切片大小为[BS, MB, KB]。它还融合了两个后置操作，一个ReLU操作后跟一个重排序操作。这两个操作都插入在后置操作锚点#1处。该重排序操作将C张量的内存布局从分块因子MB和NB更改为MB2和NB2。

V. 图IR优化

图IR优化流程。oneDNN Graph Compiler的图IR优化将图转换为使用低精度计算和预处理过的常量张量，然后将图准备成一系列融合操作（Fused ops），以便进行优化的代码生成。图IR首先被分解为基本DNN操作的图以简化优化过程，然后聚类形成细粒度的融合操作，最后使用融合OP模板降级为张量IR。

图5. 图IR优化
图5. 图IR优化

低精度转换。图5以一个量化的多层感知机（MLP）为例说明了这些优化过程。图5.1中的输入DNN量化MLP包含两个matmul操作（为简化说明省略了激活操作）。每个FP32 matmul操作都被两个反量化（dequantize）操作和一个量化（quantize）操作包围，分别用DQ和Q表示。反量化将Int8数据类型的张量转换为FP32，而量化操作则相反。它使用非对称量化方案，因此第一个反量化操作用a_s缩放A输入张量，然后用a_z进行偏移以调整零点，另一个反量化操作用b_s缩放权重矩阵B。图5.2中优化后的DNN图显示了低精度转换优化的效果。它首先将量化和反量化操作分解为简单的加法和乘法操作，然后将图转换为使用Int8 matmul操作的数学等价形式。低精度转换带来了显著的加速，因为它减少了计算一个深度学习模型所需的计算量和内存带宽。

常量权重预处理。图5.2还展示了常量权重预处理优化的效果。该优化识别出常量权重张量及其相关计算，并构建一个特殊的初始化函数，该函数预处理常量权重并在运行时重用预处理后的权重。对于静态量化推理用例，权重张量和量化参数是常量，因此在运行时可以完全避免对常量权重、缩放因子和零点的计算。由于在编译期间权重数据缓冲区可能不可用，因此编译后的代码需要生成一个函数，在执行时首次接收到常量权重时对其进行预处理。量化参数a_s、b_s、c_s和c_z作为反量化操作的属性传递，是常量，可以在降级到张量IR时被折叠到生成的代码中。

细粒度融合。细粒度融合优化将图聚类为细粒度的图区域，并将它们封装为融合操作。它首先考虑可调优操作的直接后续操作作为后置操作（post-op）的候选者，并不断扩大融合OP区域。后置操作可以是逐元素、广播、归约和重排序操作，并且多个后置操作可以被添加到一个融合OP区域中。例如，matmul操作后的激活和归一化操作被分解为基本操作并添加到融合OP区域。当达到某个限制时，区域停止增长，例如，后置操作序列只能有一个重排序、一个归约、一定数量的总操作数或总附加输入大小。然后，它寻找前面的操作作为前置操作（pre-op）的候选者。前置操作融合仅支持有限的情况，如重排序和转置操作，并且仅在图的入口点使用。对于位于两个可调优操作之间的可融合操作，将其作为第一个可调优操作的后置操作进行融合通常更有利，因此融合优化首先添加后置操作，然后添加前置操作到融合OP区域。

布局传播。布局传播优化通过允许可调优操作使用最理想的分块布局，来发掘跨可调优操作的额外性能优势。由于可调优操作依赖于分块布局以在CPU上实现最佳性能，两个可调优操作之间最佳性能的分块布局通常可能不同。该优化允许图中的可调优操作使用分块布局，但保持图输入/输出张量为普通布局。它首先在图的边界插入重排序操作，以确保入口和出口点使用普通布局。然后，它遍历DNN计算图，并在两个使用不同分块布局的可调优操作之间插入重排序操作。它首先查询一个可调优操作其期望的分块布局，如果所有期望的布局都与当前布局不一致，它就在该可调优操作之前插入一个重排序操作。图5.3展示了细粒度融合区域以及在图边界和可调优操作之间内部新插入的重排序操作。两个可调优操作之间的重排序操作被添加到前一个可调优操作的融合OP区域的末尾。用于输入常量权重的重排序被转换为预处理权重，称为预打包权重（prepacked weight），这是通过常量权重预处理优化实现的。

粗粒度融合。粗粒度优化发生在优化的最后阶段，此时图被转换为按拓扑顺序排列的融合操作列表，并降级为一系列parallel-for。它将由融合操作降级而来的多个parallel-for循环合并成一个parallel-for循环。合并多个parallel-for的机制在张量IR级别得到支持，但合并的决定是在图IR级别作为降级的一部分完成的。粗粒度优化极大地改善了两个融合操作之间的数据局部性，并可应用于许多图模式。例如，由于MLP中的两个matmul操作使用相同的批次（batch），降级后的两个嵌套并行循环具有相同的最外层循环，该循环遍历批次维度，这可以被进一步合并为一个并行循环。当启发式方法为一个融合操作选择参数时，它会尝试选择与核心数量最匹配的最外层循环分块因子，因此实例化的融合操作很可能与其邻居具有相同的分块因子。当粗粒度融合优化决定合并两个融合操作时，它会在降级过程中将张量IR中的两个嵌套循环标记为“可合并”。然后，张量IR在图IR优化的指导下机械地合并这两个嵌套循环。

VI. 张量IR优化

张量IR概述。张量IR是oneDNN Graph Compiler中最低级别的中间表示。在张量IR级别，DNN计算图被降级为一个类似C的程序，其中包括函数、语句、表达式和内在函数。融合操作（Fused op）被降级为一个函数，其中包含嵌套循环。复杂语句描述了像循环这样的结构，而简单语句则执行计算。Var和Tensor分别代表标量变量和多维数组。

支持图IR优化。张量IR通过按照图IR的指示合并循环来支持图IR的优化。图6展示了图4中伪代码对应的张量IR示例。在张量IR中，对张量切片的计算由嵌套循环或对微内核的函数调用来表示。插入的前置操作和后置操作被降级为嵌套循环。两个后置操作，即ReLU和重排序操作，根据图IR传递的提示被合并为一个嵌套循环。

张量大小优化。张量IR上的主要优化是张量大小优化和内存缓冲区优化。张量大小优化试图减少每个临时张量的尺寸。临时张量是在前置和后置操作融合过程中引入的。在降级过程中，临时张量最初被引入为全尺寸张量，然后通过张量大小优化来减小其尺寸。在图6的示例代码中，后置操作被融合成一个循环嵌套在张量IR中。由于对临时张量C''和C'''的访问是局限于最内层循环体的，因此临时张量可以被一个标量变量替换，以获得更小的内存占用和更好的缓存局部性。由前置操作融合引入的临时张量也可以通过分析张量使用范围来类似地减小。例如，$A'[MSN, BS, MB, KB]$ 可以被减小为 $A'[BS, NB, KB]$，因为A'的生产者和消费者都在“msi”循环内部，所以没有必要保存A'第二个维度的结果。

内存缓冲区优化。在张量大小优化之后，多维张量表示被扁平化为一维数组以代表内存缓冲区。内存缓冲区优化试图重用临时张量的内存缓冲区，以为编译后的代码实现最小的整体临时缓冲区大小，并试图改善临时缓冲区使用的局部性。内存缓冲区优化的主要目标是重用在融合操作之间为临时张量创建的内存缓冲区。在推理用例中，输出张量仅被下一个融合操作消费，因此一旦下一个融合操作执行完成，该缓冲区就可以被回收。由于输入张量的大小在编译过程中是已知的，因此可以在编译时计算和优化内部内存缓冲区的使用，以提高效率。内存缓冲区优化使用类似于传统编译器基于定义-使用链（def-use chain）进行寄存器分配的生命周期分析。该算法既考虑重用热点内存，也考虑减少整体峰值内存。在每个需要中间缓冲区的点，它会尝试重用已经分配但不再使用的空闲中间缓冲区。在多个可重用内存缓冲区的选择中，它会选择最近使用的那个，这样数据很可能仍存在于缓存系统中。

Var Const MPN, NPN, MSN, NSN, BS, KSN, MB, NB;
Tensor FP32[M, K] A;
Tensor FP32[M/MB, K/KB, MB, KB] A';
Tensor FP32[K/KB, N/NB, NB, KB] B;
Tensor FP32[NSN, MB, NB] C';
Tensor FP32[M/MB,N/NB, MB, NB] C'', C''';
Tensor FP32[M/MB2, N/NB2, MB2, NB2] C;
Var int* A_addr[BS], B_addr[BS];
Var Index mp, np, ms, ks, ns, nps, mps;
Parallel loop mpi = 0, MPN, 1 {
    Parallel loop npi = 0, NPN, 1 {
        Loop msi = 0, MSN, 1 {
            mpsi = mpi * M/MPN + msi;
            Loop nsi = 0, NSN, 1 {
                Loop mbi = 0, MB, 1 {
                    Loop nbi = 0, NB, 1 {
                        C’[0:NSN, 0:MB, 0:NB] = 0;
                    }
                }
            }
            Loop ksi = 0, KSN, BS {
                Loop bsi = 0, BS, 1 {
                    ksbi = ksi * BS + bsi;
                    Loop mbi = 0, NB, 1 {
                        Loop kbi = 0, KB, 1 {
                            A’[mpsi, ksbi, mbi, kbi] = A[mpsi*MB + mbi, ksbi*KB+kbi];
                        }
                    }
                }
                Loop nsi = 0, NSN, 1 {
                    npsi = npi * N/NPN + nsi;
                    Loop bsi = 0, BS, 1 {
                        A’_addr[bsi] = &A’[mpsi, ksi, 0, 0];
                        B_addr[bsi] = &B[ksi, npsi, 0, 0];
                    }
                    C'_addr = &C’[nsi,0, 0] ;
                    Batch_reduce_gemm(A’_addr, B_addr, C'_addr, MB, NB, KB, Batch = BS);
                }
            } //end of loop ksi
            Loop nsi = 0, NSN, 1 {
                npsi = npi * N/NPN + nsi;
                Loop mbi = 0, MB, 1 {
                    Loop nbi = 0, NB, 1 {
                        C’’[mpsi, npsi, mbi, nbi] = C’[nsi, mbi, nbi];
                        C’’’[mpsi, npsi, mbi, nbi]= max(C’’[mpsi, npsi, mbi, nbi], 0);
                        C[(mpsi*MB+mbi)/MB2, (npsi*NB+nbi)/NB2, (mpsi*MB+mbi)%MB2, (npsi* NB+nbi)%NB2] = C’’’[mpsi, npsi, mbi, nbi];
                    }
                }
            } // end of loop nsi
        } // end of loop msi
    }
} // end of parallel loop mpi, npi

图6. 张量IR示例

A4 实验环境与结果

实验环境

oneDNN Graph Compiler作为oneDNN库的嵌入式组件，用于加速DNN计算子图。oneDNN是业界标准的、性能领先的库，已被多个DL框架集成为默认的性能库，以在CPU上加速深度学习。

实验结果

单个Matmul算子性能比较 (Fig. 7)

为了模拟真实工作负载中的缓存局部性效应，我们在MLP子图的上下文中连续运行matmul操作，而不是孤立地测试。

图7. oneDNN primitives、TVM和oneDNN Graph Compiler的Matmul内核执行时间比较
图7. oneDNN primitives、TVM和oneDNN Graph Compiler的Matmul内核执行时间比较
子图性能比较 (Fig. 8)
图8. oneDNN Graph Compiler在MLP和MHA子图上的性能评估
图8. oneDNN Graph Compiler在MLP和MHA子图上的性能评估
端到端模型性能 (Fig. 9)

由于TVM的自动调优时间过长，无法进行端到端模型级别的实验。

图9. oneDNN Graph Compiler带来的端到端模型加速
图9. oneDNN Graph Compiler带来的端到端模型加速

A5 结论

本文提出了一种混合方法来应对深度学习编译的独特挑战。该方法提炼了专家为计算密集型DNN操作（如矩阵乘法）调优算子时的关键要素，并结合编译器技术在DNN计算图层面充分挖掘性能机会。模板使用了专家开发的微内核、算法和启发式方法，以确保编译器生成的代码能够达到与专家调优算子相当的性能。编译器使用图操作（DNN op graph）和类C程序（C program）两个级别的中间表示，以支持深度学习计算所需的领域特定优化，包括低精度计算、常量权重处理、张量内存布局、细粒度融合、粗粒度融合以及张量内存缓冲区的重用。性能评估结果显示，在CPU推理场景中，对于性能关键的DNN计算图和端到端模型，该方法相比现有的张量编译器和算子库取得了显著的性能增益。