Unity : ComputeShader のシンプルなサンプル(2)

Unityのロゴ。

Unity でコンピュートシェーダ(ComputeShader)を使う方法について、シンプルに解説します。 ここでは適当な計算を GPU 上で実行して、その結果をテクスチャに書き込みます。

サンプルには次のような操作が含まれます。

  1. コンピュートシェーダを使って、テクスチャに情報を書き込む。
  2. 多次元(2次元)のスレッドを有効に活用する。

このページの解説は「Unity : ComputeShader のシンプルなサンプル(1)」に目を通していることを前提としています。 同様に、サンプルには先に目を通していることを前提としています。

カーネルの実装

全体の実装についてはサンプルを参照してください。このサンプルでは概ね次のようなコードをコンピュートシェーダで実行します。

このサンプルではカーネルは多次元スレッドで実行される点に注目してください。 (8, 8, 1) なので、8 * 8 * 1 = 16 スレッドで実行されます。

RWTexture2D<float4> textureBuffer;

[numthreads(8, 8, 1)]
void KernelFunction_A(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    float width, height;
    textureBuffer.GetDimensions(width, height);

    textureBuffer[dispatchThreadID.xy] = float4(dispatchThreadID.x / width,
                                                dispatchThreadID.x / width,
                                                dispatchThreadID.x / width,
                                                1);
}

テクスチャの用意

以前のサンプルでは、GPU で処理した結果を保存しておく領域をバッファとして配列を用意しましたが、 ここではテクスチャに処理結果を保存します。

RWTexture2D<float4> textureBuffer;

特殊な引数 SV_DispatchThreadID

以前のサンプルでは SV_DispatchThradID セマンティクスは使いませんでした。

SV_DispathThreadID は、SV_Group_ID * numthreads + SV_GroupThreadID で算出される値です。 SV_Group_ID はあるグループを (x, y, z) で示し、SV_GroupThreadID は、 あるグループに含まれるスレッドを (x, y, z) で示します。

例えば、(2, 2, 1) グループで、(4, 1, 1) スレッドで実行される、カーネルを実行するとします。 すると、その内の 1 つのカーネルは、(0, 1, 0) 番目のグループの、(2, 0, 0) 番目のスレッドで実行されます。

このとき、SV_DispatchThreadID は、(0, 1, 0) * (4, 1, 1) + (2, 0, 0) = (0, 1, 0) + (2, 0, 0) = (2, 1, 0) になります。

少々分かり難いですが、要するに、「そのカーネルを実行するスレッドが、 すべてのスレッドの中でどこに位置するか(x,y,z)」を示しています。

テクスチャ(ピクセル)に値を書き込む

スクリプト側の設定も重要になりますが、dispatchThreadID.xy は、テクスチャ上のあるピクセルを示します。 dispatchThreadID.xy がすべてのピクセルを示すように、上手くグループとスレッドを設定します。 スクリプトについては後述します。

textureBuffer[dispatchThreadID.xy] = float4(dispatchThreadID.x / width,
                                            dispatchThreadID.x / width,
                                            dispatchThreadID.x / width,
                                            1);

このサンプルでは仮に 512x512 のテクスチャを用意していますが、dispatchThreadID.x が 0 ~ 511 を示すとき、 dispatchThreadID / width は、0 ~ 0.998… を示します。

つまり、dispatchThreadID.xy の値( = ピクセル座標)が大きくなるにつれて、黒から白に塗りつぶしていくことになります。

テクスチャは、RGBA チャネルから構成され、それぞれ 0 ~ 1 で設定します。 すべて 0 のとき、完全に黒くなり、すべて 1 のとき、完全に白くなります。

スクリプトの実装

テクスチャの用意

以前のサンプルでは、コンピュートシェーダの計算結果を保存するために配列のバッファを用意しました。 今回のサンプルでは、代わりにテクスチャを用意します。

…
RenderTexture renderTexture_A;
…
void Start()
{
    this.renderTexture_A = new RenderTexture(512, 512, 0, RenderTextureFormat.ARGB32);
    this.renderTexture_A.enableRandomWrite = true;
    this.renderTexture_A.Create();

解像度とフォーマットを指定して(レンダー)テクスチャを初期化します。 このとき、"RandomWrite" を有効にしてテクスチャを生成する点に注意します。

スレッド数の取得

カーネルのインデックスが取得できるように、カーネルがどれくらいのスレッド数で実行できるかも取得できます(スレッドサイズ)。

struct ThreadSize
{
    public int x;
    public int y;
    public int z;

    public ThreadSize(uint x, uint y, uint z)
    {
        this.x = (int)x;
        this.y = (int)y;
        this.z = (int)z;
    }
}

ThreadSize kernelThreadSize_KernelFunction_A;

void Start()
{
…
    uint threadSizeX, threadSizeY, threadSizeZ;

    this.computeShader.GetKernelThreadGroupSizes
	(this.kernelIndex_KernelFunction_A,
	 out threadSizeX, out threadSizeY, out threadSizeZ);

    this.kernelThreadSize_KernelFunction_A
	= new ThreadSize(threadSizeX, threadSizeY, threadSizeZ);
…

カーネルインデックスと同様に、スレッドサイズについても定数で指定することはできますが、 現実的にはスクリプトから取得する方が良いでしょう。

カーネルの実行

コンピュートシェーダでテクスチャに書き込んだ結果。

Dispath メソッドで処理を実行します。このとき、グループ数の指定方法に注意します。 この例では、グループ数は「テクスチャの水平(垂直)方向の解像度 / 水平(垂直)方向のスレッド数」で算出しています。

水平方向について考えるとき、テクスチャの解像度は 512、スレッド数は 8 ですから、 水平方向のグループ数は 512 / 8 = 64 になります。同様に垂直方向も 64 です。 したがって、合計グループ数は 64 * 64 = 4096 になります。

void Update()
{
    this.computeShader.Dispatch
        (this.kernelIndex_KernelFunction_A,
         this.renderTexture_A.width / this.kernelThreadSize_KernelFunction_A.x,
         this.renderTexture_A.height / this.kernelThreadSize_KernelFunction_A.y,
         this.kernelThreadSize_KernelFunction_A.z);

    plane_A.GetComponent<Renderer>().material.mainTexture = this.renderTexture_A;

言い換えれば、各グループは 8 * 8 * 1 = 64 (= スレッド数) ピクセルずつ処理することになります。 グループは 4096 あるので、4096 * 64 = 262,144 ピクセル処理します。

画像は、512 * 512 = 262,144 ピクセルなので、ちょうどすべてのピクセルを並列で処理できたことになります。

異なるカーネルの実行

もう一方のカーネルは、x ではなく、 y 座標を使って塗りつぶしていきます。 このとき 0 に近い値、黒い色が下のほうに表れている点に注意します。 テクスチャを操作するときは原点を考慮しなければならないこともあります。

多次元スレッド、グループの利点

このサンプルのように、多次元の結果が必要な場合、あるいは多次元の演算が必要な場合には、 多次元のスレッドやグループが有効に働きます。

もしこのサンプルを 1 次元のスレッドで処理しようとすると、縦方向のピクセル座標を任意に算出する必要があります。

画像処理でいうところのストライド、例えば 512x512 の画像があるとき、 その 513 番目のピクセルは、(0, 1) 座標になると算出する必要がある。

演算数は削減したほうが良いですし、高度な処理を行うにしたがって複雑さは増します。 コンピュートシェーダで処理を実行するときは、多次元にするべきかどうかを検討したほうが良いでしょう。

厳密には、データの読み書き効率にも影響が出るようですが割愛。