使用OpenMP与AVX优化矩阵乘法

使用OpenMP与AVX优化矩阵乘法

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

前两篇的传送门:

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];
        }
}

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

img

最后得到了近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。

img

可以看到,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后的性能变化。

img

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


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

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

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

img


本文发布于2023年1月4日

最后修改于2023年1月4日

posted @ 2023-01-04 22:39  千松  阅读(988)  评论(0编辑  收藏  举报