在 C 中使用 AVX 实现矩阵运算

Question

我正在尝试使用 AVX 实现以下操作：

for (i=0; i<N; i++) {
  for(j=0; j<N; j++) {
    for (k=0; k<K; k++) {
      d[i][j] += 2 * a[i][k] * ( b[k][j]- c[k]);
    }
  }
}

for (int i=0; i<N; i++){
   f+= d[ind[i]][ind[i]]/2;
}    

其中 d 是 NxN 矩阵，a 是 NxK，ba KxN 和 ca 长度为 K 的向量。它们都是双精度数。 当然，所有数据都是对齐的，我使用#pragma vector aligned来帮助编译器（gcc）。

我知道如何将 AVX 扩展与一维数组一起使用，但对我来说用矩阵来做这件事有点棘手。 目前，我有以下内容，但我没有得到正确的结果：

    for (int i=0; i< floor (N/4); i++){
        for (int j=0; j< floor (N/4); j++){
            __m256d D, A, B, C;
            D = _mm256_setzero_pd();
            #pragma vector aligned
            for (int k=0; k<K_MAX; k++){
                A = _mm256_load_pd(a[i] + k*4);
                B = _mm256_load_pd(b[k] + j*4);
                C = _mm256_load_pd(c + 4*k);
                B = _mm256_sub_pd(B, C);
                A = _mm256_mul_pd(A, B);
                D = _mm256_add_pd(_mm256_set1_pd(2.0), A);
                _mm256_store_pd(d[i] + j*4, D);
            }

        }
    }


    for (int i=0; i<N; i++){
        f+= d[ind[i]][ind[i]]/2;
    }    

我希望有人能告诉我错误在哪里。

提前致谢。

注意：我不愿意介绍 OpenMP，只使用 SIMD 英特尔指令

Answer 1

假设 N 和 K 数字都相对较大（远大于硬件矢量大小的 4），这是矢量化主循环的一种方法。 未经测试。

主要思想是矢量化中间循环而不是内部循环。 这样做有两个原因。

1. 这避免了水平操作。 当仅对内部循环进行向量化时，我们必须计算向量的水平和。

2. 当加载 4 个连续的k值时， b[k][j]加载具有不幸的 RAM 访问模式，需要 4 个单独的加载指令，或者收集加载，这两种方法都相对较慢。 为 4 个连续的j值加载元素是一个全向量加载指令，非常有效，尤其是在您对齐输入时。

    const int N_aligned = ( N / 4 ) * 4;
    for( int i = 0; i < N; i++ )
    {
        int j = 0;
        for( ; j < N_aligned; j += 4 )
        {
            // Load 4 scalars from d
            __m256d dv = _mm256_loadu_pd( &d[ i ][ j ] );

            // Run the inner loop which only loads from RAM but never stores any data
            for( int k = 0; k < K; k++ )
            {
                __m256d av = _mm256_broadcast_sd( &a[ i ][ k ] );
                __m256d bv = _mm256_loadu_pd( &b[ k ][ j ] );
                __m256d cv = _mm256_broadcast_sd( &c[ k ] );

                // dv += 2*av*( bv - cv )
                __m256d t1 = _mm256_add_pd( av, av );   // 2*av
                __m256d t2 = _mm256_sub_pd( bv, cv );   // bv - cv
                dv = _mm256_fmadd_pd( t1, t2, dv );
            }
            // Store the updated 4 values
            _mm256_storeu_pd( &d[ i ][ j ], dv );
        }

        // Handle remainder with scalar code
        for( ; j < N; j++ )
        {
            double ds = d[ i ][ j ];
            for( int k = 0; k < K; k++ )
                ds += 2 * a[ i ][ k ] * ( b[ k ][ j ] - c[ k ] );
            d[ i ][ j ] = ds;
        }
    }

如果您想进一步优化，请尝试将内部循环展开像 2 这样的小因子，使用 2 个用_mm256_setzero_pd()初始化的独立累加器，在循环之后添加它们。 可能是在某些处理器上，此版本在 FMA 指令的延迟上停止，而不是使加载端口或 ALU 饱和。 多个独立的累加器有时会有所帮助。

Answer 2

b[k][j]是你的问题：元素b[k + 0..3][j]在内存中不连续。 使用 SIMD（以合理/有用的方式）不是您可以放入经典的幼稚 matmul 循环的东西。 看看每个程序员都应该知道什么关于内存的知识？ - 有一个带有 SSE2 matmul 示例（带有缓存阻塞）的附录，它显示了如何以对 SIMD 友好的不同顺序进行操作。

Soonts 的回答显示了如何通过对中间循环j进行矢量化来进行矢量化。 但这留下了相对较差的内存访问模式，循环内有 3 次加载 + 3 次 ALU 操作。 （这个答案开始是对它的评论，请参阅我正在谈论的代码并提出更改。）

