使用OpenMP与AVX优化矩阵乘法

使用OpenMP与AVX优化矩阵乘法

由于课设内容做的太过简（mo）单（yu），于是在去年12月初的时候就计划写三篇博客随笔作为实验报告，前两篇简单介绍了OpenMP和SIMD指令进行铺垫，本篇将会介绍他的应用场景——优化矩阵乘法。

前两篇的传送门：

1.OpenMP优化for循环的基础运用

2.SSE与AVX指令基础介绍与使用

---

1.循环顺序调整

先来看看最基本的矩阵乘法。

void MatrixMul(int **A, int **B, int **C, int size)
{
    for (int i = 0; i < size; i++)
        for (int j = 0; j < size; j++)
            for (int k = 0; k < size; k++)
                C[i][j] += A[i][k] * B[k][j];
}

这是一个普通的方阵的乘法函数，用于计算矩阵C = A x B。

线性代数中的矩阵乘法原理告诉我们，只要遵循下面这个原则进行计算就能得到正确结果。

C[i][j] += A[i][k] * B[k][j]

所以，改变循环中ijk的遍历顺序并不会影响结果的正确性。又因为后续使用AVX时需要载入一段连续的内存空间，当前ijk的顺序下最内层循环改变的是k，这就导致B是逐行也就是非连续访问。这既不利于AVX载入数据，也不利于Cache访问。因此我们这里将循环顺序改为ikj，这样就能实现内存的连续访问。

void MatrixMul(int **A, int **B, int **C, int size)
{
    // 改为ikj的循环顺序
    for (int i = 0; i < size; i++)
        for (int k = 0; k < size; k++)
            for (int j = 0; j < size; j++)
                C[i][j] += A[i][k] * B[k][j];
}

---

2.使用AVX指令优化

2.1分析

在开始之前，若是不熟悉AVX指令的用法的话，可以先回顾上一篇的内容SSE与AVX指令基础介绍与使用。

在上一步进行调整循环顺序之后，内部已经能够连续访问数据了。现在我们就需要分析，如何用AVX将多次循环合并以提升效率。

for (int j = 0; j < size; j++)
    C[i][j] += A[i][k] * B[k][j];

我们可以看到，循环最内层的操作有两步。第一步，将A[i][k]乘上B[k][0]_{B[k][size-1]。第二部，将第一步乘法的结果加到C[i][0]}C[i][size-1]中。

2.2乘法优化

由于最内层循环中A[i][k]并没有改变，所以只需要在循环开始的时候进行一次载入。并且每个数字B[k][j]都是与相同的A[i][k]相乘，所以每个单元都要写入相同的A[i][k]，这就需要使用_mm256_set1_epi32方法进行复制型的加载。

__m256i ra = _mm256_set1_epi32(A[i][k]);

A[i][k]加载完毕后就需要加载B矩阵的内容，这里使用上篇介绍的类型转换就能完成。

// 两种取址方法
__m256i rb = *(__m256i *)(B[k] + j);
__m256i rb = *(__m256i *)(&B[k][j]);

乘法方面注意要使用mullo才能得到每一位相乘的结果。mullo是保存每一个结果的低位，mulhi是保存高位，而mul则是用双倍的空间同时保存低位和高位，这样会导致只能保存一半的乘法结果。

rb = _mm256_mullo_epi32(ra, rb);

2.3加法优化

加法和乘法同理，读取C矩阵的内容。

// 两种取址方法
__m256i rc = *(__m256i *)(C[i] + j);
__m256i rc = *(__m256i *)(&C[i][j]);

和乘法的结果相加。

rc = _mm256_add_epi32(rb, rc);

最后用类型转换的方式写回。

*(__m256i *)(C[i] + j) = rc;

2.4最终结果

AVX一次处理256位数据，也就是8个int，因此循环的步进也需要改为8。另外为了防止越界访问，还需要判断剩余的数据是否大于8个，当内容不足8个时则使用普通的逐个相乘计算。

综上所属，初步的优化结果如下。别忘了32位内存对齐！

void MatrixMul(int **A, int **B, int **C, int size)
{
    for (int i = 0; i < size; i++)
        for (int k = 0; k < size; k++)
        {
            int j = 0;

            // AVX优化
            __m256i ra = _mm256_set1_epi32(A[i][k]);
            for (; j <= size - 8; j += 8)
            {
                // 乘法
                __m256i rb = *(__m256i *)(B[k] + j);
                rb = _mm256_mullo_epi32(ra, rb);
                // 加法
                __m256i rc = *(__m256i *)(C[i] + j);
                rc = _mm256_add_epi32(rb, rc);
                // 写回
                *(__m256i *)(C[i] + j) = rc;
            }

            // 处理余下内容
            for(;j<size;j++)
                C[i][j] += A[i][k] * B[k][j];
        }
}

