【Unity】compute shader的介绍和使用

RuntimeMapMaker3D-Pro

---

参考原视频：

https://www.youtube.com/watch?v=BrZ4pWwkpto

 Compute shader 属于shader的一种，它使用CG language，但是和传统的Vertex & fragement shader或者是Unity的SurfaceShader又有所不同。
Compute Shaders是在GPU运行却又在普通渲染管线之外的程序，通过Compute Shader我们可以将大量可以并行的计算放到GPU中计算从而节省CPU资源，下面就和大家介绍下Compute Shader是如何使用的。
 首先我们创建一个Compute shader：

// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel CSMain

// Create a RenderTexture with enableRandomWrite flag and set it
// with cs.SetTexture
RWTexture2D<float4> Result;

[numthreads(8,8,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
    // TODO: insert actual code here!

    Result[id.xy] = float4(id.x & id.y, (id.x & 15)/15.0, (id.y & 15)/15.0, 0.0);
}

 第一行kernel CSMain了定义方法入口，“CSMain”的名字必须和主方法的名字一制。第二行到CSMain之间是输入参数列表，用来存储和返回给c#数据。
 主方法之上的[numthreads(X,Y,Z)]涉及到多维度的GPU线程调用，本文中暂不涉及。本案例暂时只使用x维度，其他两维保持为1。
然后我们创建一个C#脚本：

using System.Collections;
using System.Collections.Generic;
using UnityEngine;

public class ComputeShaderTest : MonoBehaviour
{
    public ComputeShader computeShader;
    public RenderTexture renderTexture;

    private void OnRenderImage(RenderTexture source, RenderTexture destination)
    {
        if (renderTexture == null)
        {
            renderTexture = new RenderTexture(256, 256, 24);
            renderTexture.enableRandomWrite = true;
            renderTexture.Create();
        }


        computeShader.SetTexture(0, "Result", renderTexture);
        computeShader.SetFloat("Resolution", renderTexture.width);
        computeShader.Dispatch(0, renderTexture.width / 8, renderTexture.height / 8, 1);

        Graphics.Blit(renderTexture, destination);
    }

}

OnRenderImage函数会由Unity渲染时被调用，创建一张renderTexture。并把computeShader中的数据显示在renderTexture中，并传送给destination显示在当前脚本所挂载的摄像机。如果顺利的话，你将看到一张分型图：

那么现在我们尝试利用Compute shader进行一些并行计算。
首先我们创建一些立方体，并且利用CPU给这些立方体随机一下颜色。

using System.Collections.Generic;
using UnityEngine;

public class ComputeShaderTest2 : MonoBehaviour
{
    public int repeitions;
    public List<GameObject> objects;
    public ComputeShader computeShader;

    public int count = 50;
    private void Start()
    {
        CreateCubes();
    }
    private void CreateCubes()
    {
        for (int i = 0; i < count; i++)
        {
            for (int j = 0; j < count; j++)
            {
                var go = GameObject.CreatePrimitive(PrimitiveType.Cube);
                go.transform.position = new Vector3(i, j, 0);
                Color color = Random.ColorHSV();
                go.GetComponent<MeshRenderer>().material.SetColor("_Color", color);
                objects.Add(go);
                go.transform.SetParent(this.transform);
            }
        }
    }

    [ContextMenu("OnRandomize CPU")]
    public void OnRandomizeCPU()
    {
        for (int i = 0; i < repeitions; i++)
        {
            for (int c = 0; c < objects.Count; c++)
            {
                GameObject obj = objects[c];
                obj.transform.position = new Vector3(obj.transform.position.x, obj.transform.position.y, Random.Range(-0.1f, 0.1f));
                obj.GetComponent<MeshRenderer>().material.SetColor("_Color", Random.ColorHSV());
            }
        }
    }
 }

当我们将重复次数上升到1000时，可以感受到明显的延迟。现在我们将随机计算移植到ComputeShader中：

// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel CSMain

struct Cube {
    float3 postion;
    float4 color;
};

RWStructuredBuffer<Cube> cubes;
float resolution;
float repeitions;

float rand(float2 co){
    return(frac(sin(dot(co.xy, float2(12.9898,78.233))) * 43758.5453))* 1;
}

[numthreads(10,1,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
    float xPos = id.x / resolution;
    Cube cube = cubes[id.x];
    for(int i = 0; i < repeitions; i++){
        float zPos = rand(float2(xPos, cube.postion.z));
        cube.postion.z = zPos;
        float r = rand(float2(cube.color.r, cube.color.g));
        float g = rand(float2(cube.color.g, cube.color.b));
        float b = rand(float2(cube.color.b, cube.color.r));
        cube.color = float4(r, g, b, 1.0);
    }
    cubes[id.x] = cube;
}

对原来的C#代码做一些修改，创建一个结构体Cube用来传递位置和颜色数据信息，通过ComputeBuffer 将Cube数组的数据传递给ComputeShader。利用Dispatch运行一个compute shader，在X、Y和Z维度上启动指定数量的计算着色器线程组。（Dispatch是指定线程组的数量，numthreads是指定每个组的独立线程数）最后通过GetData获取到计算结果，释放掉ComputeBuffer。

using System.Collections;
using System.Collections.Generic;
using UnityEngine;

public struct Cube
{
    public Vector3 position;
    public Color color;
}
public class ComputeShaderTest2 : MonoBehaviour
{
    public int repeitions;
    public List<GameObject> objects;
    public ComputeShader computeShader;

    public int count = 50;
    private Cube[] data;

    private void Start()
    {
        CreateCubes();
    }
    private void CreateCubes()
    {
        data = new Cube[count * count];

        for (int i = 0; i < count; i++)
        {
            for (int j = 0; j < count; j++)
            {
                var go = GameObject.CreatePrimitive(PrimitiveType.Cube);
                go.transform.position = new Vector3(i, j, 0);
                Color color = Random.ColorHSV();
                go.GetComponent<MeshRenderer>().material.SetColor("_Color", color);
                objects.Add(go);
                go.transform.SetParent(this.transform);

                Cube cubeData = new Cube() { position = go.transform.position, color = color };
                data[i * count + j] = cubeData;
            }
        }
    }

    [ContextMenu("OnRandomize CPU")]
    public void OnRandomizeCPU()
    {
        for (int i = 0; i < repeitions; i++)
        {
            for (int c = 0; c < objects.Count; c++)
            {
                GameObject obj = objects[c];
                obj.transform.position = new Vector3(obj.transform.position.x, obj.transform.position.y, Random.Range(-0.1f, 0.1f));
                obj.GetComponent<MeshRenderer>().material.SetColor("_Color", Random.ColorHSV());
            }
        }
    }
    [ContextMenu("OnRandomize GPU")]
    public void OnRandomizeGPU()
    {
        int colorSize = sizeof(float) * 4;
        int vector3Size = sizeof(float) * 3;
        int totalSize = colorSize + vector3Size;

        ComputeBuffer cubesBuffer = new ComputeBuffer(data.Length, totalSize);

        cubesBuffer.SetData(data);

        computeShader.SetBuffer(0, "cubes", cubesBuffer);
        computeShader.SetFloat("resolution", data.Length);
        computeShader.SetFloat("repeitions", repeitions);
        computeShader.Dispatch(0, data.Length / 10, 1, 1);


        cubesBuffer.GetData(data);
        for (int i = 0; i < objects.Count; i++)
        {
            GameObject obj = objects[i];
            Cube cube = data[i];
            obj.transform.position = cube.position;
            obj.GetComponent<MeshRenderer>().material.SetColor("_Color", cube.color);
        }
        cubesBuffer.Dispose();
    }

}

通过实践，可以明显的感觉到ComputeShader计算速度比CPU的方式要快很多。以上就是我对ComputeShader的介绍和使用分享，如果有什么疑问欢迎留言。