循环反转应该可以将j作为最内层循环。 这意味着在最里面的循环中为d[i][j] += ...存储，但是 OTOH 它在2 * a[i][k] * ( b[k][j]- c[k] ) ，因此您可以有效地转换为d[i][j] += (2*a_ik) * b[k][j] - (2*a_ik*c_k) ，即一个 VFMSUBPD 和一个 VADDPD每个加载和存储。 （将bv加载折叠到 FMSUB 作为内存源操作数，并将dv加载折叠到 VADDPD 中，因此希望前端只有 3 个微指令，包括单独的存储，不包括循环开销。）

编译器必须展开并避免索引寻址模式，以便存储地址 uop 可以保持微融合并在 Intel CPU（Haswell 到 Skylake 系列）的端口 7 上运行，而不是与两个负载竞争。 Ice Lake 没有这个问题，它有两个完全独立的存储 AGU，与两个负载 AGU 分开。 但可能仍需要一些循环展开以避免前端瓶颈。

这是一个未经测试的示例（由 Soonts 提供的原始版本，谢谢）。 它以不同的方式优化循环中的 2 个 FP 数学运算：只需将2*a提升到循环之外，然后为dv += (2av)*(sub_result)执行 SUB 然后 FMA。 但是bv不能作为vsubpd的源操作数，因为我们需要bv - cv 。 但是我们可以通过否定cv以允许(-cv) + bv在内循环中解决这个问题，其中bv作为内存源操作数。 有时编译器会为你做类似的事情，但在这里他们似乎没有，所以我手动做了。 否则，我们会通过前端获得单独的vmovupd负载。

#include <stdint.h>
#include <stdlib.h>
#include <immintrin.h>

// This double [N][N] C99 VLA syntax isn't portable to C++ even with GNU extensions
// restrict tells the compiler the output doesn't overlap with any of the inputs
void matop(size_t N, size_t K, double d[restrict N][N], const double a[restrict N][K], const double b[restrict K][N], const double c[restrict K])
{
    for( size_t i = 0; i < N; i++ ) {
        // loop-invariant pointers for this outer iteration
        //double* restrict rowDi = &d[ i ][ 0 ];
        const double* restrict rowAi = &a[ i ][ 0 ];
        for( size_t k = 0; k < K; k++ ) {
            const double* restrict rowBk = &b[ k ][ 0 ];
            double* restrict rowDi = &d[ i ][ 0 ];
#if 0  // pure scalar
            // auto-vectorizes ok; still a lot of extra checking outside outermost loop even with  restrict
            for (size_t j=0 ; j<N ; j++){
                rowDi[j] += 2*rowAi[k] * (rowBk[j] - c[k]);
            }
#else   // SIMD inner loop with cleanup

         // *** TODO: unroll over 2 or 3 i values
         // and maybe also 2 or 3 k values, to reuse each bv a few times while it's loaded.
            __m256d av = _mm256_broadcast_sd( rowAi + k );
            av = _mm256_add_pd( av, av );   // 2*a[ i ][ k ] broadcasted
            const __m256d cv = _mm256_broadcast_sd( &c[ k ] );
            const __m256d minus_ck = _mm256_xor_pd(cv, _mm256_set1_pd(-0.0));    // broadcasted -c[k]

            //const size_t N_aligned = ( (size_t)N / 4 ) * 4;
            size_t N_aligned = N & -4;    // round down to a multiple of 4 j iterations
            const double* endBk = rowBk + N_aligned;
            //for( ; j < N_aligned; j += 4 )
            for ( ; rowBk != endBk ; rowBk += 4, rowDi += 4) {  // coax GCC into using pointer-increments in the asm, instead of j+=4
                // Load the output vector to update
                __m256d dv = _mm256_loadu_pd( rowDi );
                // Update with FMA
                __m256d bv = _mm256_loadu_pd( rowBk );
                __m256d t2 = _mm256_add_pd( minus_ck, bv );   // bv - cv
                dv = _mm256_fmadd_pd( av, t2, dv );
                // Store back to the same address
                _mm256_storeu_pd( rowDi, dv );
            }
            // rowDi and rowBk point to the double after the last full vector

            // The remainder, if you can't pad your rows to a multiple of 4 and step on that padding
            for(int j=0 ; j < (N&3); j++ )
                rowDi[ j ] += _mm256_cvtsd_f64( av ) * ( rowBk[ j ] + _mm256_cvtsd_f64( minus_ck ) );
#endif
        }
    }
}

如果不展开（ https://godbolt.org/z/6WeYKbnYY），GCC11的内循环 asm 看起来像这样，所有单 uop 指令即使在 Haswell 及更高版本的后端也可以保持微融合。

