用 AVX 编写 SIMD 程序

Accelerate the Program with SIMD using AVX

Posted by lzzmm on September 9, 2021
About 10 minutes to read

DCS242 - Parallel and Distributed Computing 2021 Fall

SIMD

SIMD(Single Instruction Multiple Data,单指令多数据流),是一种实现空间上的并行性的技术。这种技术使用一个控制器控制多个处理单元,同时对一组数据中的每一个数据执行相同的操作。在 SIMD 指令执行期间,任意时刻都只有一个进程在运行,即 SIMD 没有并发性,仅仅只是同时进行计算。 在 Intel 的 x86 微架构处理器中,SIMD 指令集有 MMX、SSE、SSE2、SSE3、SSSE3、SSE4.1、SSE4.2、AVX、AVX2、AVX512。

SSE

SSE(Streaming SIMD Extensions)是英特尔在 AMD 的 3D Now! 发布一年之后,在其计算机芯片 Pentium III 中引入的指令集,是继 MMX 的扩展指令集。SSE 指令集提供了70条新指令。AMD 后来在 Athlon XP 中加入了对这个新指令集的支持。

SSE5 是 AMD 为了打破 Intel 垄断在处理器指令集的独霸地位所提出的,目前AMD已放弃下一代 Bulldozer 核心内置 SSE5 指令集,改内置 Intel 授权 SSE4 系列指令集。

AVX

Intel AVX(Advanced Vector Extensions,高级矢量扩展)指令集是 Sandy Bridge 和 Larrabee 架构下的新指令集。AVX 沿用了的 MMX/SSE 指令集并进行扩展和加强,不过和 MMX/SSE 的不同点在于增强的 AVX 指令从指令的格式上就发生了很大的变化,形成了一套新一代的完整 SIMD 指令集规范。AVX 将 SSE 的 XMM 128bit 寄存器升级成了 YMM 256bit 寄存器,同时浮点运算命令扩展至 256 位,运算效率提升了一倍。另外,AVX 还添加了三操作数指令,以减少在编码时先复制再运算的动作,实现了单指令多数据流计算性能增强。

AVX2 将大多数整数运算命令扩展至 256 位,同时支持 FMA(Fused Multiply-Accumulate,融合乘法累加)运算,可以在提高运算效率的同时减少运算时的精度损失。

AVX512 将 AVX 指令进一步扩展至 512 位。

使用 AVX 编写 SIMD 程序

数据类型

数据类型 描述
__m128 包含4个float类型数字的向量
__m128d 包含2个double类型数字的向量
__m128i 包含若干个整型数字的向量
__m256 包含8个float类型数字的向量
__m256d 包含4个double类型数字的向量
__m256i 包含若干个整型数字的向量

由2个下划线开头,接一个m,然后是vector的位长度。

__m128__m256 每个数由 4bytefloat 构成;

d后缀表示双精度浮点数,例如__m128d__m256d,每个数由 8bytedouble 构成。

__m128i__m256i是由整型构成的向量,charshortintlong均属于整型(以及 unsigned 以上类型),所以例如 __m256i 就可以由32个 char,或者16个 short ,或者8个 int ,又或者4个 long 构成,这些整型可以是有符号类型也可以是无符号类型。

函数

0
_mm<bit_width>_<name>_<data_type>

函数名称由 _mm 开头。

<bit_width> 表明了向量的位长度,对于128位的向量,这个参数为空,对于256位的向量,这个参数为256。

<name> 描述了内联函数的算术操作。

<data_type> 标识函数主参数的数据类型,参数含义如下:

  1. ps:里面都是float,把32bits当成一个数看
  2. pd:里面都是double,把64bits当成一个数看
  3. epi8 / epi16 / epi32 / epi64:向量里每个数都是整型,一个整型8bit/16bit/32bit/64bit
  4. epu8 / epu16 / epu32 / epu64:向量里每个数都是无符号整型(unsigned),一个整型8bit/16bit/32bit/64bit
  5. m128 / m128i / m128d / m256 / m256i / m256d:输入值与返回类型不同时会出现 ,例如__m256i_mm256_setr_m128i(__m128ilo,__m128ihi),输入两个__m128i向量 ,把他们拼在一起,变成一个__m256i返回 。另外这种结尾只见于load
  6. si128 / si256:不care向量里到底都是些啥类型,反正128bit/256bit,例如:__m256i _mm_broadcastsi128_si256 (__m128i a)

AVX指令的函数可以通过Intrinsics Guide查询

实现第一个使用 AVX 的程序

通过学习决定先实现两个向量的加法,具体如下:

0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
// hw1.c

// Copyright (c) 2021 CHEN Yuhan
// Date: 2021-09-17
//
// Test using avx
//
#include <immintrin.h>
#include <stdio.h>

int main(){
    __m256 float_vec_0 = _mm256_set1_ps(102.4);
    __m256 float_vec_1 = _mm256_set1_ps(1.024);

    __m256 float_result = _mm256_add_ps(float_vec_0, float_vec_1);
    float *f_arr = (float *)&float_result;

    printf("\nfloat_result:\n%f %f %f %f\n%f %f %f %f", f_arr[0],f_arr[1],f_arr[2],f_arr[3],f_arr[4],f_arr[5],f_arr[6],f_arr[7]);
    return 0;
}

