prime's diary

そすうの日々を垂れ流しちゃうやつだよ

『フカシギの数え方』の数え上げを、全探索のままGPGPUで高速化した 【KMCアドベントカレンダー 2021 4日目】

ネタ

この記事はKMCアドベントカレンダー2021の4日目の記事です。

adventar.org

概要

元ネタ：『フカシギの数え方』


www.youtube.com

問題としてはN×Nマスのグリッド（頂点としては (N+1) \times (N+1) 頂点）を左上から右下まで移動する経路であって、同じところを二度通らないものの数を数える、というものになっています。

『フカシギの数え方』でお姉さんが数え上げていた問題を、実際に全探索で解いてみて、組合せ爆発の凄さを体感しつつ、GPGPUで頑張って高速化してみます。

CPUで全探索

これを全探索で解くためには、これまでに通った場所を覚えておいて、そこを通らないようにしつつ、右下の頂点までたどり着けるものを探せばよいです。 これはバックトラックによって簡単に数え上げることができます。

RustでCPU向けに実装したものがこちらになります

github.com

バックトラック部分の本体は

fn solve_impl(n: isize, visit: &mut BitVec, row: isize, col: isize) -> usize {
    if row == n && col == n {
        return 1;
    }
    let mut ans = 0;
    let d = [(1, 0), (0, 1), (-1, 0), (0, -1)];
    for (dr, dc) in d {
        let new_row = row + dr;
        let new_col = col + dc;
        if min(new_row, new_col) < 0 || max(new_row, new_col) > n {
            continue;
        }
        let new_pos = new_row * (n + 1) + new_col;
        if visit.get(new_pos as usize).unwrap() {
            continue;
        }
        visit.set(new_pos as usize, true);
        ans += solve_impl(n, visit, new_row, new_col);
        visit.set(new_pos as usize, false);
    }
    ans
}

こんな感じで容易に書くことができます。 今いる場所から4方向に対して、枠外になっていないか、既に通っていないかを確認して、大丈夫なら移動する、ということを再帰的に繰り返しています。

並列化のため、最初の2Nステップほどを並列に実行するようになっています（async-stdとfutures-rsを利用しています）。

実行するとこんな感じです（oneesan(N)でN×Nマスでの数え上げを行います）

$  cargo run --release
    Finished release [optimized] target(s) in 0.54s
     Running `target/release/kazoeage-oneesan`
oneesan(1) = 2, elapsed: 0.001s
oneesan(2) = 12, elapsed: 0.001s
oneesan(3) = 184, elapsed: 0.001s
oneesan(4) = 8512, elapsed: 0.002s
oneesan(5) = 1262816, elapsed: 0.019s
oneesan(6) = 575780564, elapsed: 6.891s
oneesan(7) = 789360053252, elapsed: 12195.893s

実行時間はCPUの性能等に依存します。 この調子で行くと、oneesan(8)の計算には1年半ほどかかる計算になります。

GPUで全探索

今回はoneesan(7)の計算をGPGPUで高速化してみます。 まず、GPUは並列化しないとその性能を発揮できないため、最初の一定ステップをCPUで全列挙し、そこからゴールまでの経路をそれぞれGPUで数え上げることにします。 しかし、GPUは再帰関数の処理が苦手なため、並列化しただけではとても効率が悪いです。 手元で試したところ、oneesan(6)の計算に40秒ほどかかってしまいました。

効率が悪くなる主な原因は、GPUの内部構造にあります。GPUは内部的には複数スレッドを一定数(32や64程度)束ねて、すべてのスレッドで同一の命令を実行する仕組みになっています(SIMT)。 このままだとスレッド間で計算内容が異なる場合に対応できないので、一方のスレッドの計算を行っている間は、別の命令を実行したいスレッドを待機させます。 その結果、バックトラックでの分岐が発生するたびに、一度に実行できるスレッドが絞られていき、束ねたスレッドのうち、最終的には同時に走るのは1スレッドのみになってしまいます。

今回は、再帰関数を明示的に用意した配列で実装したスタックを用いたループに展開することで、分岐して別行動になったスレッド同士が再び同時に実行できるチャンスを増やすことで、効率化を行いました。 その実装がこちらになります。

github.com

細かい点としては、ローカルメモリ消費量を抑えるために、再帰1段分のスタック用配列サイズを2bitに抑える工夫を入れてあります。

これらの工夫の結果、手元のGPUではoneesan(6)を1.8秒、oneesan(7)を1728秒で求めることができるようになりました。 oneesan(7)の計算は実に7倍高速になりました！ このペースだとoneesan(8)も83日程度で求められる計算になります。ぎりぎり希望が持てる時間になりました（？） CPUとGPUの性能比を考えるとそう悪くない結果だと思います。

『フカシギの数え方』ではスーパーコンピュータは1秒間に2000億通り数えられる想定でしたが、このプログラムでは手元のGPU(GeForce RTX 3080)1枚で毎秒約4.5億通り数えられます。 現代のGPUを500枚程度用意すれば、おねえさんのスーパーコンピュータより速く数え上げることができるかもしれません。