使用 SIMD 和 openMP 加速矩阵乘法

Question

我正在尝试加快矩阵乘法的速度，因此它的性能比幼稚的实现要好得多。 我的目标是将它加速到 150 倍。 到目前为止，我在实现中尝试了以下内容：

• 1. 在连续块中分配矩阵元素以提高缓存效率。
• 2. 转置要按列访问的第二个矩阵，以将列排列到连续的存储块中。
• 3. 使用 SIMD 指令。
• 4. 使用 openMP 并行化 for 循环。

在第 1 步和第 2 步之后，我的 mat mul 比简单的实现快了 4 倍。 使用 SIMD 后，速度提高了 17 倍。 使用 openMP 后，速度提高了 56 倍。 我期待看到 openMP 的速度有更大的提升，比如至少 6 到 8 倍的提升。 我可能会错过什么？ 我的代码大致如下：

#pragma omp parallel for collapse(2)
    for (int i = 0; i < result.rows; i += 1) {
        for (int j = 0; j < result.cols; j += 1) {
          double product = 0.0;
          for (int k = 0; k < matrix1.cols / 4 * 4; k += 4) {
             //Use _mm256_mul_pd and _mm256_add_pd to process 4 elements at a time. 
          }
          //Do _mm256_store_pd and set the product. 
          result.mat[r][c] = product;
          for (int k = matrix1.cols / 4 * 4; k < matrix1.cols; k += 1) {
             //Tail case 
          }
        }
     }

我想将我的速度提高到至少 100 倍。 即比我当前的基准测试快 2 倍。 我还应该如何优化我的代码？

Answer 1

并行性只能给你这么多。 此外，顺序代码优化得越多，您可能从并行性中获得的感知收益就越少。 尽管如此，您可以改进的地方——我过去已经这样做了并且帮助改进了很多——是将矩阵乘法划分为更小的块。 因此，矩阵乘法被细分为较小矩阵（瓦片）的乘法。

因此，通过将矩阵乘法划分为较小的块，在其中执行较小子矩阵的矩阵乘法，您可以改进缓存的使用，包括temporal locality和spatial locality 。 您需要根据您使用的架构的缓存级别（例如， L1 、 L2和L3 ）的大小来划分矩阵。 您可以在这些幻灯片中查看有关缓存阻塞和矩阵乘法的更多详细信息。 关于内存，每个程序员都应该知道什么？ 在附录中还有一个矢量化缓存阻塞 matmul。

例如，如果您有一个在内核之间共享的 Cache L3 ，您可以将矩阵B多列加载到L3缓存中，并重用这些值来执行适合缓存L1和L2的较小切片的矩阵乘法. 您可以走得更远，在这些瓦片内部进一步划分瓦片，以便您可以利用寄存器。

优化矩阵乘法的内存使用后，您可以尝试从多线程中获得额外的加速。 如果你有一个多核集群，你可以尝试使用MPI + OpenMP进行并行化，当然此时你会遇到另一个瓶颈，进程之间的通信。

这一切都取决于您的代码运行所在的架构，如果您有一个NUMA架构，那么您还必须考虑本地和非本地内存等因素。 您还可以探索 GPU 路线： 使用 Nvidia CUDA 在 GPU 上进行矩阵-矩阵乘法。

看看BLAS可以很好地了解高效代码。