如何进一步优化矩阵乘法的性能？

Posted 2023-02-16

【中文标题】如何进一步优化矩阵乘法的性能？

【英文标题】：How to further optimize performance of Matrix Multiplication?

【发布时间】：2019-11-26 12:12:53

【问题描述】：

我正在尝试优化在单核上运行的矩阵乘法代码。如何进一步提高循环展开、FMA/SSE 的性能？我也很想知道为什么在内循环中使用四个而不是两个和时性能不会提高。

问题大小是 1000x1000 矩阵乘法。 gcc 9 和 icc 19.0.5 都可用。 Intel Xeon @ 3.10GHz，32K L1d 高速缓存，Skylake 架构。用gcc -O3 -mavx编译。

void mmult(double* A, double* B, double* C)

    const int block_size = 64 / sizeof(double);

    __m256d sum[2], broadcast;

    for (int i0 = 0; i0 < SIZE_M; i0 += block_size) 
    for (int k0 = 0; k0 < SIZE_N; k0 += block_size) 
    for (int j0 = 0; j0 < SIZE_K; j0 += block_size) 

        int imax = i0 + block_size > SIZE_M ? SIZE_M : i0 + block_size;
        int kmax = k0 + block_size > SIZE_N ? SIZE_N : k0 + block_size;
        int jmax = j0 + block_size > SIZE_K ? SIZE_K : j0 + block_size;

        for (int i1 = i0; i1 < imax; i1++) 

            for (int k1 = k0; k1 < kmax; k1++) 

                broadcast = _mm256_broadcast_sd(A+i1*SIZE_N+k1);
                for (int j1 = j0; j1 < jmax; j1+=8) 
                    sum[0] = _mm256_load_pd(C+i1*SIZE_K+j1+0);
                    sum[0] = _mm256_add_pd(sum[0], _mm256_mul_pd(broadcast, _mm256_load_pd(B+k1*SIZE_K+j1+0)));
                    _mm256_store_pd(C+i1*SIZE_K+j1+0, sum[0]);

                    sum[1] = _mm256_load_pd(C+i1*SIZE_K+j1+4);
                    sum[1] = _mm256_add_pd(sum[1], _mm256_mul_pd(broadcast, _mm256_load_pd(B+k1*SIZE_K+j1+4)));
                    _mm256_store_pd(C+i1*SIZE_K+j1+4, sum[1]);

                    // doesn't improve performance.. why?
                    // sum[2] = _mm256_load_pd(C+i1*SIZE_K+j1+8);
                    // sum[2] = _mm256_add_pd(sum[2], _mm256_mul_pd(broadcast, _mm256_load_pd(B+k1*SIZE_K+j1+8)));
                    // _mm256_store_pd(C+i1*SIZE_K+j1+8, sum[2]);

                    // sum[3] = _mm256_load_pd(C+i1*SIZE_K+j1+12);
                    // sum[3] = _mm256_add_pd(sum[3], _mm256_mul_pd(broadcast, _mm256_load_pd(B+k1*SIZE_K+j1+12)));
                    // _mm256_store_pd(C+i1*SIZE_K+j1+4, sum[3]);
                
            
        
    
    
    

【问题讨论】：

一般来说，只有在展开的循环有未使用的寄存器时，循环展开才能提高性能。如果你走得太远，你展开的循环将开始寄存器溢出/填充并减慢处理速度。请注意，x86 体系结构已经很少注册，因此循环展开通常在 x86 系统上无法正常工作。但要真正了解原因，您必须对代码的两个版本进行详细分析。鉴于您发布的代码有六个级别的嵌套循环，您实际上可能会通过拆分内部循环获得更好的性能。

@AndrewHenle：x86-64 有 16 个 YMM 寄存器；这不会是个问题。甚至在具有 8 个 YMM 寄存器的 32 位代码中也不行。更有可能你看不到任何好处，因为加载+存储是瓶颈，而不是前端带宽或后端 ALU 端口。

什么CPU，什么编译器选项+版本？ 什么问题规模。 我们谈论的是适合 L1d 还是 L2 缓存的矩阵？还是 L3 或主存？

16 并不是一个庞大的寄存器数量，但足以拥有来自 B 的 12 个累加器和 3 个行向量以及来自 A 的广播标量，所以它的结果大约是足够。无论如何，上面 OP 的循环几乎不使用任何寄存器。 32bit的8个寄存器确实太少了。

@AndrewHenle：当然，很好的问题，是的，x86-64 的寄存器有点差，这就是为什么 AVX512 将矢量寄存器的数量和宽度加倍。是的，许多 RISC 具有 32 个 int / 32 fp regs，包括 AArch64，而 ARM 中的 16 个。但是这个特定的问题不需要很多 FP 寄存器，只需要 4 个用于 sum[0..3] 的向量累加器，可能还有 1 个临时的； x86 可以使用内存源 mul/add/FMA。寄存器重命名消除了任何 WAW 危害，因此我们可以重用相同的临时寄存器，而不需要软件流水线。 （而且 OoO exec 还隐藏了负载和 ALU 延迟。）