其中加上头文件 #include <immintrin.h> 才可以使用上述数据类型和函数。

我定义了两个256bit长的float向量并分别初始化为 102.4 和 1.024,定义了一个 256bit 长的 float 向量来存储 float 向量加法的结果,然后打印出来。

编译运行:

0
1
2
3
4
5
(base) PS C:\Users\cortex\Documents\code\cpp\pdc> gcc ./hw1.c -mavx -mavx2 
(base) PS C:\Users\cortex\Documents\code\cpp\pdc> .\a.exe

float_result:
103.424004 103.424004 103.424004 103.424004       
103.424004 103.424004 103.424004 103.424004  

测试使用AVX后的加速效果

我使用随机数填充了两个有着一千万元素的float数组,先用串行循环算出两个数组的和,算出时间,再用 AVX 指令计算。在使用 AVX 的时候,用到了

0
_mm256_loadu_ps(float const * mem_addr)

这个未对齐读取函数来读取数组中的连续8个浮点,然后使用

0
__m256 _mm256_add_ps (__m256 a, __m256 b)

这个函数运算出结果,并使用

0
void _mm256_storeu_ps (float * mem_addr, __m256 a)

这个未对齐写回函数将结果写入结果数组中。

由于avx指令每次操作8个元素,在最后还剩下 $NUM \% 8$ 个元素,直接使用串行的方法将其循环相加。

最后输出时间消耗和加速比。

注意,在函数中声明一千万长度的局部数组会发生运行时错误,解决方法是采用全局数组或者动态内存分配,我使用了全局数组。

0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
// hw1_avx.c

// Copyright (c) 2021 CHEN Yuhan
// Date: 2021-09-17
//
// Computing vector sum using avx and test the speedup
//
#include <immintrin.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#define NUM 100000000

float vec_1[NUM], vec_2[NUM], vec_res[NUM];

int main() {
    long long i;
    clock_t start_time, end_time;

    for (i = 0; i < NUM; i++) {
        vec_1[i] = (float)rand() / 1024;
        vec_2[i] = (float)rand() / 1024;
    }

    // serial
    start_time = clock();
    for (i = 0; i < NUM; i++) {
        vec_res[i] = vec_1[i] + vec_2[i];
    }
    end_time = clock();
    double s_time = (double)(end_time - start_time) / CLOCKS_PER_SEC;
    printf("\nserial time cost: %lfs\n", s_time);
    printf("\n%f %f %f %f\n", vec_1[0], vec_1[1], vec_1[2], vec_1[3]);
    printf("\n%f %f %f %f\n", vec_2[0], vec_2[1], vec_2[2], vec_2[3]);
    printf("\n%f %f %f %f\n", vec_res[0], vec_res[1], vec_res[2], vec_res[3]);

    // avx parallel
    start_time = clock();
    __m256 avx_vec_1, avx_vec_2;
    for (i = 0; i < NUM; i += 8) {
        int *x = vec_1 + i, *y = vec_2 + i, *z = vec_res + i;
        avx_vec_1 = _mm256_loadu_ps((const __m256i *)x);
        avx_vec_2 = _mm256_loadu_ps((const __m256i *)y);
        avx_vec_1 = _mm256_add_ps(avx_vec_1, avx_vec_2);
        _mm256_storeu_ps((__m256i *)z, avx_vec_1);
    }
    for (; i < NUM; i++) { // NUM % 8
        vec_res[i] = vec_1[i] + vec_2[i];
    }
    end_time = clock();
    double p_time = (double)(end_time - start_time) / CLOCKS_PER_SEC;
    printf("\nparallel time cost: %lfs\n", p_time);
    printf("\n%f %f %f %f\n", vec_1[0], vec_1[1], vec_1[2], vec_1[3]);
    printf("\n%f %f %f %f\n", vec_2[0], vec_2[1], vec_2[2], vec_2[3]);
    printf("\n%f %f %f %f\n", vec_res[0], vec_res[1], vec_res[2], vec_res[3]);

    printf("\nSpeedup: %f\n", s_time / p_time);

    return 0;
}

结果如下:

0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
(base) PS C:\Users\cortex\Documents\code\cpp\pdc> gcc ./hw1_avx.c -mavx -mavx2
(base) PS C:\Users\cortex\Documents\code\cpp\pdc> .\a.exe

serial time cost: 0.791000s

0.040039 6.185547 18.719727 11.208984

18.034180 25.878906 15.355469 28.669922

18.074219 32.064453 34.075195 39.878906

parallel time cost: 0.274000s

0.040039 6.185547 18.719727 11.208984

18.034180 25.878906 15.355469 28.669922

18.074219 32.064453 34.075195 39.878906

Speedup: 2.886861

可以看到加速比达到了接近2.89,是一个比较正常的值。

使用 AVX 加速 GEMM

结语

以上是对 AVX 和其他 SIMD 指令集的初探,了解不深,欢迎批评。

References

Intel® Intrinsics Guide

https://github.com/chen0031/AVX-AVX2-Example-Code

Creative Commons License本作品采用知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议进行许可。
This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.