Skip to content

SIMD 向量化

Warning

⭐题目摘自2024-zju-hpc-summercamp，解答仅供参考。

基础知识

现代处理器一般都支持向量化指令，x86 架构下 Intel 和 AMD 两家的处理器都提供了诸如 SSE，AVX 等 SIMD 指令集，一条指令可以同时操作多个数据进行运算，大大提高了现代处理器的数据吞吐量。

现代编译器在高优化等级下，具有自动向量化的功能，对于结构清晰，循环边界清晰的程序，编译器的自动向量化已经可以达到很优秀的程度了。然而，编译器的优化始终是保守的，很多情况下编译器无法完成使用 SIMD 指令进行向量化的工作，为了追求性能，高性能计算领域经常需要手写 SIMD 代码进行代码优化。

显然直接手写汇编指令过于困难，在 C 语言环境下，Intel 提供了一整套关于 SIMD 指令的函数封装接口和指令相关行为的参照手册，可以在实验文档的参考资料中找到。

使用这些函数 API 需要 include 对应的头文件，不同 SIMD 指令集需要的头文件不同，具体需要参考 Intel 相关文档。

C
#include <smmintrin.h>
#include <emmintrin.h>
#include <immintrin.h>

另外深入到这个级别的优化已经开始需要考虑具体处理器的体系结构细节了，如某个架构下某条指令的实现延时和吞吐量是多少，处理器提供了多少向量寄存器，访存的对齐等等。这种时候编译器具体产生的汇编代码能比 C 语言代码提供更多的信息，你能了解到自己使用了多少寄存器，编译器是否生成了预期外的代码等等。

参考资料中提供的 godbolt 是一款基于 web 的研究不同编译器编译产生汇编代码的工具，在进行本实验的时候可以学习使用。

实验任务

用SIMD改写进行 MAXN 次 \(D_{4*4}=A_{4*12}*B_{12*4}\) 的矩阵乘法，实现向量化。

提示

其中 MAXN 个矩阵的内容不同，计算的时候要注意位移。

因为该题只是为了使用手写 SIMD, 由此不接受其他优化方法，包括 omp 等。

初始代码详见starter_code

基准代码

C
for (int n = 0; n < MAXN; n++)
    {
        /* 可以修改的代码区域 */
        // -----------------------------------
        for (int k = 0; k < 4; k++)
        {
            for (int i = 0; i < 4; i++)
            {
                for (int j = 0; j < 12; j++)
                {
                    *(d + n * 16 + i * 4 + k) += *(a + n * 48 + i * 12 + j) * \
                                                    *(b + n * 48 + j * 4 + k);
                }
            }
        }
        // -----------------------------------
    }

编译

在编译时添加以下选项可以允许编译器生成使用 AVX2 和 FMA 指令集的代码，如果你使用了其它不在编译器默认范围内的指令集，类似的编译选项是必要的。

Bash
-mavx2 -mfma

参照 Intel 文档优化完成后就可以开始测试优化前和优化后的性能差距，最后对比前后编译器产生的汇编代码的不同即可完成实验。
⭐能快速生成汇编代码的网站：Compiler Explorer

参考方案

思路

思路非常简单，主要是熟悉intel指令集查找。

代码

分析待优化代码：循环实现的是d[ i , k ] = a[ i , j ] * b[ j , k ],这时候原代码数组 b 的内存访问是不连续的。
因此首先调整循环顺序。提升访存命中率，加快数据读取速度。

C++
for (long long i = 0; i < 4; i++)
{
    for (long long j = 0; j < 12; j++)
    {
        for (long long k = 0; k < 4; k++)
        {
            *(d + n * 16 + i * 4 + k) += *(a + n * 48 + i * 12 + j) * *(b + n * 48 + j * 4 + k);
        }
    }
}
继续考虑SIMD优化，发现最内层循环刚好是四个64位数据一组，共256位。可以直接将最内层循环作SIMD处理，一次加载4个数据。
⭐由于数组 d 取值与 j 无关，故在 i 层循环进行加载，避免在 j 层重复读取数据，造成内存访问变慢。
⭐注意到数组 a 在运算中每四个元素取相同值。故用 _mm256_set1_pd 进行加载，读取四个相同数据。

得到代码如下：

C++
for (long long i = 0; i < 4; i++)
{
    __m256d rd = _mm256_loadu_pd(d + n * 16 + i * 4);
    for (long long j = 0; j < 12; j++)
    {
        // AVX优化
        __m256d ra = _mm256_set1_pd(*(a + n * 48 + i * 12 + j));
        __m256d rb = _mm256_loadu_pd(b + n * 48 + j * 4);
        rd = _mm256_add_pd(rd, _mm256_mul_pd(rb, ra));
        _mm256_storeu_pd(d + n * 16 + i * 4, rd);
    }
}

可以进一步合并指令，减少中间变量，从而减少对内存的读取，得到最终代码：

C++
for (long long i = 0; i < 4; i++)
{
    __m256d rd = _mm256_loadu_pd(d + n * 16 + i * 4);
    for (long long j = 0; j < 12; j++)
    {
        // AVX优化
        _mm256_storeu_pd(d + n * 16 + i * 4, _mm256_add_pd(rd, _mm256_mul_pd(_mm256_loadu_pd(b + n * 48 + j * 4), _mm256_set1_pd(*(a + n * 48 + i * 12 + j)))));
        }
    }

参考资料