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的一些用法和注意的问题。