【Unity】コンピュートシェーダにサクッと入門 - GPGPUの説明から最小コードを書くまで

Unityのコンピュートシェーダについて簡単にまとめました。

GPUを描画以外の用途で使うGPGPU
DirectComputeとCompute Shader
UnityのCompute Shader
コンピュートシェーダを書く
スクリプトを書く
実行する
参考

GPUを描画以外の用途で使うGPGPU

GPUは描画をするために作られた処理装置です。
フラグメントシェーダでは画面上の各ピクセルごとに計算を行うので、
CPUとは異なり大量に並列で計算を行うことができるように作られています。

このような並列処理に強い特性を活かして、GPUは次第に描画処理以外の計算、
例えばシミュレーションの計算や仮想通貨のマイニングなどにも使われるようになりました。

このように描画処理以外にGPUを活用することを
GPGPU(General Purpose computing on Graphics Processing Units)といいます。

DirectComputeとCompute Shader

さてGPGPUは初めはグラフィックスAPIをうまく利用して実装していたようですが、
GPGPUに注目が集まるにつれて次第に専用のAPIが開発されていきます。

そのうちの一つとしてDirectComputeがあります。
これを使ってCompute Shaderと呼ばれるプログラムを書くと、
GPUの計算資源をグラフィックスではなくGPGPUに適した形で使うことができます。

UnityのCompute Shader

さてUnityのコンピュートシェーダは、Unity上でGPGPUを実現するための機能です。
マニュアルを見ると、これはDirectComputeとほぼ同じものであるという説明がされています。

Compute shaders in Unity closely match DirectX 11 DirectCompute technology.
https://docs.unity3d.com/Manual/class-ComputeShader.html

対応しているプラットフォームは今のところ以下の通りです。

DirectX11かDirectX12かつシェーダモデル5.0以上のWindows
MetalのiOSやmacOS
VulkanのAndroidやWindows
OpenGL ES 3.1以上のAndroid
PS4とか

コンピュートシェーダを書く

それでは実際にUnityでコンピュートシェーダを書いてみます。
Unityでコンピュートシェーダを作成するにはAssets > Create > Shader > Compute Shaderを選択します。

f:id:halya_11:20191119132636p:plain

すると.computeという拡張子のファイルが生成されるので、中身を以下のように書き換えます。

// カーネルを指定
#pragma kernel CSMain

// コンピュートシェーダにより読み書きするバッファ
RWStructuredBuffer<int> intBuffer;

// スレッド数を指定
[numthreads(3, 1, 1)]
void CSMain (uint3 id : SV_GroupThreadID)
{
    intBuffer[id.x] += id.x;
}

上記のソースコードでは、まず#pragma kernel CSMainで並列実行する処理を指定しています。
ここではCSMain関数を並列実行します。
また、この一つの処理をカーネルと呼びます。

そしてカーネルを処理する単位をスレッドと呼びます。
例えば3つのスレッドを同時に走らせる場合は3つのカーネルが並列実行されます。
また、このスレッド数は三次元の値で表され、これらの値を掛け合わせたものが実行されるスレッド数になります。
具体的には、上記のソースコードで[numthreads(3, 1, 1)]としている部分がスレッド数の定義です。
今回は(3, 1, 1)と指定しているため、3 x 1 x 1 = 3スレッドが並列実行されます。

RWStructuredBuffer<int> intBuffer;はコンピュートシェーダが読み書きするためのバッファです。
今回のカーネルでは、SV_GroupThreadIDセマンティクスを使って今処理しているスレッドのインデックスを取得し、それに対応するバッファを適当に書き換えています。

なおセマンティクスを使うと、他にもいろいろな情報をカーネルで扱えます。
これについては以下の記事がわかりやすくまとまっていたのでご参照ください。

scrapbox.io

スクリプトを書く

さてこうして作ったコンピュートシェーダはスクリプトから制御する必要があります。
まずは全文を掲載します。

using UnityEngine;

public class ComputeShaderExample : MonoBehaviour
{
    public ComputeShader computeShader;
    private int _kernelIndex;
    private ComputeBuffer _computeBuffer;

    [SerializeField]
    private int[] _results;

    void Start()
    {
        _results = new int[3];

        // カーネルを検索
        _kernelIndex = computeShader.FindKernel("CSMain");
        // バッファを作ってセット
        _computeBuffer = new ComputeBuffer(3, sizeof(int));
        _computeBuffer.SetData(_results);
        computeShader.SetBuffer(_kernelIndex, "intBuffer", _computeBuffer);
    }

    private void Update()
    {
        if (Input.GetKeyDown(KeyCode.Space)) {
            // グループを指定して実行
            computeShader.Dispatch(_kernelIndex, 1, 1, 1);
            // バッファからデータを取得
            _computeBuffer.GetData(_results);
        }
    }

    private void OnDestroy()
    {
        _computeBuffer.Release();
    }
}

さて上記のソースコードでは、まずStart()でコンピュートシェーダの初期化を行っています。
具体的には、computeShader.FindKernel()でカーネルのインデックスを取得しています。
また並列実行するサイズに応じたComputeBufferを生成し、コンピュートシェーダにセットしています。

コンピュートシェーダの実行はUpdate()で行っています。
まず、ComputeShader.Dispatch()でカーネルを指定して並列処理を実行しています。
ここで、このメソッドの第二～第四引数は「グループ」を指定しています。
グループはスレッドを並列実行する単位で、並列数はスレッドと同じく3次元の値で管理されます。
例えば3スレッドを並列実行するコンピュートシェーダを2グループ実行すると3 x 2 = 6つの処理が並列実行されることになります。
なお今回はグループは一つだけにしています。

コンピュートシェーダによる計算結果はComputeBuffer.GetData()で取得します。

なお、ComputeBufferは明示的に開放する必要があるので注意が必要です。
解放するにはComputeBuffer.Release()を実行します。