C++avx256指令集加速实例

背景

avx256是x86cpu架构下实现SIMD（单指令多数据）的指令集。它能够利用cpu内部256bit的寄存器，同时对4位double或8位int类型的数操作，达到很好的加速效果。这里通过一个计算的实例来展示其威力：

linux系统下可以通过如下命令查看电脑是否支持avx256指令集：

cat /proc/cpuinfo | grep flags

image.png

sse4就表示该指令集。

//g++ avx_pi.cpp -mavx -O2 
#include <iostream>
#include <ctime>
#include <x86intrin.h>
using namespace std; 
//正常的逐个累加运算
double compute_pi_naive(size_t dt){
    double pi = 0.0;
    double delta = 1.0/dt;
    for (size_t i =0;i<dt;i++){
        double x = (double)i/dt;
        pi += delta /(1+x*x);
    }
    return pi*4.0;
}
//利用avx256指令集
double compute_pi_avx256(size_t dt){
    double pi = 0.0;
    double delta = 1.0/dt;
    __m256d ymm0,ymm1,ymm2,ymm3,ymm4;
    ymm0 = _mm256_set1_pd(1.0);
    ymm1 = _mm256_set1_pd(delta);
    ymm2 = _mm256_set_pd(delta*3,delta*2,delta,0.0);
    ymm4 = _mm256_setzero_pd();
    for (int i = 0;i<dt-4;i+=4){
        ymm3 = _mm256_set1_pd(i*delta);
        ymm3 = _mm256_add_pd(ymm3,ymm2);
        ymm3 = _mm256_mul_pd(ymm3,ymm3);
        ymm3 = _mm256_add_pd(ymm0,ymm3);
        ymm3 = _mm256_div_pd(ymm1,ymm3);
        ymm4 = _mm256_add_pd(ymm4,ymm3);        
    }
    double tmp[4] __attribute__((aligned(32)));
    _mm256_store_pd(tmp,ymm4);
    pi += tmp[0]+tmp[1]+tmp[2]+tmp[3];
    return pi*4.0;
}
int main(){
    clock_t start,end;
    size_t dt = 134217728;
    double result1,result2;
    //普通函数计时
    start = clock();
    result1 = compute_pi_naive(dt);
    end  = clock();
    cout<<"naive:\n"<< result1 <<endl<<end- start <<endl;
    //avx256计时
    start = clock();
    result2 = compute_pi_avx256(dt);
    end  = clock();
    cout<<"avx256:\n" <<result2 <<endl<<end- start <<endl;

    return 0;
}

image.png

可以看出，速度提高了两倍多。