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

Matrix Multiplication using SIMD vectors in C++

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

The templated code below implements the innermost loops that calculate a patch of size regA x regB in matrix C. The code loads regA scalars from matrixA and regB SIMD-width vectors from matrix B. The program uses Clang's extended vector syntax.

/// 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

文章的第 2 条评论链接到 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的load/broadcast函数。有了那个header,你应该可以编译文章中的代码了。

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

这篇文章本可以添加 1 个小节来解释这一点,但它似乎更侧重于性能细节,并且对解释如何使用语法并不感兴趣。

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

您可以使用多种手动向量方式中的任何一种来完成这些事情:GNU C 本机向量或 clang 非常相似的“扩展”向量,或可移植的 Intel 内在函数。