【参考方案1】：

此代码每个 FMA 有 2 次加载（如果发生 FMA 收缩），但理论上 Skylake 仅支持最多每个 FMA 一次加载（如果您想最大化 2/时钟 FMA 吞吐量） ，甚至在实践中通常也太多了。 （峰值是每个时钟 2 个负载 + 1 个存储，但通常不能完全维持）。请参阅英特尔的优化指南和https://agner.org/optimize/

循环开销不是最大的问题，主体本身强制代码以一半的速度运行。

如果k 循环是内部循环，则可以链接大量积累，而无需从C 加载/存储。这有一个缺点：使用像这样的循环携带的依赖链，将由代码来显式确保有足够的独立工作要做。

为了减少负载但有足够的独立工作，内部循环的主体可以计算来自 A 的小列向量和来自 B 的小行向量之间的乘积，例如使用 4 个标量广播从B 加载列和 2 个法向向量加载，导致 8 个独立 FMA 的加载只有 6 个（甚至可能更低的比率），这足以让 Skylake 保持快乐并且不会加载太多的独立 FMA。甚至 3x4 的足迹也是可能的，它也有足够的独立 FMA 让 Haswell 满意（它至少需要 10 个）。

我碰巧有一些示例代码，它用于单精度和 C++，但你会明白的：

sumA_1 = _mm256_load_ps(&result[i * N + j]);
sumB_1 = _mm256_load_ps(&result[i * N + j + 8]);
sumA_2 = _mm256_load_ps(&result[(i + 1) * N + j]);
sumB_2 = _mm256_load_ps(&result[(i + 1) * N + j + 8]);
sumA_3 = _mm256_load_ps(&result[(i + 2) * N + j]);
sumB_3 = _mm256_load_ps(&result[(i + 2) * N + j + 8]);
sumA_4 = _mm256_load_ps(&result[(i + 3) * N + j]);
sumB_4 = _mm256_load_ps(&result[(i + 3) * N + j + 8]);

for (size_t k = kk; k < kk + akb; k++) 
    auto bc_mat1_1 = _mm256_set1_ps(*mat1ptr);
    auto vecA_mat2 = _mm256_load_ps(mat2 + m2idx);
    auto vecB_mat2 = _mm256_load_ps(mat2 + m2idx + 8);
    sumA_1 = _mm256_fmadd_ps(bc_mat1_1, vecA_mat2, sumA_1);
    sumB_1 = _mm256_fmadd_ps(bc_mat1_1, vecB_mat2, sumB_1);
    auto bc_mat1_2 = _mm256_set1_ps(mat1ptr[N]);
    sumA_2 = _mm256_fmadd_ps(bc_mat1_2, vecA_mat2, sumA_2);
    sumB_2 = _mm256_fmadd_ps(bc_mat1_2, vecB_mat2, sumB_2);
    auto bc_mat1_3 = _mm256_set1_ps(mat1ptr[N * 2]);
    sumA_3 = _mm256_fmadd_ps(bc_mat1_3, vecA_mat2, sumA_3);
    sumB_3 = _mm256_fmadd_ps(bc_mat1_3, vecB_mat2, sumB_3);
    auto bc_mat1_4 = _mm256_set1_ps(mat1ptr[N * 3]);
    sumA_4 = _mm256_fmadd_ps(bc_mat1_4, vecA_mat2, sumA_4);
    sumB_4 = _mm256_fmadd_ps(bc_mat1_4, vecB_mat2, sumB_4);
    m2idx += 16;
    mat1ptr++;

_mm256_store_ps(&result[i * N + j], sumA_1);
_mm256_store_ps(&result[i * N + j + 8], sumB_1);
_mm256_store_ps(&result[(i + 1) * N + j], sumA_2);
_mm256_store_ps(&result[(i + 1) * N + j + 8], sumB_2);
_mm256_store_ps(&result[(i + 2) * N + j], sumA_3);
_mm256_store_ps(&result[(i + 2) * N + j + 8], sumB_3);
_mm256_store_ps(&result[(i + 3) * N + j], sumA_4);
_mm256_store_ps(&result[(i + 3) * N + j + 8], sumB_4);

这意味着j-loop 和i-loop 已展开，但k-loop 未展开，即使它现在是内部循环。稍微展开k-loop 确实对我的实验有所帮助。

【参考方案2】：

查看@harold 的回答以获得实际改进。这主要是为了转发我在cmets中写的内容。

在内循环中是四个而不是两个总和。 （为什么展开没有帮助？）

sum[i] 没有循环携带的依赖。下一次迭代分配sum[0] = _mm256_load_pd(C+i1*SIZE_K+j1+0);，它不依赖于前一个值。

