読者です 読者をやめる 読者になる 読者になる

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