SIMD 演算で if 文の並列化 - その1 -

C/C++ SIMD

　今月は SIMD 演算で条件分岐（if 文や三項演算子）の並列化を行います．SIMD で条件分岐を計算するには mask の概念の理解が必要になることや，SSE 世代と AVX 世代で条件分岐を計算する関数の使用方法が変わっているので今回は簡単にご紹介します．

SSE と AVX の比較の関数

　SSE や AVX では条件分岐文の変わりに cmp 関数を用いて計算を行います．ただし，SSE 世代と AVX 世代では cmp 関数の使用方法が大きく変更されました．SSE 世代で a == b を計算したければ _mm_cmpeq_ps() 関数，a < b を計算したければ _mm_cmplt_ps() 関数など演算子ごとに専用の関数が用意されていましたが，AVX 世代では以下の関数で全ての条件演算子の計算を行います．

__m256 _mm256_cmp_ps(__m256 a, __m256 b, const int imm);//単精度 (float)
__m256d _mm256_cmp_pd(__m256d a, __m256d b, const int imm);//倍精度 (double)
//a: 左辺の値
//b: 右辺の値
//imm: 計算させたい演算子の種類を指定（詳細は以下に記載）
//返り値は mask と呼ばれるものが戻ってきます

　条件演算子と SSE 世代の関数と AVX の imm に指定する値の基本的な組み合わせの一覧を以下に記載します．
■ == 関係
・a==b，_mm_cmpeq_ps() など ，_CMP_EQ_OQ など
・a!=b ，_mm_cmpneq_ps() など ，_CMP_NEQ_UQ など

■ < 関係
・a 関係
・a>b，_mm_cmpgt_ps() など ，_CMP_GT_OS など
・a>=b，_mm_cmpge_ps() など ，_CMP_GE_OS など
・!(a>b) ，_mm_cmpngt_ps() など ，_CMP_NGT_US など
・!(a>=b) ，_mm_cmpnge_ps() など ，_CMP_NGE_OS など

■ 計算することができる数値かどうか判定 *1
・(a ord b) ，_mm_cmpord_ps() など ，_CMP_ORD_Q など
・(a unord? b)，_mm_cmpunord_ps() など ，_CMP_UNORD_Q など
数値なのか，nan などなのかを判定してくる関数です（詳細不明）．
他のサイトでは「順序付けされないか比較」と書いているサイトもありますが，おそらく誤訳と思われます＾＾；
試しに実行したところ以下のような結果が得られました．

  • inf unord? inf => 0x00000000
  • inf unord? nan => 0x00000000
  • nan unord? nan => 0xffffffff

mask の考え方

　cmp 関数の返り値は mask と呼ばれるもので得られます．_mm256_cmp_ps 関数の場合，条件が成立すれば (TURE ならば) 0xFFFFFFFF，成立しなければ (ELSE ならば) 0x00000000 が得られます．SIMD 演算で条件文を並列処理する場合，得られた mask を別の SIMD 処理系の関数に渡すことによって並列化が可能となります．

　mask を元に SIMD 演算を行う関数については Intel AVX の Instrinsics Guide for Intel Advanced Vector Extensions 2 をご覧頂くのがよいと思います *2．次節のサンプルコードでは _mm256_blendv_ps() と _mm256_blendv_pd() を使用して代入を行っています．

サンプルコード

　配列 \bf{x} 内の要素で閾値以下の値をすべて閾値と同じ値にするサンプルコードを以下に記載します．

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

std::size_t const mem_align = 32;

void check_value(std::size_t const n, float* x, float const val){
  std::size_t const para = mem_align / sizeof(float);
  std::size_t const end = (n / para) * para;

  __m256 const vval = _mm256_set1_ps(val);
  for(std::size_t i=0; i<end; i+=para){
    __m256 vx = _mm256_load_ps(x+i);
    __m256 mask = _mm256_cmp_ps(vx, vval, _CMP_LT_OS);//if ( x < val ) 
    _mm256_stream_ps(x+i, _mm256_blendv_ps(vx, vval, mask) );//mask の値を元に代入
  }

  //端数処理
  for(std::size_t i=end; i<n; ++i)
    x[i] = (x[i] < val) ? val : x[i];
}

void check_value(std::size_t const n, double* x, double const val){
  std::size_t const para = mem_align / sizeof(double);
  std::size_t const end = (n / para) * para;

  __m256d const vval = _mm256_set1_pd(val);
  for(std::size_t i=0; i<end; i+=para){
    __m256d vx = _mm256_load_pd(x+i);
    __m256d mask = _mm256_cmp_pd(vx, vval, _CMP_LT_OS);//if ( x < val ) 
    _mm256_stream_pd(x+i, _mm256_blendv_pd(vx, vval, mask) );//mask の値を元に代入
  }

  //端数処理
  for(std::size_t i=end; i<n; ++i)
    x[i] = (x[i] < val) ? val : x[i];
}

int main(void){
  std::size_t const n = 35;
  float* x = static_cast<float*>( _mm_malloc(sizeof(float) * n, mem_align) );
  double* y = static_cast<double*>( _mm_malloc(sizeof(double) * n, mem_align) );

  srand( static_cast<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;

  check_value(n, x, 0.5);
  std::cout << "Single Precision\n";
  for(std::size_t i=0; i<n; ++i)
    std::cout << x[i] << std::endl;

  check_value(n, y, 0.5);
  std::cout << "\nDouble Precision\n";
  for(std::size_t i=0; i<n; ++i)
    std::cout << y[i] << std::endl;

  _mm_free(x);
  _mm_free(y);

  return 0;
}
SIMD 演算で if 文の並列化

*1:参考サイト：advanced optimization SSE

*2: Intel Instruction Set Architecture Extensions のページ下部