江苏龙网

  1. 首页 > 资讯 > 科技 > >

|CPU推理性能提高数十倍，旷视天元计算图、MatMul优化深度解读( 三 )


将下图中红色 (访问重复次数最多的 A 的行块 ， 计算时需要的 B 的一个列块以及计算结果的 C 的小块) 部分都保存在 L1 中 。
由于计算完每一个 C 的行块 ， 都需要重复遍历一次整个 B 矩阵 ， 因此将 B 存放在 L2 中使得每次读取 B 的一个列块都是从 L2 中读取 ， 代价最小 。
将重复访问率最高的 C 的累加中间结果保存在速度最快的 CPU 处理器的寄存器中 。
|CPU推理性能提高数十倍，旷视天元计算图、MatMul优化深度解读
本文插图

通过上面的分配策略 ， 并结合 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(mr*Kr)*B(Kr*nr)=C(mr*nr)的分块矩阵乘 ， 决定了整个 MatMul 的计算性能 ， 因此需要极致地优化 。
Kernel 计算
最核心的计算 Kernel 进行的是尺寸为 (mr*Kr)x(Kr*nr)-->(mr*nr) 的小尺寸 MatMul ， 其计算示意图如下：
|CPU推理性能提高数十倍，旷视天元计算图、MatMul优化深度解读
本文插图

如上图所示 ， Kernel 在计算时会读取 A 中一列 ，B 中一行 ， 进行矩阵乘 ， 得到 大小为 mr*nr 的 C ， 然后和原来 C 中的值相加 ， 如此循环 Kr 次 ， 完成该 Kernel 的计算 。
在该 Kernel 的计算过程中乘法和加法的计算量为 2mr*nr*Kr ， 访存量为 (mr+nr)*Kr+2mr*nr ， 可以根据处理器来判断该计算能否隐藏数据的访存 。 下面以 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 不友好） ， 进而获得巨大收益 。


推荐阅读


上一篇：青年|150斤女生的5套冰淇淋色系穿搭，平价百搭又显瘦

下一篇：|做数据分析还在死磕Excel？用这个简单工具，摆脱复杂函数和公式