MegEngine计算图、MatMul优化解析

MegEngine计算图、MatMul优化解析

本文针对天元在推理优化过程中所涉及的计算图优化与 MatMul 优化进行深度解读，希望能够帮助广大开发者在利用天元 MegEngine「深度学习，简单开发」的同时，也能够了解 CPU 优化的相关知识。从而帮助大家在模型部署的整体流程中更好地进行加速；在实际模型部署时能够评估模型在特定平台上运行所能达到的性能以及内存使用情况；以及在算法设计时可以设计出更利于 CPU 优化加速的卷积 Opr 等。

本文针对旷视天元深度学习框架在推理优化过程中所涉及的计算图优化与 MatMul 优化进行深度解读。

背景及引言

在深度学习大规模落地边缘端场景的今天，如何最大程度降本增效，是企业与开发者共同关注的话题。其中，模型的训练与推理是两个关键环节。

天元（MegEngine）深度学习框架凭借「训练与推理一体化」的独特范式，能够极大程度上（90%）节省模型从研发到部署的整体成本，降低转换难度，真正实现小时级转化；同时，天元（MegEngine）在 CPU 推理方面所做的大量优化工作，也使得开发者在推理时能够发挥出处理器的最佳性能。

在之前对天元的极致推理优化进行了综述《 工程之道，MegEngine 推理性能极致优化之综述篇》。本文则针对天元在推理优化过程中所涉及的计算图优化与 MatMul 优化进行深度解读，希望能够帮助广大开发者在利用天元 MegEngine「深度学习，简单开发」的同时，也能够了解 CPU 优化的相关知识。

从而帮助大家在模型部署的整体流程中更好地进行加速；在实际模型部署时能够评估模型在特定平台上运行所能达到的性能以及内存使用情况；以及在算法设计时可以设计出更利于 CPU 优化加速的卷积 Opr 等。

CPU 推理优化概览

对于产业应用而言，CPU 推理的性能优化至关重要，如下表所示，经过优化的推理性能可以较未经优化的原始性能提升数十倍。

 

 在进行模型推理阶段的优化时，首先需要对模型的计算图进行优化，以避免冗余的计算与访存，确保计算图在推理时是最优的；其次，在大多 CV 相关模型中，卷积的计算比重最高，达到模型总计算量 90% 以上，因此对模型推理的优化主要聚焦在对卷积计算的优化上。

通常而言，卷积计算的实现方式有 3 种：direct 卷积，Im2col 卷积以及 Winograd 卷积。

• direct 卷积：根据卷积的计算公式直接对 FeatureMap 上进行滑窗计算；
• Im2col 卷积：根据卷积计算需要在输入通道上进行 reduce sum 的特点，将卷积运行转化为 MatMul 计算；
• winograd：在保证计算无误的前提下，使用加法替代乘法，达到优化卷积乘法计算量的目的，在中间过程需要使用 MatMul 进行计算。

由上文可知，以 Im2col 或 Winograd 方法进行的卷积计算会频繁使用到 MatMul，因此对 MatMul 这种基础算子的优化就显得尤为重要。

基于上述考量，本文将首先介绍模型优化中的图优化，然后介绍基础算子 MatMul 在 CPU 上的优化方法。

推理计算图优化

在训练阶段定义模型的计算图，主要是为了满足模型参数的训练需求。当训练结束，模型参数固定后，对计算图的进一步优化能够帮助模型推理更加高效。

推理和训练的计算图是一张有向无环图 (DAG)，在天元中，开发者能够以类似 LLVM 的方式对 DAG 计算图定义许多优化方法，这里简称 OptPass。OptPass 可以根据用户的配置有选择性的加入到图优化的 OptPass 列表中，从而帮助用户灵活地为图优化定义 OptPass。

计算图优化

天元定义了多个为推理进行计算图优化的 OptPass，开发者使用这些 OptPass 之后，将得到一张用于推理的最优计算图。下面以 MobileNetV1 中， Convolution+Batch Norm+Relu 这样的典型结构经过计算图优化之后 Fuse 为 ConvBias 的过程为例，介绍天元的计算图优化过程。

由于 Batch Norm 在推理阶段除了输入 Tensor 外其他都是常数，因此可以简化为多个 elemwise 的组合，天元实现了一个 ConvertBatchNormToElemwisePass 的 Optpass，这个 OptPass 将模型中的所有 Batch Norm 转化为 Elemwise，具体转化原理如下：

 

 紧接这一过程，天元将运行 OptPass ParamRedistributePass，该 OptPass 会将上述 Batch Norm 转化而来的 Elemwise 中的 scale 融合到 convolution 的权重中，具体实现原理如下：

 

 最后天元将运行 OptPass FuseConvBiasNonlinPass，该 OptPass 会将计算图中的 Convolution+Elemwise 转化为天元内部实现的 ConvBias Op 中，同时，它还会设置 ConvBias Op 中 NonelineaMode 参数。

