了解各种 BLAS 实现

计算科学 线性求解器 图书馆 布拉斯
2021-11-30 09:40:50

我不断遇到诸如“高度优化的 BLAS 内核”和“特定于架构的优化”之类的短语,但一直无法找到这些优化到底是什么,以及如何编写自己的特定于架构的优化。

了解不同的 BLAS 实现、它们的低级优化以及架构特定性能调整和基准测试的方法有什么好的方法?

我对这个领域相当陌生,只完成了一门本科课程,要求我们使用 OpenMP、MPI 和 CUDA C 实现简单的 sgemm/dgemm 子例程。如果我能以任何方式改进这个问题,请告诉我。

2个回答

我是许多面向“架构特定优化”的 Julia 库的主要作者,包括LoopVectorization.jlOctavian.jl

对于类似 BLAS 的操作,LoopVectorization.jl 所做的最重要的优化之一是“寄存器平铺”。虽然 CPU 可能有大量的实际寄存器(用于寄存器重命名以启用乱序执行),但它们只有少量的命名寄存器。下面,当我谈到数字时,我特别指的是有多少被命名。对于浮点向量寄存器:

  • Apple M1 (ARM) CPU 具有 32 个 128 位寄存器 (v0-v31)
  • 带有 AVX 的 x86_64 CPU 具有 16 个 256 位寄存器 (ymm0-ymm15)
  • 带有 AVX512 的 x86_65 CPU 具有 32 个 512 位寄存器 (zmm​​0-zmm31)

我们可以做的第一个明显优化是使用可用的最宽向量寄存器。在大多数 CPU 上,指令/周期的吞吐量不会随着寄存器宽度的变化而变化。因此,使用 256 位寄存器而不是 128 位寄存器几乎可以“免费”为您提供 2 倍的性能。W在这里指的是寄存器宽度。对于双精度(64 位)浮点数,W上述 CPU 分别为 2、4 和 8。

此外,并非所有 CPU 都支持“融合乘加”(fma)指令。它们执行乘法和加法。一个简单的优化是在它们可用时使用它们。大多数最新的 CPU 都有它们,所以我假设我们在下面的讨论中也有。

当我在下面写区间时,区间是闭闭区间,就像在 Julia 中一样(与 Python 不同)。所以m:m+W-1意思是m, m+1, ..., m+W-1我还将假设这A是一个以M x K列为主的矩阵,这B是一个以K x N列为主的矩阵(所以C = A*BM x N)。

为了使 CPU 能够对数据进行操作,必须将其加载到寄存器中**。即使内存在最高缓存级别的 L1 缓存中可用,加载可能仍然相对较慢。许多 CPU 可以在每个时钟周期执行 2 次融合乘加指令,而每个时钟周期执行 1-2 次加载(例如,2 次加载不跨越缓存线边界,或者 1 次加载)。因此,如果您想计算A[m:m+W-1,k] * B[k,n] + C[m:m+W-1,n]A[m+W:m+2W-1,k] * B[k,n] + C[m+W:m+2W-1,n],这些 4 W 算术运算(2 个融合乘加)将平均** 1 个时钟周期。然而,这 6 次加载可能需要 3-6 个周期,即使所有内存都在 L1 缓存中!

那么,我们可以在这里做什么呢?答案是寄存器平铺。我们以 AVX2 CPU 为例。我们可以初始化一个 8x6 块C

# 8 elements from column 1
C00 = C[m  :m+ W-1,n]
C10 = C[m+W:m+2W-1,n]
# 8 elements from column 2
C01 = C[m  :m+ W-1,n+1]
C11 = C[m+W:m+2W-1,n+1]
# 8 elements from column 3
C02 = C[m  :m+ W-1,n+2]
C12 = C[m+W:m+2W-1,n+2]
# 8 elements from column 4
C03 = C[m  :m+ W-1,n+3]
C13 = C[m+W:m+2W-1,n+3]
# 8 elements from column 5
C04 = C[m  :m+ W-1,n+4]
C14 = C[m+W:m+2W-1,n+4]
# 8 elements from column 6
C05 = C[m  :m+ W-1,n+5]
C15 = C[m+W:m+2W-1,n+5]

这使用了 16 个命名寄存器中的 12 个,剩下 4 个未使用的寄存器。现在我们遍历内积维度,K在每次迭代中,我们从 的列中加载 8 个元素A,使用另外 2 个寄存器:

A0 = A[m  :m+ W-1,k]
A1 = A[m+W:m+2W-1,k]

然后我们展开n:n+5

