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