如此便完成了从 Convolution+Batch Norm+Relu 到 ConvBias 的转换，整体转换过程如下图：

 

 实验验证图优化之后的性能

在推理之前，如果要对完成训练的模型进行图优化，则需要在模型 dump 的期间进行，当然，也可以在 SDK 运行模型之前进行。下面是在模型 dump 时进行图优化的代码：

from megengine.jit import trace

@trace(symbolic=True)

def pred_fun(data, *, net):

    net.eval()

    pred = net(data)

    pred_normalized = F.softmax(pred)

    return pred_normalized

 

# 使用 trace 类的 trace 接口无需运行直接编译

pred_fun.trace(data, net=xor_net)

# 使用 trace 类的 dump 接口进行部署

pred_fun.dump("xornet_deploy.mge", arg_names=["data"], optimize_for_inference=True, enable_fuse_conv_bias_nonlinearity=True)

上面的 optimize_for_inference=True 将在 dump 模型时候针对 inference 进行优化， enable_fuse_conv_bias_nonlinearity=True 将在模型中进行 Op Fuse 的优化，此外天元还支持其他的优化参数，具体可见天元的文档。

下表展示的，是在实验过程中对模型进行图优化前、后，模型运行性能的测试对比。

 

 可以看出图优化对模型性能的提升效果显著，具体提升比例由模型自身决定。

MatMul 优化

如前文所述，MatMul 作为卷积运算的基础算子，会频繁地被 Im2col、Winograd 以及 FullyConnect 使用。因此在天元中，MatMul 既被封装为单独的 Op，也可以作为单独的 algo 供卷积的实现使用。

此外，由于 MatMul 也是计算最为密集的算子，因此天元对它进行了极致的优化。

优化

MatMul 是线性代数中的矩阵乘，假设矩阵 A 大小为 MK，矩阵 B 大小为 KN，则得到矩阵 C 大小为 M*N，其中 C 的每个元素的计算公式如下：

 

 可以发现，在 MatMul 的计算中乘法和加法的计算量为 2MNK (计算 C 中每个元素时，加法和乘法计算量分别为 K，C 的总元素个数为 MN)，访存量为 2MNK (计算每个 C 中元素需要 2K 访存)+ 2MN(整个 C 矩阵读一次和写一次)。由于计算量固定 (排除 Strassen)，所以只能优化访存，使得乘法和加法运算达到处理器的极限性能，从而实现 MatMul 的最佳性能。

MatMul 分块

关于减少 MatMul 计算时的访存量，最有效的方法是对 MatMul 的计算进行分块，并将分块之后的数据保存在 CPU 的 Cache 中。

如下图所示，将 A 按照 mr 进行行分块，将 B 按照 nr 进行列分块，计算时将需要用到的分块保存在 CPU 的各级 Cache 中，从而极大的减少了内存的读写。

 

 进一步细化上面的分块计算过程如下图所示，A 中的一个行块都要重复地和 B 中的每一个列块进行小分块的 MatMul，写入到 C 的小块中，为了使得分块尽量的大 (如上所述，能够减少内存访存量)，Cache miss 率尽量低，因此需要根据 CPU 的 Cache 结构特点 (速度：L1D>L2>L3，容量 L1D<L2<L3) 来分配 Cache 和数据块之间的存放关系，天元中进行的分配如下：

• 将下图中红色 (访问重复次数最多的 A 的行块，计算时需要的 B 的一个列块以及计算结果的 C 的小块) 部分都保存在 L1 中。
• 由于计算完每一个 C 的行块，都需要重复遍历一次整个 B 矩阵，因此将 B 存放在 L2 中使得每次读取 B 的一个列块都是从 L2 中读取，代价最小。
• 将重复访问率最高的 C 的累加中间结果保存在速度最快的 CPU 处理器的寄存器中。

 

 通过上面的分配策略，并结合 CPU 中资源 (寄存器数量，L1D 和 L2 的大小)，便可以确定最佳的 MatMul 计算中的 Nr，Kr：

• 可以根据 CPU 处理器的寄存器数量得到 mr 和 nr 的具体大小，寄存器容量 > mr*nr
• 根据 L1D Cache 的大小结合 mr 和 nr 计算出 Kr，Kr=L1D/(mr+nr)
• 再根据 L2 的大小计算出 B 矩阵中的 Nr，Nr=(L2-L1D)/Kr

