SIMDを用いて評価関数計算を高速化したかった話

　コンピュータ将棋において評価関数を計算する際にSIMDによって高速化する手法が知られています。

　基本的には野田さんのブログが詳しいので、深く知りたい方はそちらを参照してください。

　今僕が開発している将棋ソフトでは手番なしの絶対2駒(kp, pp)のみを評価項目として利用しており、パラメータはすべてint型(32bit)で確保しているという違いがあるため、そのままコピペでは動きません。

　命令レベルの高速化なんてさっぱりやったことないのですが、まぁやるしかないようです。

　とりあえずいきなり将棋ソフトに実装する前にテスト用のコードを書いてみました。

#include<vector>
#include<random>
#include<iostream>
#include<chrono>
#include<cassert>
#include<immintrin.h>

using namespace std;
using namespace chrono;

//疑似的なKP, PP
constexpr int K_NUM = 81, P_NUM = 1534;
int kp[K_NUM][P_NUM];
int pp[P_NUM][P_NUM];
constexpr int LIST_SIZE = 38;

int calcByForLoop(int k1, int k2, const vector<int>& list) {
    //forループで計算
    int sum = 0;
    for (int i = 0; i < LIST_SIZE; i++) {
        //KPの計算
        sum += kp[k1][list[i]];
        sum += kp[k2][list[i]];

        for (int j = i; j < LIST_SIZE; j++) {
            sum += pp[list[i]][list[j]];
        }
    }
    return sum;
}

int calcBySIMD(int k1, int k2, const vector<int>& list) {
    //SIMD計算
    int sum = 0;

    //0初期化したものを準備
    __m256i sum256 = _mm256_setzero_si256();
    for (int i = 0; i < LIST_SIZE; i++) {
        //kpはそのまま足す
        sum += kp[k1][list[i]];
        sum += kp[k2][list[i]];

        int j = i;
        for (; j + 8 < LIST_SIZE; j += 8) {
            //まずはlist[j]から8要素ロード
            __m256i indexes = _mm256_load_si256(reinterpret_cast<const __m256i*>(&list[j]));

            //indexesをオフセットとしてppから8要素ギャザー
            __m256i gathered = _mm256_i32gather_epi32(reinterpret_cast<const int*>(&pp[list[i]]), indexes, 4);

            //足し合わせる
            sum256 = _mm256_add_epi32(sum256, gathered);
        }

        for (; j + 4 < LIST_SIZE; j += 4) {
            //list[j]から4要素ロード
            __m128i indexes = _mm_load_si128(reinterpret_cast<const __m128i*>(&list[j]));

            //indexesをオフセットとしてppから4要素ギャザー
            __m128i gathered = _mm_i32gather_epi32(reinterpret_cast<const int*>(&pp[list[i]]), indexes, 4);

            //256bitに拡張
            __m256i expanded = _mm256_insertf128_si256(_mm256_setzero_si256(), gathered, 0);

            //足し合わせる
            sum256 = _mm256_add_epi32(sum256, expanded);
        }
        for (; j < LIST_SIZE; j++) {
            sum += pp[list[i]][list[j]];
        }
    }

    //64bitずらして足し合わせる
    sum256 = _mm256_add_epi32(sum256, _mm256_srli_si256(sum256, 8));

    //128bitずらして足し合わせる
    __m128i sum128 = _mm_add_epi32(_mm256_extracti128_si256(sum256, 0), _mm256_extracti128_si256(sum256, 1));

    //上位64bitを32bitずつにバラして足す
    int sum32[2];
    _mm_storel_epi64(reinterpret_cast<__m128i*>(sum32), sum128);
    sum += sum32[0];
    sum += sum32[1];

    return sum;
}

int main() {
    //適当にkp,ppの値を設定
    random_device seed_gen;
    default_random_engine engine(seed_gen());
    uniform_int_distribution<int> dist_param(-2000, 2000);

    int cnt = 1;
    for (int i = 0; i < K_NUM; i++) {
        for (int j = 0; j < P_NUM; j++) {
            kp[i][j] = dist_param(engine);
        }
    }
    for (int i = 0; i < P_NUM; i++) {
        for (int j = 0; j < P_NUM; j++) {
            pp[i][j] = dist_param(engine);
        }
    }

    //ランダムにK, Pを設定
    uniform_int_distribution<int> dist_K(0, 81);
    uniform_int_distribution<int> dist_P(0, 1534);

    constexpr int num = 1000000;

    for (int i = 0; i < num; i++) {
        int k1 = dist_K(engine), k2 = dist_K(engine);
        vector<int> list(LIST_SIZE);
        for (int i = 0; i < LIST_SIZE; i++) {
            list[i] = dist_P(engine);
        }
        assert(calcByForLoop(k1, k2, list) == calcBySIMD(k1, k2, list));
    }

    double time_of_ForLoop = 0, time_of_SIMD = 0;
    for (int i = 0; i < num; i++) {
        int k1 = dist_K(engine), k2 = dist_K(engine);
        vector<int> list(LIST_SIZE);
        for (int i = 0; i < LIST_SIZE; i++) {
            list[i] = dist_P(engine);
        }

        auto start = high_resolution_clock::now();
        calcByForLoop(k1, k2, list);
        auto end = high_resolution_clock::now();
        auto elapsed = end - start;
        time_of_ForLoop += duration_cast<duration<double>>(elapsed).count();
    }

    for (int i = 0; i < num; i++) {
        int k1 = dist_K(engine), k2 = dist_K(engine);
        vector<int> list(LIST_SIZE);
        for (int i = 0; i < LIST_SIZE; i++) {
            list[i] = dist_P(engine);
        }

        auto start = high_resolution_clock::now();
        calcBySIMD(k1, k2, list);
        auto end = high_resolution_clock::now();
        auto elapsed = end - start;
        time_of_SIMD += duration_cast<duration<double>>(elapsed).count();
    }

    cout << "ForLoop : " << time_of_ForLoop / num << endl;
    cout << "SIMD    : " << time_of_SIMD / num << endl;
}

　値が同じかどうかのチェックは抜けてくれるので計算としては合っているようです。しかし速度の結果は

ForLoop : 1.787e-08
SIMD : 1.566e-08

　と、ものすごく速くなっているわけではない感じでした。キャッシュの影響とかがあるかもしれないので導入してみないとわからないところですが、"NPS全体で約9%の高速化"というほどには至らないような気がします。手番なしの2駒関係だと恩恵が薄いということなのでしょうか。プログラムが遅いと学習にも時間がかかるためできるだけ高速化したいところですが、なかなか上手くはいかないようです。