AVX 命令を用いる SIMD 演算で総和を計算してみる
AVX 命令を用いる SIMD 演算で総和を計算する
今回は AVX 命令を用いる SIMD 演算で総和を計算してみます.SIMD 演算や AVX 命令については下記の URL を参照してください.
2012-03-03 - kawa0810の日記
さて,自分で総和計算のコードを書く場合は以下のようになる気がします.
/*-------------------- sum.cpp --------------------*/ #include <iostream> #include <cstdlib> #include <ctime> #include <numeric> #include <immintrin.h> template<class T> T vec_sum(const std::size_t n, const T* x){ T sum = 0; for(std::size_t i=0; i<n; ++i) sum += x[i]; return sum; } int main(void){ const std::size_t n = 10; float* x = static_cast<float*>( _mm_malloc(sizeof(float) * n, 32) ); double* y = static_cast<double*>( _mm_malloc(sizeof(double) * n, 32) ); srand( (unsigned)time(NULL) ); for(std::size_t i=0; i<n; ++i) x[i] = static_cast<float>( rand() ) / RAND_MAX; for(std::size_t i=0; i<n; ++i) y[i] = static_cast<double>( rand() ) / RAND_MAX; std::cout << "Single Precision\n"; std::cout << "-Accumulate: " << std::accumulate(x, x+n, 0.0) << std::endl; std::cout << "-SIMD: " << vec_sum(n, x) << std::endl; std::cout << "\nDouble Precision\n"; std::cout << "-Accumulate: " << std::accumulate(y, y+n, 0.0) << std::endl; std::cout << "-SIMD: " << vec_sum(n, y) << std::endl; _mm_free(x); _mm_free(y); return 0; }
上記のコードでは 0〜1 の乱数を生成し,生成した乱数の総和を計算しています.また,計算結果が正しいか確認するため,std::accumulate() を用いて確認しています.
SIMD 演算で総和を計算する際の注意事項
AVX を用いる場合の SIMD 演算は 128bit 単位か 256bit 単位で計算を行います.総和を計算する場合,__m128 型や __m256 型を用意しておき,用意したメモリにデータを足し込むことで総和を計算することができます.ただし,最後に __m128 型や __m256 型の中身をすべて足し合わせる作業が必要です.つまり,
sum x
[0] += [0] + [4] + .... + [n^4]
[1] += [1] + [5] + .... + [n-3]
[2] += [2] + [6] + .... + [n-2]
[3] += [3] + [7] + .... + [n-1]
と計算し,最後に sum[0] + sum[1] + sum[2] + sum[3] を計算します.
SIMD 演算を用いた総和の高速化
さっそくですが,ソースコードみていきましょう
/*-------------------- sum_simd.cpp --------------------*/ #include <iostream> #include <cstdlib> #include <ctime> #include <numeric> #include <immintrin.h> float vec_sum(const std::size_t n, float* x){ static const std::size_t para = 8; const std::size_t end = n / para; const std::size_t beg = end * para; __m256 vsum = _mm256_setzero_ps(); float* sum = reinterpret_cast<float*>(&vsum); __m256* vx = reinterpret_cast<__m256*>( x ); for(std::size_t i=0; i<end; ++i) vsum = _mm256_add_ps(vsum, vx[i]); for(std::size_t i=beg; i<n; ++i) sum[0] += x[i]; sum[0] += sum[1] + sum[2] + sum[3] + sum[4] + sum[5] + sum[6] + sum[7]; return sum[0]; } double vec_sum(const std::size_t n, double *x){ static const std::size_t para = 4; const std::size_t end = n / para; const std::size_t beg = end * para; __m256d vsum = _mm256_setzero_pd(); double* sum = reinterpret_cast<double*>(&vsum); __m256d* vx = reinterpret_cast<__m256d*>( x ); for(std::size_t i=0; i<end; ++i) vsum = _mm256_add_pd(vsum, vx[i]); for(std::size_t i=beg; i<n; ++i) sum[0] += x[i]; sum[0] += sum[1] + sum[2] + sum[3]; return sum[0]; } int main(void){ const std::size_t n = 10; float* x = static_cast<float*>( _mm_malloc(sizeof(float) * n, 32) ); double* y = static_cast<double*>( _mm_malloc(sizeof(double) * n, 32) ); srand( (unsigned)time(NULL) ); for(std::size_t i=0; i<n; ++i) x[i] = static_cast<float>( rand() ) / RAND_MAX; for(std::size_t i=0; i<n; ++i) y[i] = static_cast<double>( rand() ) / RAND_MAX; std::cout << "Single Precision\n"; std::cout << "-Accumulate: " << std::accumulate(x, x+n, 0.0) << std::endl; std::cout << "-SIMD: " << vec_sum(n, x) << std::endl; std::cout << "\nDouble Precision\n"; std::cout << "-Accumulate: " << std::accumulate(y, y+n, 0.0) << std::endl; std::cout << "-SIMD: " << vec_sum(n, y) << std::endl; _mm_free(x); _mm_free(y); return 0; }
単精度の場合は
__m256 vsum = _mm256_setzero_ps(); float* sum = reinterpret_cast<float*>(&vsum);
とすることで SIMD 演算を行う場合は vsum,逐次計算する場合は sum にアクセスすることができます.倍精度の場合もほぼ同様の方法で実装できます.
他にも何か方法がある気がしますが,今回はこんな感じで実装してみましたorz