Intel AVX を使用して SIMD 演算を試してみる - その4 -

2012-03-05 - kawa0810の日記の続きです．
今回はその2 (2012-03-04 - kawa0810の日記) で紹介した AVX 命令を試します．

計算する問題

　今回は SIMD 演算を用いてベクトルのスカラ倍 $\bf{x}^{'} = \alpha \bf{x}$ を計算する方法を考えます．ただし，

$\bf{x} = [0, 1, \cdots, n-1]^{ \rm{t} }$
$n$ はベクトルの要素数
$\alpha$ = 2.0

とします．また，端数は逐次計算で処理を行います．

普通にプログラミングすると以下のようになると思われます．ただし，今回はプログラミング言語を C++ にし，テンプレートを使ったり関数をオーバーロードさせたりして単精度と倍精度のソースが1つに収まるようにしています．

/*--------------------
  scal.cpp
  --------------------*/

#include <iostream>
#include <immintrin.h>

template<class T> 
void scal(const std::size_t n, const T alpha, T* x){
  for(std::size_t i=0; i<n; ++i)
    x[i] = alpha * x[i];
}

int main(void){
  const std::size_t n = 12;
  const double alpha = 2.0;
  double *x = static_cast<double *>( _mm_malloc(sizeof(double) * n, 32) );

  for(std::size_t i=0; i<n; ++i) 
    x[i] = i;
  
  scal(n, alpha, x);

  for(std::size_t i=0; i<n; ++i) 
    std::cout << x[i] << "\n";
  
  _mm_free(x);

  return 0;
}

C++ の場合でもコンパイル方法はほぼ同様です．Mac でコンパイルする場合を以下に記載しておきます．

$ g++-mp-4.6 -S main.cpp -std=c++0x -mavx
$ clang++-mp-3.1 main.s -mavx

SIMD 演算を行うにあたり気をつけること

　AVX 命令は必ず 128bit (16byte) 単位もしくは 256bit (32byte) 単位で計算を行います．しかしながら， $\alpha$ はスカラであり，単精度ならば 4byte 倍精度ならば 8byte であるため計算を行う際に少し工夫が必要です．今回は _mm256_set_ps() もしくは _mm256_set_pd を使用して __m256 型に複数 $\alpha$ を格納して計算を行うことにします．つまり，

単精度で計算を行う場合

__m256 valpha = _mm256_set_ps(alpha, alpha, alpha, alpha, alpha, alpha, alpha, alpha);
//valpha[0] == alpha
//valpha[1] == alpha
//valpha[2] == alpha
//valpha[3] == alpha
//valpha[4] == alpha
//valpha[5] == alpha
//valpha[6] == alpha
//valpha[7] == alpha

倍精度で計算を行う場合

__m256d valpha = _mm256_set_pd(alpha, alpha, alpha, alpha);
//valpha[0] == alpha
//valpha[1] == alpha
//valpha[2] == alpha
//valpha[3] == alpha

となります．

SIMD 演算を用いて高速化

　上述の点に気をつけてコードを書くと以下のようになるかと思います．

/*--------------------
  scal_simd.cpp
  --------------------*/

#include <iostream>
#include <immintrin.h>

void scal(const std::size_t n, const float alpha, float* x){
  static const std::size_t para = 8;
  __m256 *vx = reinterpret_cast<__m256*>( x );
  const __m256 valpha = _mm256_set_ps(alpha, alpha, alpha, alpha, alpha, alpha, alpha, alpha);
  
  const std::size_t end = n / para;
  const std::size_t beg = end * para;

  for(std::size_t i=0; i<end; ++i)
    vx[i] = _mm256_mul_ps(valpha, vx[i]);
  for(std::size_t i=beg; i<n; ++i)
    x[i] = alpha * x[i];
}

void scal(const std::size_t n, const double alpha, double* x){
  static const std::size_t para = 4;
  __m256d *vx = reinterpret_cast<__m256d*>( x );
  const __m256d valpha = _mm256_set_pd(alpha, alpha, alpha, alpha);
  
  const std::size_t end = n / para;
  const std::size_t beg = end * para;

  for(std::size_t i=0; i<end; ++i)
    vx[i] = _mm256_mul_pd(valpha, vx[i]);
  for(std::size_t i=beg; i<n; ++i)
    x[i] = alpha * x[i];
}

int main(void){
  const std::size_t n = 12;
  const double alpha = 2.0;
  double *x = static_cast<double *>( _mm_malloc(sizeof(double) * n, 32) );

  for(std::size_t i=0; i<n; ++i) 
    x[i] = i;
  
  scal(n, alpha, x);

  for(std::size_t i=0; i<n; ++i) 
    std::cout << x[i] << "\n";
  
  _mm_free(x);

  return 0;
}