SIMD命令を試してみました。

はじめに

Halideを試していたのですが、parallel以外のtile、unroll、vectorize パラメータを変更しても、全然高速化できない(どころかむしろ遅くなってしまう)という状況におちいってしまいました。メモリのキャッシュが重要でない単純なプログラムだったので、tileは効かないであろうことは想定内だったのですが、vectorize、unrollが全然速くならないという現象がぜんぜん理解できなかったので、SIMD 命令の動作を調べてみました。

SIMD命令

x86_64のCPUでは、SSEやAVX、ARMのCPUでは neon という名前で実装されている並列処理命令で、１度の命令で複数データをまとめて処理できてしまうという素敵な機能です。AVX-512では、1度の命令で、1バイトのデータなら64個も同時に処理できるようで、うまく使うえると大幅に高速化できそうな予感がします。

AVX 動作確認

以下のようなコードでAVXを使った場合と使わなかった場合でどのくらい違いが出るか調べてみました。

AVXを使わない場合

#include <iostream>
#include <chrono>

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

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

int main(int argc, char* argv[]) {

  const size_t n = 100 * 1024 * 1024;

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

  srand((unsigned int)time(NULL));
  for (size_t i=0; i<n; i++) x[i] = (uint8_t)rand() / RAND_MAX;
  for (size_t i=0; i<n; i++) y[i] = (uint8_t)rand() / RAND_MAX;
  for (size_t i=0; i<n; i++) z[i] = 0;

  auto t1 = std::chrono::system_clock::now();
  vec_add(n, z, x, y);
  auto t2 = std::chrono::system_clock::now();
  auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();
  std::cout << "sse\t" << duration << std::endl;

  free(x);
  free(y);
  free(z);
  return 0;
}

AVXを使った場合

#include <iostream>
#include <chrono>

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

const static size_t malign = 16;

void vec_add_sse(const size_t n, uint8_t *z, const uint8_t *x, const uint8_t *y) {
    for (size_t i=0; i<n; i+=malign) {
    /*
       _mm_stream_si128((__m128i*)(z + i), 
                        _mm_add_epi8(
                            _mm_load_si128((__m128i*)(x + i)), 
                            _mm_load_si128((__m128i*)(y + i)))
                        );
     */
       _mm_store_si128((__m128i*)(z + i), 
                       _mm_add_epi8(
                            _mm_loadu_si128((__m128i*)(x + i)), 
                            _mm_load_si128((__m128i*)(y + i)))
                       );
    }
}

int main(int argc, char* argv[]) {

  const size_t n = 100 * 1024 * 1024;

  uint8_t* x = (uint8_t *)_mm_malloc(n * sizeof(uint8_t), malign);
  uint8_t* y = (uint8_t *)_mm_malloc(n * sizeof(uint8_t), malign);
  uint8_t* z = (uint8_t *)_mm_malloc(n * sizeof(uint8_t), malign);

  srand((unsigned int)time(NULL));
  for (size_t i=0; i<n; i++) x[i] = (uint8_t)rand() / RAND_MAX;
  for (size_t i=0; i<n; i++) y[i] = (uint8_t)rand() / RAND_MAX;
  for (size_t i=0; i<n; i++) z[i] = 0;

  auto t1 = std::chrono::system_clock::now();
  vec_add_sse(n, z, x, y);
  auto t2 = std::chrono::system_clock::now();
  auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();
  std::cout << "sse\t" << duration << std::endl;

  // check
  for (size_t i=0; i<n; i++) {
    if (z[i] != (x[i] + y[i])) {
        std::cout << "error!" << std::endl;
        break;
    }
  }

  _mm_free(x);
  _mm_free(y);
  _mm_free(z);
  return 0;
}

immintrin.h をインクルードして、ロード、ストア、加算命令をAVXを使うよう変更するほか、メモリ確保関数をAVX用に変更しています。

速度比較

コンパイルオプション	AVXなし	AVXあり
g++ -std=c++11	240 ms	32 ms
g++ -std=c++11 -O2	45 ms	20 ms
g++ -std=c++11 -O3	21 ms	20 ms

CPU: Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz
コンパイラ: g++ (Ubuntu 5.4.0-6ubuntu1~16.04.10) 5.4.0 20160609

最適化オプションをつけない場合、AVX命令を使った場合とそうでない場合に大きな速度差がありますが、最適化オプションをつけた場合は、特に-O3オプションをつけた場合にほぼ差がなくなっていることが分かります。

そこで、アセンブラを出力して、コードを比較してみました。

AVXあり -O3

        movdqa  0(%r13,%rax), %xmm0
        paddb   (%r12,%rax), %xmm0
        movaps  %xmm0, 0(%rbp,%rax)

AVXなし -O3 13KB

        movdqu  (%r12,%r10), %xmm0
        paddb   0(%r13,%r10), %xmm0
        movups  %xmm0, (%r8,%r10)

AVXなしの方でも、並列命令が使われていることがわかりました。アライメントが揃っている時に使われる、movdqa/movapsの方が速いようですが、今回の実験ではその差は最終的な実行時間には、ほとんど影響していません。

まとめ

コンパイルオプションで最適化をつけない場合は、10倍近い差がありますが、最適化オプションをつけた場合には、コンパイラが勝手に並列命令を使ってくれていることが分かりました。

しかし結局、Halideでvectorize指定して遅くなる理由はわからず。vectorizeの時に、ちゃんとハードで実装しているSIMDのサイズと合ってないと行けないのかなと思う一方、根本的にvectorizeを勘違いしている気もします。

##　参考URL
組み込み関数（intrinsic）によるSIMD入門
https://www.slideshare.net/FukushimaNorishige/simd-10548373