在上面计算 N 时，用 L2-L1D 的原因是，由于当前 CPU 使用的 Cache 是 Inclusive 的，且 L2 是指令缓存和数据缓存合并的，另外上面 M 没有明确限制。

在得到上面最佳 Nr、Kr、mr 和 nr 之后，进一步便可以首先对 MatMul 计算中的 N、K 进行 Nr 和 Kr 分块，然后在 Nr、Kr 的基础上再进行 mr 和 nr 分块。如此，MatMul 计算中的计算访存比达到最高，且 CPU 处理器的资源也得到充分利用。

经过分块之后，由于在最内层的计算 Kernel 为 A(mrKr)B(Kr*nr)=C(mr*nr) 的分块矩阵乘，决定了整个 MatMul 的计算性能，因此需要极致地优化。

Kernel 计算

最核心的计算 Kernel 进行的是尺寸为 (mrKr)x(Krnr)–>(mr*nr) 的小尺寸 MatMul，其计算示意图如下：

 

 如上图所示，Kernel 在计算时会读取 A 中一列， B 中一行，进行矩阵乘，得到 大小为 mr*nr 的 C，然后和原来 C 中的值相加，如此循环 Kr 次，完成该 Kernel 的计算。

在该 Kernel 的计算过程中乘法和加法的计算量为 2mrnrKr，访存量为 (mr+nr)Kr+2mrnr，可以根据处理器来判断该计算能否隐藏数据的访存。下面以 ARM cortex A76 为例进行分析，根据 A76 的数据手册得到：

• FP32 SIMD Load throughput=2，即单周期可以 load 8 个 float 数据
• FP32 SIMD FMLA throughput=2，即单周期可以进行 16 个乘加运算

因此，当 Kernel 的尺寸 mr=8、nr=12、Kr=256，计算量为 49152 次乘加运算，访存量为 5312 个 float 数据时，该计算访存量为 9.25，大于处理器的计算访存比 2。因此可以得出结论，如果 A 和 B 均在 L1 中，则该 Kernel 的计算不会因为数据的 Load 被阻塞，所以计算单元能够发挥出处理器的最佳性能。

虽然在对 MatMul 进行分块时，已经计划将 A 和 B 这样的小矩阵置于 L1D Cache 中，但是数据在真正运行时却不一定都在 L1D 中，原因在于，B 矩阵的列块在原来的大矩阵中内存并不连续，其次 A 中的一列由于内存不连续，也不能使用 SIMD 进行 load，为此需要对 A 和 B 进行数据 PACK。

MatMul 数据 PACK

上文 kernel 在计算过程中，需要同时读取 A 矩阵 mr 行数据，而每行数据之间内存不连续，因此这种访问对 Cache 不友好，同样，在读取矩阵 B 中 nr 列的时候也存在数据读取不连续的问题，加之 A 的所有行块和 B 中的所有列块将被读取多次，因此可以通过对 A 和 B 提前进行数据 PACK， 实现在后续计算中频繁多次访问数据时是连续的（只在第一次 PACK 时对 Cache 不友好），进而获得巨大收益。

 

 对矩阵 A 进行数据 PACK 是将 A 中 mr 行数据的相同列拷贝到一起，如上图中将 A PACK 到 A’ 的步骤。重复完所有 A 中的行块便完成了 A 矩阵的数据 PACK。B 矩阵的 PACK 操作是，将 nr 列数据拷贝到连续的内存地址中，它对应上图 B PACK 到 B’ 的过程。

实验

按照文介绍方式方式，天元在 X86 和 ARM 上分别对 MatMul 进行了优化。下表展示了 ARM64 上的性能测试结果，实验平台为 kirin 980。

首先，对该处理器进行分析可以看到，其主频为 2.6 GHz，每个周期能够进行 16 次乘加计算，因此其理论计算峰值为 16*2.6=41.6 Gflops。

 

 可以看到，经天元优化的 MatMul 计算，发挥出了该处理器 90% 以上的计算性能。

总结

本文以 Batch Norm 为例介绍了推理计算图的具体实现，以及 MatMul 在 CPU 上的优化细节。作为 CPU 推理优化的基石，最优的推理计算图是实现高性能 CPU 推理的前提条件，极致性能的 MatMul 计算基础算子将为实现卷积计算中的 Im2col 和 Winograd 提供性能保障。

在后面的文章中，将在详细介绍卷积计算中 Im2col 和 Winograd 的优化细节。

参考文献

1. Anatomy of High-Performance Matrix Multiplication
2. Fusing batch normalization and convolution in runtime
3. Arm® Cortex®-A76 Software Optimization Guide

了解更多信息可访问：

• MegEngine WebSite： https://megengine.org.cn
• MegEngine GitHub（欢迎 Star）： https://github.com/MegEngine/MegEngine