【Unity】ComputeShader入門

やりたいこと

Compute Shaderに入門するために、まずはComputeShaderで簡単な計算をやってみる。

横方向に256個、縦方向に256個のオブジェクトを並べたとして、それぞれの色をComputeShaderで決定してみる。少し具体的に書くと、
(0.0, 0.0, 1.0) ~ (1.0, 1.0, 1.0)の値が返ってくるようにする。（ここでは実際にオブジェクト生成はせずに値をDebug.Logで出力するだけ）

GPU(ComputeShader)

Create > Shader > Compute ShaderでComputeShaderを作成して以下のように変更する。

.shader

#pragma kernel CSMain

// 結果を格納するバッファ。CPU側に確保してもらう。
RWStructuredBuffer<float3> Result;

// スレッド数。1 * 1 * 1 = 1のスレッド
[numthreads(1,1,1)]
void CSMain (uint2 id : SV_DispatchThreadID)
{
    // SV_DispatchThreadIDでは実行中のスレッドIDが渡される。
    // idは(0, 0) ~ (255, 255)なので、これを利用してバッファのindexを取得する。
    int index = id.x + id.y * 256;
    // 
    Result[index] = float3(id.x / 255.0, id.y / 255.0, 1.0);
}

カーネルを宣言する

#pragma kernel CSMain

カーネルは関数ひとつを表しCSMainが関数名になる。後ほど、この名前に対応した関数を記述する。
以下のように、ひとつのCompute Shader内に複数のカーネルを記述できる。

#pragma kernel CSSomeKernel
#pragma kernel CSOtherKernel

結果を格納するバッファ

RWStructuredBuffer<float3> Result;

これは結果を格納するためのバッファで、C#（CPU）側で確保してもらう。値を格納する場合はRWStructuredBufferに変更する。

C#(CPU)から値を受け取る

ここでは使用していないが、以下のようにしてC#から値を受け取ることができる。

int intVal;

カーネル（メイン関数）

こちらがメイン関数でここに実際の処理を記述する。

.shader

[numthreads(1,1,1)]
void CSMain (uint2 id : SV_DispatchThreadID)
{
    // SV_DispatchThreadIDでは実行中のスレッドIDが渡される。
    // idは(0, 0) ~ (255, 255)なので、これを利用してバッファのindexを取得する。
    int index = id.x + id.y * 256;
    // 
    Result[index] = float3(id.x / 255.0, id.y / 255.0, 1.0);
}

CPU(C#)

C#はこちら。

.cs

using System.Runtime.InteropServices;
using UnityEngine;

public class Test01 : MonoBehaviour
{
    [SerializeField] private ComputeShader _computeShader;

    private ComputeBuffer _computeBuffer;
    
    void Start()
    {
        var kernelIndex = _computeShader.FindKernel("CSMain");
        var countX = 256;
        var countY = 256;
        var count = countX * countY;

        // ComputeShaderの計算結果を格納するバッファを作成。個数と1要素のサイズを指定。
        _computeBuffer = new ComputeBuffer(count, Marshal.SizeOf<Vector3>());
        // 作成したバッファをComputeShaderのバッファに結びつける。
        _computeShader.SetBuffer(kernelIndex, "Result", _computeBuffer);
        // ComputeShaderを実行する。第二引数以降でグループ数を指定。ここでは(countX, countY, 1)。
        _computeShader.Dispatch(kernelIndex, countX, countY, 1);

        // 結果取り出し。
        var data = new Vector3[count];
        _computeBuffer.GetData(data);
        foreach (var x in data)
        {
            // 結果出力。小数点以下2桁表示だと0になってしまう場合があるので、桁数指定。
            Debug.Log(x.ToString("F4"));
        }
    }

    private void OnDestroy()
    {
        // 確保したバッファはGCの対象にならないので明示的に破棄する必要がある。
        _computeBuffer.Release();
    }
}

結果を格納するバッファを確保する

.cs

_computeBuffer = new ComputeBuffer(count, Marshal.SizeOf<Vector3>());
_computeShader.SetBuffer(kernelIndex, "Result", _computeBuffer);

GPU側ではバッファを確保できないらしいのでCPU側で確保したうえでGPUに渡してあげる。バッファを確保する際は、要素数と1要素あたりのサイズを指定して必要なバッファの合計サイズを指定する。

確保したバッファはSetBufferでGPUに渡す。第二引数でGPU側のバッファ名を指定してする。

処理を実行する

.cs

_computeShader.Dispatch(kernelIndex, countX, countY, 1);

SV_DispatchThreadID

SV_DispatchThreadIDには実行中のスレッドIDが入る。IDなのだが数値が一つ渡ってくるわけではなくベクトル型になる。

ComputeShaderを実行する際には、

• [numthreads(x,y,z)]でスレッド数を指定し、
• Dispatch(kernelIndex, X, Y, Z)でグループ数を指定

している。どちらも3次元ベクトルで、SV_DispatchThreadIDは(0, 0, 0) ~ (x * X -1, y * Y - 1, z * Z - 1)となる。

今回は[numthreads(1,1,1)]、_computeShader.Dispatch(kernelIndex, countX, countY, 1);としているので、(0, 0, 0) ~ (255, 255, 0)が渡ってくる（countX, countYは256）。

int index = id.x + id.y * 256;

xとyの値を利用してResultバッファ内でのindexを取得している。

参考

https://edom18.hateblo.jp/entry/2017/05/10/083421
https://www.wwwmaplesyrup-cs6.work/entry/MapleComputeTutorial1