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

2012-03-04 - kawa0810の日記の続きです．
今回は SIMD 演算を行う際に問題となる端数の計算方法です．

SIMD 演算を行う際に発生する端数の処理方法

　AVX 命令を用いる SIMD 演算は 128bit 単位もしくは 256bit 単位で計算を行うため任意の長さのデータを計算する際は端数がでることが多いと思います．端数がでるデータを計算する際はしっかりと端数の処理をしなければ正しい計算結果が得られないだけでなく，予期しないメモリにデータを書き込んでしまう問題も起こりえます．特に High Performance Computing (HPC) の分野等では任意の長さのデータを計算することが多く，プログラマが上述の問題に対応するしかありません．

　端数の処理方法には以下の2通りの方法が考えられます．

計算結果に影響を与えないデータを混入させ SIMD 演算が行いやすい任意の長さにデータを拡張する (データパディング方法)
SIMD 演算で計算できるところまで計算を行い，端数は逐次処理する (端数を逐次処理方法)

上記の2通りの方法にはそれぞれメリット・デメリットが存在するため一概にどちらが良いとは言えません．今回は「Intel AVX を使用して SIMD 演算を試してみる - その1 -」(2012-03-03 - kawa0810の日記) で取り扱ったベクトル同士の加算を例に比較を行っていきます．

データパディング方法

　以下に単精度の場合のソースコードを記載します．倍精度の場合は single_size = 8 を single_size = 4 に変更し，__m256 を __m256d に変更すれば計算可能だと思います．

/*--------------------
  simd_single_pading.c
  --------------------*/

#include <stdio.h>
#include <stdlib.h>
#include <immintrin.h>//AVX: -mavx

const size_t single_size = 8; //単精度は8つずつ計算

void vec_add(const size_t n, float *z, const float *x, const float *y){
  const size_t end = n / single_size; 

  //AVX 専用の型にデータをロードする
  __m256 *vz = (__m256 *)z;
  __m256 *vx = (__m256 *)x;
  __m256 *vy = (__m256 *)y;
  
  for(size_t i=0; i<end; ++i)
    vz[i] = _mm256_add_ps(vx[i], vy[i]); //AVX を用いる SIMD 演算
}

int main(void){
  const size_t n = 12;
  size_t space;
  float *x, *y, *z;

  if(n % single_size != 0)
    space = (n / single_size) * single_size + single_size;
  else
    space = n;
  printf("space = %d\n", space);

  //メモリのアライメントを 32byte 境界に揃える
  x = (float *)_mm_malloc(sizeof(float) * space, 32);
  y = (float *)_mm_malloc(sizeof(float) * space, 32);
  z = (float *)_mm_malloc(sizeof(float) * space, 32);

  for(size_t i=0; i<n; ++i) x[i] = i;
  for(size_t i=n; i<space; ++i) x[i] = 0.0;//パディング

  for(size_t i=0; i<n; ++i) y[i] = i+1;
  for(size_t i=n; i<space; ++i) y[i] = 0.0;//パディング

  for(size_t i=0; i<space; ++i) z[i] = 0.0;

  vec_add(space, z, x, y);
  
  for(size_t i=0; i<n; ++i) printf("%g\n", z[i]);

  //_mm_malloc で確保した領域は _mm_free で解放する
  _mm_free(x);
  _mm_free(y);
  _mm_free(z);

  return 0;
}

今回は space という変数を宣言し，ベクトルの長さ $n$ が8（単精度は8個ずつ計算）で割り切れない場合に確保するデータ長を変更するようにしています．変更するデータ長は8の倍数になれば良いので $(n \div 8) \times 8 + 8$ とすることでちょうど良い8の倍数を計算することが可能です．今回はベクトル同士の加算であるため混入させるデータはなんでも良いのですが，上記のソースコードでは0を混入させています．あとは，vec_add に渡すベクトルの長さ $n$ の代わりに space を渡すことで vec_add は変更することなく SIMD 演算を行うことができます．

データパディング方法のメリットとデメリット

メリット
- SIMD 演算のソースコード部分は変更する必要がないため理解がしやすい
- SIMD 演算だけですべて計算することができるため，端数を逐次演算方法よりも高速に計算できる可能性が高い
デメリット
- 入力データの増加や複雑化に伴い，リソース管理が困難になりやすい → バグが発生しやすくなる
- データの大規模化に伴い，無駄なメモリ空間が発生する

端数を逐次処理方法

　以下に端数を逐次処理するタイプの単精度の場合のソースコードを記載します．倍精度の場合は single_size = 8 を single_size = 4 に変更し，__m256 を __m256d に変更すれば計算可能だと思います．

/*--------------------
  simd_single_scalar.c
  --------------------*/

#include <stdio.h>
#include <stdlib.h>
#include <immintrin.h>//AVX: -mavx

void vec_add(const size_t n, float *z, const float *x, const float *y){
  static const size_t single_size = 8; //単精度は8つずつ計算
  const size_t end = n / single_size; 
  const size_t beg = end * single_size;

  //AVX 専用の型にデータをロードする
  __m256 *vz = (__m256 *)z;
  __m256 *vx = (__m256 *)x;
  __m256 *vy = (__m256 *)y;
  
  for(size_t i=0; i<end; ++i)
    vz[i] = _mm256_add_ps(vx[i], vy[i]); //AVX を用いる SIMD 演算
  for(size_t i=beg; i<n; ++i)
    z[i] = x[i] + y[i];//端数を逐次演算
}

int main(void){
  const size_t n = 12;
  float *x, *y, *z;

  //メモリのアライメントを 32byte 境界に揃える
  x = (float *)_mm_malloc(sizeof(float) * n, 32);
  y = (float *)_mm_malloc(sizeof(float) * n, 32);
  z = (float *)_mm_malloc(sizeof(float) * n, 32);

  for(size_t i=0; i<n; ++i) x[i] = i;
  for(size_t i=0; i<n; ++i) y[i] = i+1;
  for(size_t i=0; i<n; ++i) z[i] = 0.0;

  vec_add(n, z, x, y);
  
  for(size_t i=0; i<n; ++i) printf("%g\n", z[i]);

  //_mm_malloc で確保した領域は _mm_free で解放する
  _mm_free(x);
  _mm_free(y);
  _mm_free(z);

  return 0;
}

データパディング方法と異なり入力データには一切の変更を加えず，vec_add 内だけで全て計算しています．SIMD 演算で計算することができる長さが end = $n / 8$ で計算することができ，逐次処理を開始する位置は beg = end $\times 8$ で求めることができます．また，当然ながら逐次処理の終了位置は $n$ ですので今回の例題の場合，大きな変更を加えずに計算することが可能となります．

端数は逐次処理方法のメリットとデメリット

メリット
- 入力データには一切変更を加えないため，入力データが大規模な場合や複雑な場合でも対応しやすい．
- リソースの管理がデータパディング方法よりも容易になる．
デメリット
- 逐次処理が発生する分，計算に時間を要する可能性がある．
- SIMD 演算部分のソースコードが複雑になりやすい．

個人的なまとめ

　リソース管理を考えると「端数は逐次処理方法」がよく，少しでも高い演算能力が欲しい場合は「データパディング方法」が良いのかなと思います．