上面的代码为了表述思路用了很多不必要的中间步骤和变量，为了提高运行效率还能进一步压行简化。

void MatrixMul(int **A, int **B, int **C, int size)
{
    for (int i = 0; i < size; i++)
        for (int k = 0; k < size; k++)
        {
            int j = 0;

            // AVX优化
            __m256i ra = _mm256_set1_epi32(A[i][k]);
            for (; j <= size - 8; j += 8)
            {
                // 简化版
                *(__m256i *)(C[i] + j) = _mm256_add_epi32(*(__m256i *)(C[i] + j), _mm256_mullo_epi32(*(__m256i *)(B[k] + j), ra));
            }

            // 处理余下内容
            for (; j < size; j++)
                C[i][j] += A[i][k] * B[k][j];
        }
}

接下来看看优化前后的性能差距。

最后得到了近6倍的性能提升，提升幅度还是非常大的。这时候有读者可能会好奇，明明数据吞吐量是原来的8倍，为什么性能却只能提升6倍呢？这其实从简化前后的性能提升也可以看出原因，简化版减少了过程的中间变量，也就是减少了额外的读写开销，但这种开销并不能完全避免，CPU最终还是需要将数据先载入到256位寄存器中计算，然后再从中读取内容写回内存，这种的无法避免额外开销导致了程序始终无法达到8倍的性能提升。

---

3.使用OpenMP优化

相较于AVX，OpenMP在使用方面就简单了许多，第一篇中我们也对其进行了简单的介绍OpenMP优化for循环的基础运用。

在这里只需要在最外层循环上放加上一句#pragma omp parallel for 即可，当然也可以在末尾加上num_threads(N)设置使用的线程数为N。

void MatrixMul(int **A, int **B, int **C, int size)
{
    #pragma omp parallel for num_threads(4) // 设置为4线程
    for (int i = 0; i < size; i++)
        for (int k = 0; k < size; k++)
            for (int j = 0; j < size; j++)
                C[i][j] += A[i][k] * B[k][j];
}

接下来我们来测试一下OpenMP在不同线程下的性能变化，测试平台使用的是4C/8T的CPU。

可以看到，OpenMP在2和4线程时性能提升最大，8线程达到顶峰，后续开始出现下降。其中2线程时最接近理论性能提升，到4线程时则开始和理论性能差距拉大。这是由于多线程任务需要系统进行调度，线程调度和开关的损耗会随着则线程数增加而增长，并且当线程数超过物理CPU规模时也无法达到实际的并行效果，从而使提升速度放缓乃至出现下降。

---

4.综合OpenMP与AVX

在前面分别使用了OpenMP和AVX之后，现在我们来将二者同时运用起来，只需要在AVX的基础上在最外层循环上方加上OpenMP指令即可。

void MatrixMul(int **A, int **B, int **C, int size)
{
    #pragma omp parallel for num_threads(4) // OpenMP优化(4线程)
    for (int i = 0; i < size; i++)
        for (int k = 0; k < size; k++)
        {
            int j = 0;

            // AVX优化
            __m256i ra = _mm256_set1_epi32(A[i][k]);
            for (; j <= size - 8; j += 8)
            {
                // 简化版
                *(__m256i *)(C[i] + j) = _mm256_add_epi32(*(__m256i *)(C[i] + j), _mm256_mullo_epi32(*(__m256i *)(B[k] + j), ra));
            }

            // 处理余下内容
            for (; j < size; j++)
                C[i][j] += A[i][k] * B[k][j];
        }
}

然后来看看同时运用上OpenMP和AVX后的性能变化。

最后的提升基本就是二者实际提升的结果相乘，还是挺理想的。

---

5.实际应用(图像翻转程序)

最后来看一个实际应用，这里使用的是一个图像翻转程序，实现原理就是先读取图像的rgb值并用矩阵保存，之后再将图像矩阵乘一个转置矩阵实现翻转效果。而本次实验的内容，就是通过加速矩阵乘法的方式加速图像翻转程序。

图像的rgb值分别用3个char矩阵保存，但是由于Intel的SIMD指令集中并没有epi8的乘法指令，所以最低只能使用16位的short数组进行AVX加速，所以本次的AVX优化的for循环步进提升到了一次16个数据。

---

本文发布于2023年1月4日

最后修改于2023年1月4日