Intel ISPC & SPH on CPU

コミケに参加した方々、お疲れ様でした。
サイトの方では告知し忘れていましたが (…) 今回は落選でした。次回の冬コミは既に申し込み済みで、受かれば参加します。牛歩の歩みですが、来年 5 月までには何らかの形で完成した作品を出すつもりです。

前回の冬コミ版は SPH (流体シム) を CUDA で実装しており、ゆえに Radeon では動かないし、GeForce でも 500 系未満だと予想外に遅く、スペックの差をどう解決するかが頭が痛い問題だったんですが、現在 Intel ISPC による CPU 実装で行こうという方針に定まりつつあります。

この Intel ISPC なるものが面白いので紹介。
http://ispc.github.com/
GPGPU 的な並列プログラミングの思想を CPU の SIMD で実現するプログラム言語です。SIMD の各要素 (lane) を実行単位と見立て、1 つのプログラムを lane 数分 (SSE なら 4、AVX なら 8) 並列に実行します。
例えば、float の配列の合計を計算する ISPC のコードは以下のようになります。

export uniform float ArrayTotal(uniform const float a[], uniform const int n)
{
    float r = 0.0f;
    for(uniform int i=0; i<n; i+=programCount) {
        r += a[i + programIndex];
    }
    return reduce_add(r);
}

キーワード uniform は全 SIMD lane で共通の値だということを明示するものです。uniform を付けない場合、SIMD lane それぞれで別の値を持つ varing 変数となります。programCount は SIMD lane の数 (SSE なら 4、AVX なら 8)、programIndex はその時点で実行している SIMD lane のインデックスです。reduce_add() は指定の varing 変数の各 SIMD lane 要素を合計した値を返すもので、つまり返る値は uniform になります。
上の ISPC コードは以下の C++ コードと等価になります。

float ArrayTotal_sse2(const float a[], const int n)
{
    __m128 r = _mm_set1_ps(0.0f);
    for(int i=0; i<n; i+=4) {
        __m128 t = _mm_load_ps(&a[i]);
        r = _mm_add_ps(r, t);
    }

    // 以下 r の各要素の合計を算出
    __m128 r2;
    r2 = _mm_movehl_ps(r, r);   // r2: z,w,z,w
    r  = _mm_add_ps(r, r2);     // r: xz,yw,zz,ww
    r2 = _mm_shuffle_ps(r, r, _MM_SHUFFLE(2, 2, 2, 2));  // r2: yw, yw, yw, yw
    r =  _mm_add_ss(r, r2);     // r: xzyw, ...
    return (float&)r;
}

ただ、これは target が SSE2 の場合です。ISPC は SSE と AVX のほぼ全てのバージョンに対応しており、--target=avx とやるだけで AVX 用のコードを生成してくれます。--target=sse2,sse4,avx などとしておけば、実行時に最適なバージョンを自動的に選択して実行してくれます。これは intrinsic や アセンブラでがんばってた人間からすると感涙ものの便利さだと思います！

また、上のコードは配列の要素数が SIMD lane 数の倍数ではない場合結果がおかしくなりますが、ISPC はそこらへんも適切に処理してくれる foreach 文が用意されています。若干のオーバーヘッドを伴いますが、大抵はこちらを使ったほうがいいでしょう。最初の ISPC のコードを foreach で書くと以下になります。シンプル。

export uniform float ArrayTotal(uniform const float a[], uniform const int n)
{
    float r = 0.0f;
    foreach(i=0 ... n) {
        r += a[i];
    }
    return reduce_add(r);
}

SIMD の並列性を最大限に活かすには、データを SoA に並べ替えておく必要があります。
(SoA: Structure of Array。xxxx,yyyy,zzzz みたいな要素別配列のデータの並び。xyz,xyz,xyz,xyz は Array of Structure)
ISPC の標準ライブラリには xyz や xyzw の繰り返しを SoA に並べ替える関数が用意されていますが、大抵の場合もっと複雑なデータ構造を並べ替えたいはずで、C++ 側で↓のようなルーチンで並べ替えることになると思います。
https://github.com/i-saint/scribble/blob/master/SOATranspose.h

ISPC は SoA のデータ構造を言語レベルでサポートしており、例えば struct Hoge{x,y,z} みたいな構造があったとして、soa<8> Hoge hoge と書くとこの hoge は HogeSOA8{x[8],y[8],z[8]} のような構造として扱われるようになります。
具体例。以下は SoA なパーティクル (球) を総当りで当たり判定を取り、めり込んだ分逆方向に加速させるコードです。未テストにつき間違ってる可能性アリ。

struct Particle
{
    float x,y,z, radius; // float<3> とやると soa にできなくなるので x,y,z 分けないといけない
    float vx,vy,vz;
};

static inline float length3(float<3> v)
{
    return sqrt(v.x*v.x + v.y*v.y + v.z*v.z);
}

export void ParticleInteraction(soa<8> Particle particles[], uniform int num_particles)
{
    for(uniform int ri=0; ri<num_particles; ++ri) {
        uniform float<3> rpos = {particles[ri].x, particles[ri].y, particles[ri].z};
        float<3> accel = {0.0f, 0.0f, 0.0f};
        foreach(si=0 ... num_particles) {
            float<3> spos = {particles[si].x, particles[si].y, particles[si].z};
            float<3> diff = spos - rpos;
            float distance = length3(diff);
            if(distance > 0.0f) { // 距離が 0.0f == 自分自身との衝突なので無視
                float radius = particles[ri].radius + particles[si].radius;
                float<3> dir = diff / distance;
                accel += dir * min(distance-radius, 0.0f); // 実際に使用する際は反発係数とかを掛ける
            }
        }
        particles[ri].vx += reduce_add(accel.x);
        particles[ri].vy += reduce_add(accel.y);
        particles[ri].vz += reduce_add(accel.z);
    }
}