因此，将相同架构寄存器的寄存器重命名到不同的物理寄存器足以避免可能使流水线停止的写后写危险。无需使用多个 tmp 变量使源复杂化。请参阅Why does mulss take only 3 cycles on Haswell, different from Agner's instruction tables? (Unrolling FP loops with multiple accumulators)（在那个问题中，2 个数组的一个点积，是一个循环通过累加器携带依赖关系。在那里，使用多个累加器 对隐藏 FP 很有价值FMA 延迟，因此我们限制了 FMA 吞吐量，而不是延迟。）

没有寄存器重命名的流水线（大多数有序 CPU）将受益于“软件流水线”，以静态调度无序 exec 可以动态执行的操作：加载到不同的寄存器所以每个负载和消耗它的 FMA 之间都有距离（充满了独立的工作）。然后在那个和商店之间。

但所有现代 x86 CPU 都是 OoO；甚至 Knight's Landing 也有一些有限的用于 SIMD 向量的 OoO exec。 （Silvermont 不支持 AVX，但会按顺序运行 SIMD 指令，仅对整数执行 OoO）。

---

在没有任何隐藏延迟的多累加器情况下，展开（显式在源中或使用 -fprofile-use 启用的 -funroll-loop 或默认情况下在 clang 中）的好处是：

减少前端带宽以使循环开销进入后端。每个循环开销更多有用的工作微量。因此，如果您的“有用工作”在前端接近瓶颈，它会有所帮助。 对运行循环开销的后端执行单元需求减少。通常在 Haswell 及更高版本或 Zen 上不是问题；当指令组合包括一些整数内容和一些纯加载指令时，后端基本上可以跟上前端。 每完成一项工作的总 uops 数减少意味着 OoO exec 可以“看到”更远的内存加载/存储。 有时对短运行循环的分支预测效果更好：迭代次数越少，分支预测学习的模式就越短。因此，对于短行程计数，当执行退出循环时，更有可能正确预测最后一次迭代的未采取措施。 有时在更复杂的情况下保存mov reg,reg，编译器更容易在不同的注册表中生成新结果。同一个变量可以在两个 reg 中交替使用，而不需要将 moved 恢复到同一个变量以准备下一次迭代。特别是如果您有一个循环使用 a[i] 和 a[i+1] 以依赖的方式，或类似斐波那契的方式。

循环中有 2 个负载 + 1 个存储，这可能是瓶颈，而不是 FMA 或前端带宽。展开 2 可能有助于避免前端瓶颈，但更多的只是与另一个超线程的争用有关。

---

在 cmets 中出现了一个有趣的问题：展开不需要很多寄存器才有用吗？

哈罗德评论：

16个寄存器数量不算多，但有12个就够了 累加器和来自 B 和广播的 3 个行向量 来自 A 的标量，所以它的结果就足够了。来自 OP 的循环 无论如何，上面几乎不使用任何寄存器。 32位的8个寄存器是 确实太少了。

当然，由于问题中的代码在循环迭代中的寄存器中没有“累加器”，只添加到内存中，编译器可以优化所有 sum[0..n] 以在 asm 中重用相同的寄存器；存储后它“死”了。所以实际套准压力非常低。

---

是的，x86-64 的寄存器有点少，这就是为什么 AVX512 将矢量寄存器的数量和宽度加倍 (zmm​​0..31)。是的，许多 RISC 有 32 个 int / 32 fp regs，包括 AArch64，而 ARM 是 16 个。

x86-64有16个标量整数寄存器（包括堆栈指针，不包括程序计数器），所以普通函数可以使用15个。向量寄存器也有16个，xmm0..15。 （对于 AVX，它们的宽度是 ymm0..15 的两倍）。

---

（其中一些是在我注意到sum[0..n] 毫无意义，而不是循环携带之前写的。）

在这种情况中，将寄存器重命名到大型物理寄存器文件就足够了。在其他情况下，拥有更多架构寄存器会有所帮助，尤其是对于更高的 FP 延迟，因此 AVX512 具有 32 zmm regs。但是对于整数 16 已经足够了。 RISC CPU 通常设计为按顺序设计，无需重新命名，需要 SW 流水线。

使用 OoO exec，在减少溢出/重新加载方面，从 8 到 16 个架构 GP 整数寄存器的跳跃比从 16 到 32 的跳跃更重要。 （我看过一篇论文，用不同数量的架构寄存器测量了 SPECint 的总动态指令数。我没有再查过，但是 8->16 可能总共节省了 10%，而 16->32 只是几个%。类似的东西）。

但是这个特定的问题不需要很多 FP 寄存器，sum[0..3] 只需要 4 个向量（如果它们是循环携带的），可能还有 1 个临时的； x86 可以使用内存源 mul/add/FMA。寄存器重命名消除了任何 WAW 危害，因此我们可以重用相同的临时寄存器，而不需要软件流水线。 （而且 OoO exec 还隐藏了负载和 ALU 延迟。）

当存在循环携带的依赖项时，您需要多个累加器。这段代码被添加到内存中，而不是一些向量累加器中，所以任何依赖都是通过存储/重新加载。但这只有大约 7 个周期的延迟，所以任何合理的缓存阻塞因素都会隐藏它。