Jetson Nano 上に PyCUDA を使ってライフゲームを実装する

はじめに

Jetson Nano の使い道として、シミュレーションも忘れるわけにはいきません。JetPack にプリインストールされている CUDA ツールキットにもシミュレーション関連のサンプルプログラムが含まれていて、たとえば、N-Body シミュレーション が有名です。このサンプルプログラムはもちろん Jetson Nano でも動作します。私も最初はこの N-Body シミュレーション・プログラムを改造していろいろと試そうとしたのですが、このプログラムはかなり複雑で、ちょっと試すためのひな型として使うには難しいと判断しました。
次に考えたのが、CUDA で実装したライフゲームが世界中でいろいろ公開されているはずと思ったので、それをひな型にすることです。すばらしく完成度が高い実装例 「Jetson NanoでGPGPU (CUDA) プログラミング ~ライフゲームの開発~」 を発見しましたが、ちょっとした実験のひな型に使わせていただくにはあまりに精巧な実装という感じです。さらに調査すると 「Python に Curses ライブラリがあったので，シンプルなライフゲームを作ってみるテスト」 という記事があったのでこれを参考にさせていただき、表示の部分は Python の Curses ライブラリを使うことにして、CUDA 部分は PyCUDA を使って自分で CUDA プログラミングしました。
ということで、シミュレーションを試すひな型として自作した簡単なコードをここで紹介させていただきます。

実装

作成したコードは GitHub tsutof/game_of_life_pycuda で公開中です。

使用目的が CUDA によるシミュレーションのひな型なので、簡潔さを心掛けました。

Python プログラムですが隣接セルをチェックして状態遷移する計算部分は通常の CUDA C 言語コードです。このコードは実行時にコンパイルされるので、プログラム実行開始後、最初にこのコードが CUDA コアで実行されるときは一呼吸待たされる感じです。もちろん、その後は高速に実行されます。ご覧いただいてお分かりいただけるように、このコードは一つのセルだけ対象にしています。NVIDIA GPU が内蔵する多数の CUDA コア上でこのコードは並列実行されます。

lifegame.py

def get_next_state_gpu(state, next_state):
    height, width = state.shape

    mod = SourceModule("""
        __global__ void get_next_state(int *state, int *nextState, int height, int width)
        {
            unsigned int ix = threadIdx.x + blockIdx.x * blockDim.x;
            unsigned int iy = threadIdx.y + blockIdx.y * blockDim.y;
            unsigned int idx = iy * width + ix;
            int sum = 0;
            int val, nextVal;
            if (ix >= width || iy >= height) {
                return;
            }
            val = state[idx];
            sum += state[((iy - 1) % height) * width + ((ix - 1) % width)];
            sum += state[((iy - 1) % height) * width + (ix       % width)];
            sum += state[((iy - 1) % height) * width + ((ix + 1) % width)];
            sum += state[(iy       % height) * width + ((ix - 1) % width)];
            sum += state[(iy       % height) * width + ((ix + 1) % width)];
            sum += state[((iy + 1) % height) * width + ((ix - 1) % width)];
            sum += state[((iy + 1) % height) * width + (ix       % width)];
            sum += state[((iy + 1) % height) * width + ((ix + 1) % width)];
            if (val == 0 && sum == 3) {
                nextVal = 1;
            }
            else if (val != 0 && (sum >= 2 && sum <= 3)) {
                nextVal = 1;
            }
            else {
                nextVal = 0;
            }
            nextState[idx] = nextVal;
        }
        """)
    kernel_func = mod.get_function("get_next_state")

    blk_dim = (32, 32, 1)
    grid_dim = ((width + blk_dim[0] - 1) // blk_dim[0], \
        (height + blk_dim[1] - 1) // blk_dim[1], 1)
    kernel_func(
        drv.In(state), drv.Out(next_state), \
        np.int32(height), np.int32(width), \
        block=blk_dim, grid=grid_dim)

まとめ

このプログラムには、CPU ですべて処理するモードも加えました。コマンドライン・オプションで切り替えることができます。48 行 × 157 列＝ 7536 個のセルの場合、CUDA コア処理で約 100 倍高速化されました。（隣接セルをチェックして状態遷移する計算部分のみで、表示部分は除外した倍率です。Jeton Nano MAXNモード）CPU のみで計算する場合は、シングルスレッド処理なので、これをマルチスレッド化すれば、CPU 処理の場合も何倍か高速化されると思いますが、それでも CUDA コアによる処理は CPU 処理に比べて圧倒的に高速という結果になりました。

以上です。