(たぶんこのままだと dir = diff / distance が実行時間の半分以上を占めるので、実践では半径を一定にして逆数を事前に持たせて近似するなどのごまかしが必要になるでしょう。そのくらい割り算は遅いです)
SSE2, x86-64 の target でコンパイルすると、foreach の中は以下のようなコードになります。

	movups	(%rcx,%rbx), %xmm0
	movups	32(%rcx,%rbx), %xmm4
	movups	64(%rcx,%rbx), %xmm5
	subps	%xmm15, %xmm4
	movaps	%xmm4, %xmm6
	mulps	%xmm6, %xmm6
	subps	%xmm11, %xmm0
	movaps	%xmm0, %xmm3
	mulps	%xmm3, %xmm3
	addps	%xmm6, %xmm3
	subps	%xmm14, %xmm5
	movaps	%xmm5, %xmm6
	mulps	%xmm6, %xmm6
	addps	%xmm3, %xmm6
	sqrtps	%xmm6, %xmm3
	movaps	%xmm3, %xmm8
	cmpnleps	%xmm13, %xmm8
	movmskps	%xmm8, %eax
	testl	%eax, %eax
	je	.LBB0_25
# BB#24:                                # %safe_if_run_true.us
                                        #   in Loop: Header=BB0_23 Depth=2
	movss	(%r14), %xmm6
	movups	96(%rbx,%rcx), %xmm7
	pshufd	$0, %xmm6, %xmm6        # xmm6 = xmm6[0,0,0,0]
	divps	%xmm3, %xmm5
	divps	%xmm3, %xmm4
	addps	%xmm7, %xmm6
	divps	%xmm3, %xmm0
	subps	%xmm6, %xmm3
	minps	%xmm13, %xmm3
	mulps	%xmm3, %xmm4
	mulps	%xmm3, %xmm5
	movaps	%xmm8, %xmm10
	andnps	%xmm1, %xmm10
	movaps	%xmm8, %xmm9
	andnps	%xmm12, %xmm9
	addps	%xmm1, %xmm5
	andps	%xmm8, %xmm5
	movaps	%xmm8, %xmm1
	andnps	%xmm2, %xmm1
	orps	%xmm10, %xmm5
	addps	%xmm2, %xmm4
	andps	%xmm8, %xmm4
	orps	%xmm1, %xmm4
	mulps	%xmm0, %xmm3
	addps	%xmm12, %xmm3
	andps	%xmm8, %xmm3
	orps	%xmm9, %xmm3
	movaps	%xmm5, %xmm1
	movaps	%xmm4, %xmm2
	movaps	%xmm3, %xmm12

SIMD 命令がずらっと並んでるのが見て取れて気持ちよくなれると思います。あと、条件式 if(distance > 0.0f) が全 SIMD lane で false であれば以降の処理をスキップする処理まで入っているようです。このケースでは不利に働く最適化ではありますが。

で、それなりにがんばって ISPC で実際にパーティクルを衝突させてみた図

bin: https://github.com/i-saint/ispc_SPH/blob/master/bin/test_impulse.zip
src: https://github.com/i-saint/ispc_SPH

hash grid、Intel TBB による並列化、あとは単純に周囲のパーティクルを見てめり込んだ分押し返すだけの処理。
私のマシンでは 115000 パーティクル出しても 13ms 程度で計算できています。SPH だともっと重くなり、60000 くらいで 16ms 前後行く感じです。これは CUDA 版とあまり変わらない速度で、若干遅い程度です。推測ですが、hash grid の cell 内の要素数は大抵そんなに多くないので、CUDA 版は GPGPU の並列性をあまり活かせていなかったのかもしれません。あと、衝突結果を CPU 側に転送するところが相当時間を食っていたので、総じてこのスケールと用途では GPGPU に向いてない問題だったのかもしれません。

(ちなみに、ISPC では SIMD lane 数分のプログラムの集合を gang と呼んでおり、これは丁度 NVIDIA 系 GPU で言う warp に相当します。 NVIDIA は warp を 32 スレッドの集合と言い張っていますが、プログラムカウンタを共有する性質を考えると 32 lane の SIMD と言った方が実態に即しています)

SPH は濃度や速度を hash grid にキャッシングして近似することでかなり速くできることを教わったので、アルゴリズムの改良次第では CPU で 100000 パーティクルも夢ではなさそうです。加えて CPU は GPU と比べてスペック差が少ないので、最軽量モードでも 30000 パーティクルくらいは出していいと予想されます。そのくらい出せればそこそこいい見栄えにできるでしょう。
なんだか CUDA 版が抱えていた諸問題を一挙に解決できそうです。ありがとう ISPC。GPGPU なんていらんかったんや。

余談ですが、Intel は OpenCL の CPU 実装 & SDK も提供しています。
http://software.intel.com/en-us/articles/vcsource-tools-opencl-sdk/
SDK には LLVM ベースのコンパイラや、コンパイル結果のアセンブラを即座に見れるツールなんかもついています。こちらも適切に書けばそこそこいい感じに SIMD 化されたコードを生成してくれるようです。ただ、OpenCL は ISPC よりジェネリックな言語 (=言語仕様の制限が緩い) である分、些細なうっかりで効率の悪いコードを生成しがちな感触です。OpenCL には GPU による実装もあるという強力なアドバンテージがありますが、がんばって CPU に最適化したい場合 ISPC に分があると思います。