ComputeShader基礎用法系列之二

上一節我們介紹了如何讓一個最簡單的ComputeShader跑起來並且使我們看到效果。接下來我們看看如何向ComputeShader中傳值以及如何回讀到CPU。

ComputeShader類提供了一些傳值方法，基礎類型如SetFloat,SetInt,SetVector,SetMatrix，集合如SetMatrixArray,SetFloats,SetInts,SetBuffer等。我們具體看一下怎么使用。

在我們上一節compute文件基礎上，在RWTexture2D<float4> Result;下聲明一個float4 color;如下：

#pragma kernel CSMain
RWTexture2D<float4> Result;
float4 color;

[numthreads(8,8,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
    Result[id.xy] = color;
}

然后再C#中SetVector，如下：

public ComputeShader compute;
    // Start is called before the first frame update
    public Renderer renderer;

    public Color color;
    private int kernel;
    void Start()
    {
        kernel = compute.FindKernel("CSMain");
        RenderTexture rt = new RenderTexture(512,512,32);
        rt.enableRandomWrite = true;
        rt.Create();
        compute.SetTexture(kernel,"Result",rt);
        renderer.material.SetTexture("_BaseMap",rt);
    }

    // Update is called once per frame
    void Update()
    {
        compute.SetVector("color",color);
        compute.Dispatch(kernel,64,64,1);
    }

當我們不斷的在CPU傳入不同的顏色值時，貼圖的顏色就會發生變化：

 

在舉個傳float的例子，我們加入一個時間變量，在Compute文件中加入time屬性，返回結果修改為如下：

Result[id.xy] = color + sin(time).xxxx;

c#中：

compute.SetFloat("time",Time.time);

紋理就會隨着時間不斷變化顏色：

 

 

 如果我們要傳一個集合呢，比如說我們要做一個顏色列表，讓GPU根據線程ID取不同的顏色值，可以如下（color可以隱式轉換為vec4，但是color數組不行，需要手動轉換成vec4數組）：

compute.SetVectorArray("colors",colorVectors);

compute文件中如下：

#pragma kernel CSMain
RWTexture2D<float4> Result;
float4 colors[10];
float time;
[numthreads(8,8,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
    int colorIndex = id.x%10;
    Result[id.xy] = colors[colorIndex] + sin(time).xxxx;
}

最后結果如下：

 

 

 為了進行下面的內容，我們先做一個例子，通過ComputeShader進行物體的Transform設置，這里面涉及到傳入數據到GPU以及回讀到CPU。

Transform這種結構數據需要使用StructuredBuffer類型進行傳遞，因為需要讀寫，所以我們用RWStructuredBuffer。但是我們會發現ComputeShader類只給了設置數據的方法，但是沒有返回數據到CPU的方法。這時我們就需要使用ComputeBuffer進行數據傳遞和回讀。

c#代碼如下：

using UnityEngine;

public class ComputeShaderTest : MonoBehaviour
{
    struct InstanceData
    {
        public Vector3 position;
        public Vector3 rotation;
        public Vector3 scale;
    }
    public ComputeShader compute;
    public GameObject prefab;
    public int count = 32;
    // Start is called before the first frame update
    private int kernel;
    void Start()
    {
        kernel = compute.FindKernel("CSMain");
        ComputeBuffer cb = new ComputeBuffer(count,36);
        var instanceData = new InstanceData[count];
        cb.SetData(instanceData);
        compute.SetBuffer(kernel,"data",cb);
        compute.Dispatch(kernel,4,1,1);
        cb.GetData(instanceData);
        for (int i = 0; i < count; i++)
        {
            var go = GameObject.Instantiate(prefab);
            go.transform.position = instanceData[i].position;
            go.transform.localEulerAngles = instanceData[i].rotation;
            go.transform.localScale = instanceData[i].scale;
        }
        cb.Release();
    }
}

compute文件如下：

// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel CSMain
struct InstanceData{
    float3 position;
    float3 rotation;
    float3 scale;
};
RWStructuredBuffer<InstanceData> data;
[numthreads(8,1,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
    int idx = id.x;
    data[idx].position = float3(id.x,0,0);
    data[idx].rotation = float3(id.x,0,0);
    data[idx].scale = float3(idx,idx,idx);
}

最后效果如圖：

 

 可以看到我們通過ComputeBuffer的SetData和GetData傳遞和回讀數據，需要注意的是ComputeBuffer是有大小限制的，大小為不能超過2048的、4的倍數的一個自然數。

 我們最后看一下傳遞float數組的例子，根據一個float數組划線：

c#代碼如下：

using UnityEngine;

public class ComputeShaderTest : MonoBehaviour
{
    public ComputeShader compute;
    // Start is called before the first frame update
    public Renderer renderer;
    private int kernel;
    void Start()
    {
        float[] fArray = new float[]
        {
           1, 1, 0, 1, 0, 0, 0, 1,
        };
        kernel = compute.FindKernel("CSMain");
        RenderTexture rt = new RenderTexture(512,512,32);
        rt.enableRandomWrite = true;
        rt.Create();
        compute.SetTexture(kernel,"Result",rt);
        compute.SetFloats("fValues",fArray);
        compute.Dispatch(kernel,64,64,1);
        renderer.material.SetTexture("_BaseMap",rt);
    }
}

compute代碼如下：

#pragma kernel CSMain
RWTexture2D<float4> Result;
float fValues[8];

[numthreads(8,8,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
    uint idx = id.x;
    float colorFactor = fValues[idx%8];
    Result[id.xy] = float4(colorFactor,colorFactor, colorFactor, 0);
}

結果如下：

 

 我們看到computeshader的結果並沒有根據我們傳入的結果進行分布，而是隔了8個像素畫一次。那么問題出在了哪里呢？因為內存對齊！

我們設置float數組必須滿足內存對齊16字節，因為GPU運算單元是4-component的，是基於向量運算的（即float4），所以我們傳入的float數組（長度為8）到GPU的時候，只有第1個和第5個元素會被當做float值，其他元素都當做補齊元素了，這樣就導致了數據丟失。而最后的結果之所以是8個像素一條線，是因為我們再compute shader中聲明的float數組是8長度的。但是我們只向GPU傳入了2長度。所以剩下的會被自動補0。那么我們就必須修改CPU傳入數組為：

float[] fArray = new float[]
        {
           1, 0,0,0,
           1,0,0,0,
           0,0,0,0,
           1,0,0,0,
           0,0,0,0,
           0,0,0,0,
           0,0,0,0,
           1,0,0,0
        };

結果如下：

 

 這樣結果就正確了。接下來我們會學習基於ComputeShader的一些用法和注意的問題。