在 C++ 中使用 SIMD 向量进行矩阵乘法

Question

我目前正在 github 上阅读一篇关于使用 Clang 的扩展向量语法进行性能优化的文章。作者给出如下代码sn-p：

下面的模板化代码实现了计算矩阵 C 中大小为 regA x regB 的补丁的最内层循环。代码从矩阵 A 加载 regA 标量，从矩阵 B 加载 regB SIMD 宽度向量。程序使用 Clang 的扩展向量语法。

/// Compute a RAxRB block of C using a vectorized dot product, where RA is the
/// number of registers to load from matrix A, and RB is the number of registers
/// to load from matrix B.
template <unsigned regsA, unsigned regsB>
void matmul_dot_inner(int k, const float *a, int lda, const float *b, int ldb,
                      float *c, int ldc) {
  float8 csum[regsA][regsB] = {{0.0}};
  for (int p = 0; p < k; p++) {

    // Perform the DOT product.
    for (int bi = 0; bi < regsB; bi++) {
      float8 bb = LoadFloat8(&B(p, bi * 8));
      for (int ai = 0; ai < regsA; ai++) {
        float8 aa = BroadcastFloat8(A(ai, p));
        csum[ai][bi] += aa * bb;
      }
    }
  }

  // Accumulate the results into C.
  for (int ai = 0; ai < regsA; ai++) {
    for (int bi = 0; bi < regsB; bi++) {
      AdduFloat8(&C(ai, bi * 8), csum[ai][bi]);
    }
  }
}

下面概述的代码最让我困惑。我读了整篇文章，理解了使用阻塞和计算小补丁背后的逻辑，但我不能完全理解这个位是什么意思：

    // Perform the DOT product.
for (int bi = 0; bi < regsB; bi++) {
  float8 bb = LoadFloat8(&B(p, bi * 8)); //the pointer to the range of values?
  for (int ai = 0; ai < regsA; ai++) {
    float8 aa = BroadcastFloat8(A(ai, p));
    csum[ai][bi] += aa * bb;
  }
}

}

谁能详细说明这里发生了什么？ 文章可以找到here

Answer 1

文章的第二条评论链接到https://github.com/pytorch/glow/blob/405e632ef138f1d49db9c3181182f7efd837bccc/lib/Backends/CPU/libjit/libjit_defs.h#L26，它将float8类型定义为

typedef float float8 __attribute__((ext_vector_type(8)));

（类似于 immintrin.h 定义 __m256 的方式）。并且定义了类似于_mm256_load_ps和_mm256_set1_ps的加载/广播功能。有了这个头文件，你应该可以编译文章中的代码了。

见Clang's native vector documentation。 GNU C 本机向量语法是获得重载 * 运算符的好方法。我不知道clang的ext_vector_type GCC/clang/ICC float __attribute__((vector_width(32)))（32 字节宽度）不会。

这篇文章本来可以加一小节来解释这一点，但似乎更侧重于性能细节，并没有真正对解释如何使用语法感兴趣。

文章中的大部分讨论都是关于如何使用 SIMD 向量手动向量化 matmul 以提高缓存效率。从我快速浏览的角度来看，这部分看起来不错。

您可以使用多种手动向量方法中的任何一种来执行这些操作：GNU C 原生向量或 clang 非常相似的“扩展”向量，或可移植的 Intel 内在函数。