.L7:                                             # do{
        vaddpd  ymm0, ymm2, YMMWORD PTR [rax]     # -c[k] + rowBk[0..3]
        add     rax, 32                           # rowBk += 4
        add     rdx, 32                           # rowDi += 4
        vfmadd213pd     ymm0, ymm1, YMMWORD PTR [rdx-32]  # fma(2aik, Bkj-ck, Dij)
        vmovupd YMMWORD PTR [rdx-32], ymm0        # store FMA result
        cmp     rcx, rax
        jne     .L7                              # }while(p != endp)

但它总共有 6 个 uops，其中 3 个循环开销（指针增量和融合 cmp+jne），因此通过 Skylake 的 Haswell 只能以每 1.5 个时钟 1 次迭代运行它，在前端的 4-wide issue 阶段遇到瓶颈。 （这不会让 OoO exec 提前执行指针增量和循环分支，以便在后端仍在处理旧负载和 FP 数学时尽早注意到并恢复。）

---

所以循环展开应该会有所帮助，因为我们设法哄骗 GCC 使用索引寻址模式。 没有它，它在英特尔 Haswell/Skylake CPU 上的 AVX 代码相对没用，每个vaddpd ymm5, ymm4, [rax + r14]解码为 1 个微融合 uop，但在后端分解成 2 个有问题的，对我们没有帮助通过前端最窄的部分获得更多的工作。 （很像如果我们使用单独的vmovupd加载，就像我们使用_mm256_sub_pd(bv, cv)而不是add(bv, -cv)一样。）

vmovupd ymmword ptr [rbp + r14], ymm5存储保持微融合但不能在端口 7 上运行，这将我们限制为每个时钟总共 2 个内存操作（其中最多 1 个可以是存储。）所以每个向量 1.5 个周期的最佳情况。

使用 GCC 和 clang -O3 -march=skylake -funroll-loops在https://godbolt.org/z/rd3rn9zor上编译。 GCC 确实使用指针增量，负载折叠成 8x vaddpd和 8x vfmadd213pd 。 但是 clang 使用索引寻址模式并且不会展开。 （您可能不希望整个程序都使用-funroll-loops ，因此要么单独编译它，要么手动展开。GCC 的展开完全剥离了在进入实际 SIMD 循环之前执行 0..7 次向量迭代的序言，因此它非常激进.)

GCC 的循环展开对于大 N 在这里看起来很有用，在多个向量上分摊指针增量和循环开销。 （例如，GCC 不知道如何在点积中为 FP dep 链发明多个累加器，在这种情况下使其展开毫无用处，这与 clang 不同。）

不幸的是，clang 并没有为我们展开内部循环，但它确实以一种有趣的方式使用vmaskmovpd进行清理。

我们使用单独的循环计数器进行清理可能会很好，这样编译器可以轻松证明清理的行程计数是 0..3 ，因此它不会尝试使用 YMM 进行自动矢量化。

另一种方法是使用实​​际的j变量进行内部循环及其清理，更像是 Soonts 的编辑。 IIRC，编译器确实尝试为此自动矢量化清理，浪费代码大小和一些总是错误的分支。

            size_t j = 0;     // used for cleanup loop after
            for( ; j < N_aligned; j += 4 )
            {
                // Load the output vector to update
                __m256d dv = _mm256_loadu_pd( rowDi + j );
                // Update with FMA
                __m256d bv = _mm256_loadu_pd( rowBk + j );
                __m256d t2 = _mm256_sub_pd( bv, cv );   // bv - cv
                dv = _mm256_fmadd_pd( av, t2, dv );
                // Store back to the same address
                _mm256_storeu_pd( rowDi + j, dv );
            }
            // The remainder, if you can't pad your rows to a multiple of 4
            for( ; j < N; j++ )
                rowDi[ j ] += _mm256_cvtsd_f64( av ) * ( rowBk[ j ] - _mm256_cvtsd_f64( cv ) );

对于现代 CPU（ https://agner.org/optimize/和https://uops.info/ ），尤其是英特尔，我们可以进行 2 次加载和1 次存储，这具有相当好的加载和存储与 FP 数学组合。 我认为 Zen 2 或 3 也可以做 2 个加载 + 1 个存储。 不过，它需要命中 L1d 缓存以维持这种吞吐量。 （即便如此，英特尔的优化手册称 Skylake 上的最大持续 L1d 带宽低于所需的完整 96 字节/周期。更像是 80 年代中期 IIRC，因此我们不能完全期望每个周期有一个结果向量，即使有足够的展开以避免前端瓶颈。）

没有延迟瓶颈，因为我们每次迭代都转到一个新的dv ，而不是在循环迭代中累积任何东西。