B0 = B[k,n] # broadcast load
C00 = A0 * B0 + C00
C10 = A1 * B0 + C10
B0 = B[k,n+1] # broadcast load
C01 = A0 * B0 + C01
C11 = A1 * B0 + C11
B0 = B[k,n+2] # broadcast load
C02 = A0 * B0 + C02
C12 = A1 * B0 + C12
B0 = B[k,n+3] # broadcast load
C03 = A0 * B0 + C03
C13 = A1 * B0 + C13
B0 = B[k,n+4] # broadcast load
C04 = A0 * B0 + C04
C14 = A1 * B0 + C14
B0 = B[k,n+5] # broadcast load
C05 = A0 * B0 + C05
C15 = A1 * B0 + C15

因此,在循环的每次迭代中,k我们现在执行来自 A 的 2 次加载和来自 B 的 6 次加载,总共 8 次加载。通过这 8 次加载,我们可以执行 12 条融合乘加指令。

现在,这足以让 fma 单元保持忙碌(而不是等待 loods),并且实际上达到接近峰值理论 FLOPS!也就是说,如果数据在高缓存级别中可用,就足以让它们保持忙碌。

这是一个特定于架构的优化,因为它依赖于诸如主机 CPU 的向量宽度和它实际可用的寄存器数量等细节。LoopVectorization.jl 有一个成本模型,它使用拉格朗日乘数求解(约束是可用寄存器的数量),在精确整数解的附近找到一个连续解。从那里它可以检查相邻的整数以找到在满足可用寄存器约束的同时最小化成本的值。

理想情况下,宏内核应尽可能大,因为它越大,我们需要传递数据的次数就越少,从而减少了通过缓存所需的内存带宽量。

LoopVectorization 还没有做的事情(但总有一天会做!)是缓存平铺。Octavian.jl、OpenBLAS、BLIS、IntelMKL 等——所有真正的矩阵乘法实现——都这样做。如果您查看 Octavian README 中的基准图表,您会看到 LoopVectorization 的性能如何随着 CPU 开始耗尽缓存而变得非常不稳定,然后由于缺少缓存平铺,最终在 1000x1000 之后像石头一样下降。

缓存平铺遵循与寄存器平铺类似的理念,但它与 CPU 缓存级别 L1、L2 和 L3 中的内存有关。使用寄存器平铺,通过在块中操作,我们能够多次重用从A和加载的内存(例如,我们从 进行了 2 次加载,并将它们与从 B 加载的 6 次相乘)。类似地,我们以块方式操作,因此我们只使用来自并且(在 3 缓存级别 X86_64 CPU 上)适合 L2 和 L3 缓存的块。我们需要将矩阵的整个加载块重新加载到 B 块的每个条带的寄存器中。这可能需要大量的重新加载,但由于内存已经在 L2 缓存中,这应该相对较快。如果我们使用更大的块BAABAn:n+5A,我们会不断地从 L3 缓存或 RAM 中获取数据!高速缓存的大小再次是特定于 CPU 的细节。

** x86_64 CPU 允许您在同一指令中加载和操作,但加载仍然必须发生。

*** 平均而言,我指的是吞吐量。实际完成的时间会是几个时钟周期,一般是 4 或 5 个。但是许多融合乘加可以同时进行,因此大多数 CPU 平均每个时钟周期可以完成 2 个。Apple M1 可以做 4 个/循环。


另一类稍微不同的 CPU 优化是利用可用的特定指令。比如我写了一个exp2专门针对 AVX512 的实现。与 AVX 版本相比,除了使用更大的向量之外,它还可以使用 vscalefpd 和 imperm2pd 等 AVX512 特定指令来进行快速 16 元素查找表,而 AVX 版本使用慢速收集指令和 256 元素表。随机播放指令(以及足够宽的向量以将它们用作表格)是 AVX512 的一项功能,因此“小表格”作为实现策略只是 AVX512 的一个选项。不过,它需要 2 个 fma 单位才能真正提高基于较大表收集的版本的性能。较小的表意味着您需要较大的多项式,这需要大量的 fma。vscalefpd、vgetmantpd、vgetexppd 之类的东西有点像“fma”,因为它们是一种更快的方式,可以在没有 AVX512 的 CPU 上执行多条指令。

阿卜杜拉,感谢您为我们提供材料。我们将这些重新打包为 edX 上的大规模开放在线课程 (MOOC),标题为“LAFF-On Programming for High Performance”。它对审计员是免费的。有关信息,请参阅http://ulaff.net

BLAS 的主要开源实现是类 BLAS 库实例化软件 (BLIS)。在今年四月的 SIAM 新闻中的一篇文章中对此进行了高度讨论: https ://sinews.siam.org/Details-Page/blis-blas-and-so-much-more