Intel AVX を使用して SIMD 演算を試してみる

SIMD 演算と Intel AVX

 Single Instruction Multiple Data (SIMD) 演算とは1回の命令で複数のデータを同時に処理する演算です.近年の CPU には SIMD 演算を行うことができる SIMD 演算器が搭載されており,Intel 社の CPU ならば Streaming SIMD Extensions (SSE) を用いることで SIMD 演算を行うことが可能です.SSE は CPU に搭載されている 128bit レジスタを用いて演算を行うため単精度データならば4つ,倍精度データならば2つずつ演算を行うことができます.また,近年 SSE 後継の SIMD 拡張命令として Intel Advanced Vector eXtentions (AVX) が登場しました.AVX は第2世代 Intel Core i シリーズのプロセッサ (Sandy Bridge) から使用することが可能であり,演算幅が SSE の2倍の 256bit となっています.つまり,単精度データならば8つ,倍精度データならば4つずつ演算を行うことが可能です.

 今回は AVX を用いてベクトル  \bf{x} とベクトル  \bf{y} の加算演算を行います.ただし,

とします.また,今回に限り  n は"必ず"

  • 単精度演算を行う場合は8の倍数
  • 倍精度演算を行う場合は4の倍数

となるように設定することにします.
端数が発生する場合の計算方法は後日記載する予定です.

さて,上述の問題を C 言語で SIMD 演算を使わずにプログラミングすると以下のようになると思います (倍精度).

/*--------------------
  main.c
  --------------------*/

#include <stdio.h>
#include <stdlib.h>

void vec_add(const size_t n, double *z, const double *x, const double *y){
  for(size_t i=0; i<n; ++i)
    z[i] = x[i] + y[i];
}

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

  x = (double *)malloc(sizeof(double) * n);
  y = (double *)malloc(sizeof(double) * n);
  z = (double *)malloc(sizeof(double) * n);

  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]);

  free(x);
  free(y);
  free(z);

  return 0;
}

コンパイルは普通に

$ gcc main.c -std=c99 

としてもらえば大丈夫だと思います.

Intel AVX を使用する際のコンパイル方法

 C/C++ で AVX を使用する場合, をインクルードします.
をインクルードした場合,コンパイルオプションに -mavx を付ける必要があります.

$ gcc main.c -mavx

ただし,Mac OS X lion (10.7) でコンパイルする場合,gcc ではコンパイルすることができません.
現状の打開策として,gccアセンブラコードを生成し,生成したアセンブラコードを clang でコンパイルする方法が有力そうです.

$ gcc-mp-4.6 main.c -S -mavx
$ clang-mp-3.1 main.s -mavx

今回,Mac で用意した gcc や clang は Mac ports でインストールしたものを使用しています.

単精度演算の場合

 さっそくですが,単精度演算の場合のソースコードを見ていきましょう.

/*--------------------
  simd_single.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; 

  //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 = 1024;
  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;
}

AVX 命令を使用する場合,計算するデータのメモリアライメントを 256bit (32byte) の倍数に揃える必要があります.メモリアライメントを揃える方法はいろいろありますが,今回は に含まれる _mm_malloc を使用して動的に確保しています.静的に確保したい場合は

float __attribute__((aligned(32))) tmp[n];

等で確保することが可能です.ただし,__attribute__((aligned(32))) は gcc拡張機能であるため,gcc 以外のコンパイラを使用する場合は適宜適切な方法を使用してください.

_mm256_add_ps() を使用することで8つの単精度を同時に計算する SIMD 演算を行うことができますが,_mm256_add_ps() の引数と返り値は __mm256 と呼ばれる AVX 専用の型を使用する必要があります.今回は用意した単精度データを __mm256 にキャストしています.__m256 型にキャストすることでデータの格納形式が以下のようになります.

vx[0] = x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7];
vx[1] = x[8], x[9], x[10], x[11], x[12], x[13], x[14], x[15];
....
vx[n/8 - 1] = x[n-8], x[n-7], x[n-6], x[n-5], x[n-4], x[n-3], x[n-2], x[n-1];

vz, vy についても同様の格納方法となっています.あとは _mm256_add_ps() と for 文を用いて計算を行うだけです.また,減算,乗算,除算は _mm256_add_ps() を以下のように変更することで使用可能です.

  • 減算
    vz[i] = _mm256_sub_ps(vx[i], vy[i]); 
  • 乗算
    vz[i] = _mm256_mul_ps(vx[i], vy[i]); 
  • 除算
    vz[i] = _mm256_div_ps(vx[i], vy[i]); 

倍精度演算の場合

 今度は倍精度演算の場合のソースコードを見ていきましょう.使用する AVX 専用の型と加算命令が変化しただけで基本的な手法は同じです.

/*--------------------
  simd_double.c
  --------------------*/

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

void vec_add(const size_t n, double *z, const double *x, const double *y){
  static const size_t double_size = 4;
  const size_t end = n / double_size;

  __m256d *vz = (__m256d *)z;
  __m256d *vx = (__m256d *)x;
  __m256d *vy = (__m256d *)y;
  
  for(size_t i=0; i<end; ++i)
    vz[i] = _mm256_add_pd(vx[i], vy[i]);
}

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

  x = (double *)_mm_malloc(sizeof(double) * n, 32);
  y = (double *)_mm_malloc(sizeof(double) * n, 32);
  z = (double *)_mm_malloc(sizeof(double) * 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_free(x);
  _mm_free(y);
  _mm_free(z);

  return 0;
}

使用する AVX の型が __m256 から __m256d に変化しています.倍精度データを __m256d 型にキャストすることでデータの格納形式は以下のようになります.

vx[0] = x[0], x[1], x[2], x[3];
vx[1] = x[4], x[5], x[6], x[7];
....
vx[n/4 - 1] = x[n-4], x[n-3], x[n-2], x[n-1];

あとは _mm256_add_pd() を使用して計算していくだけで SIMD 演算を用いた並列処理が可能となります.

Intel AVX を使用して SIMD 演算を試してみる