这样做的另一个优点是对d[i][j]和b[k][j]的内存访问将是顺序的，在最内层循环中没有其他内存访问。 （中间循环将对a[i][k]和c[k]进行广播加载。如果内部循环驱逐过多，这些似乎可能会缓存未命中；随着外部循环的一些展开，一个 SIMD 加载和一些改组可能会有所帮助，但缓存阻塞可能会避免这种需要。）

对不同的b[k]行重复循环相同的d[i]行为我们正在修改的部分提供了局部性（即使用k作为中间循环，保持i作为最外层。）以k作为外部循环，我们将在整个d[0..N-1][0..N-1]上循环K次，可能需要写入 + 读取每个通道一直到任何级别的缓存或内存能坚持住。

但实际上，如果每一行真的很长，您仍然希望缓存块，因此您可以避免缓存未命中以将所有b[][]从 DRAM 中带入 N 次。 并且避免驱逐你接下来要广播加载的东西。

---

更智能的展开：缓存阻塞的第一步

如果我们在加载每个数据向量时对其进行更多处理，则上述一些关于最大化加载/存储执行单元吞吐量以及要求编译器使用非索引寻址模式的问题可能会消失。

例如，我们可以处理 2、3 或 4，而不是只处理d[][]的一行。然后每个(rowBk[j] - c[k])结果都可以使用很多次（使用不同的2aik ) 用于d[i+unroll][j + 0..vec]向量。

我们还可以加载几个不同的(rowBk+K*0..unroll)[j+0..3] ，每个都有一个对应的minus_ck0 、 minus_ck1等（或者保留一个向量数组；只要它很小并且编译器有足够的寄存器，这些元素不会存在于内存中。）

在寄存器中同时使用多个bv-cv和dv向量，我们可以在每次加载时执行更多的 FMA，而不会增加 FP 工作的总量。 但是，它需要更多的常量寄存器，否则我们可能会通过强制重新加载来达到目的。

d[i][j] += (2*a_ik) * b[k][j] - (2*a_ik*c_k)变换在这里没有用； 我们希望将bv-cv与i分开，以便我们可以将该结果作为不同 FMA 的输入重用。

b[k][j]+(-c[k])仍然可以从负载与vaddpd的微融合中受益，因此理想情况下它仍将使用指针增量，但前端可能不再是瓶颈.

不要过度使用它； 太多的内存输入流可能会成为缓存冲突未命中的问题，尤其是对于一些可能会产生别名的 N 值，以及硬件预取跟踪它们。 （尽管据说英特尔的 L2 流媒体每 4k 页跟踪 1 个前向流和 1 个后向流，但 IIRC。）大概 4 到 8 个 ish 流是可以的。 但是，如果 L1d 中没有丢失d[][] ，那么它实际上并不是来自内存的输入流。 但是，您不希望您的b[][]输入行驱逐d数据，因为您将重复循环 2 到 4 行d数据。

---

相比之下：Soonts 的循环 - 清理频率较低，但内存访问模式较差。

Soonts 的当前循环具有 3 次加载和 3 次 ALU 操作，虽然每次 FMA 操作 1 次加载如果它们在缓存中命中就已经可以了（大多数现代 CPU 每个时钟可以执行 2 次，尽管 AMD Zen 也可以添加 2 次 FP与 mul/fma 平行）。 如果那个额外的 ALU 操作是一个瓶颈，我们可以将a[][]预乘 2 一次，只需要O(N*K)的工作而不是O(N^2*K)的运行时间。 但这可能不是瓶颈，因此不值得。

更重要的是，Soonts 的当前答案中的内存访问模式是一次循环向前 1 双，用于广播负载的c[k]和a[i][k] ，这很好，但是bv = _mm256_loadu_pd的b[k][j + 0..3]不幸地跨过一列。

如果你要按照 Soonts 的建议展开，不要只为一个dv做两个 dep 链，至少做两个向量， d[i][j + 0..3]和4..7所以你使用您触摸的每个b[k][j]的整个 64 字节（完整缓存行）。 或者一对缓存线的四个向量。 （英特尔 CPU 至少使用相邻行预取器，它喜欢完成一对 128 字节对齐的高速缓存行，因此您将从将b[][]的行对齐 128 行中受益。或者至少对齐 64 行，并从相邻行预取中获得一些好处。

如果b[][]的垂直切片适合某个级别的缓存（以及您当前正在累积的d[i][]行），则下一组列的下一步可以从该预取中受益和地方。 如果没有，那么充分利用你接触到的线条更重要，这样它们以后就不必再被拉进去了。

因此，使用 Soonts 的向量化策略，对于不适合 L1d 缓存的大型问题，确保b的行对齐 64 可能很好，即使这意味着在每行的末尾进行填充。 （存储几何不必与实际矩阵维度匹配；您分别传递N和row_stride 。您使用一个用于索引计算，另一个用于循环边界。）