ComputeShader小例子之下雪

这是一个 ComputeShader 结合 Shader 实现下雪的例子。我们通过 ComputeShader 并行计算每个雪粒子的位置，并通过 Shader 将其绘制出来。

效果

原理

通过 ComputeShder 并行计算每个粒子位置，并将 Buffer 与 Shader 共享，Shader 在对应位置上绘制粒子。

代码

1.ComputeShder

#pragma kernel CSMain

//雪粒子属性（位置，存活时间）
struct Particle
{
    float3 position;
    float lifetime;
};
//buffer，用于设置或取回数据
RWStructuredBuffer<Particle> particleBuffer;
//帧时间间隔
float deltaTime;

[numthreads(256, 1, 1)]//方便计算时切割分组，仅使用一个维度的数值
void CSMain(uint3 id : SV_DispatchThreadID)
{
    float3 delta = float3(0, -2, 0); //重力
    //重力下落
    particleBuffer[id.x].position += delta * deltaTime;
    //更新存活时间
    particleBuffer[id.x].lifetime += deltaTime;
    //当粒子存活 25 秒以上时，将其重置回 y = 0 的点
    if (particleBuffer[id.x].lifetime > 25) {
        particleBuffer[id.x].lifetime = 0;
        particleBuffer[id.x].position.y = 0;
    }
}

2.Shader

Shader "SnowShader"
{
    Properties
    {
        _Tint("Tint", Color) = (1, 1, 1, 1)
    }

    SubShader
    {
        Pass
        {
            Blend SrcAlpha one
            CGPROGRAM

            #pragma target 5.0
            #pragma vertex vert
            #pragma fragment frag
            #include "UnityCG.cginc"

            //粒子数据（和 ComputeShader中对应）
            struct Particle
            {
                float3 position;
                float lifetime;
            };
            //input
            struct PS_INPUT
            {
                float4 position : SV_POSITION;
                float4 color : COLOR;
            };
            //与 ComputeShader 共享的 Buffer
            StructuredBuffer<Particle> particleBuffer;
            //通过编辑器设置的颜色（上文 Properties）
            uniform float4 _Tint;
            //Vertex shader
            PS_INPUT vert(uint vertex_id : SV_VertexID, uint instance_id : SV_InstanceID)
            {
                PS_INPUT o = (PS_INPUT)0;
                o.color = _Tint;
                //设置位置（通过 instance_id 从 buffer 里取不同位置的粒子数据）
                o.position = UnityObjectToClipPos(float4(particleBuffer[instance_id].position, 1.0f));
                return o;
            }
            //Pixel shader
            float4 frag(PS_INPUT i) : COLOR
            {
                return i.color;
            }
            ENDCG
        }
    }
    Fallback Off
}

3.Csharp

public class Snow : MonoBehaviour {

    //粒子数据结构体（和 ComputeShader 中的相同，用于设置 buffer）
    private struct Particle
    {
        public Vector3 position;
        public float lifetime;
    }
    //数量
    public int particleCount = 10000;
    //绘制粒子的材质（shader 为上文的 snow）
    public Material material;
    public ComputeShader computeShader;

    //单个结构体大小（字节，用于初始化 buffer， 3 * 4 + 4 = 16）
    private const int SIZE_PARTICLE = 16;

    //kernel ID
    private int mComputeShaderKernelID;

    //buffer，设置、取回数据、与 shader 共享等
    ComputeBuffer particleBuffer;

    //computeBuffer 一个 Group 中的线程数，用于 Dispatch 时计算分组
    private const int WARP_SIZE = 256;

    //需要的组数量
    private int mWarpCount;

    void Start()
    {
        if (particleCount <= 0)
            particleCount = 1;
        //计算组数量
        mWarpCount = Mathf.CeilToInt((float)particleCount / WARP_SIZE);
        //初始化粒子数据
        Particle[] particleArray = new Particle[particleCount];
        for (int i = 0; i < particleCount; ++i)
        {
            particleArray[i].position.x = Random.Range(-50f, 50f);
            particleArray[i].position.y = 0;
            particleArray[i].position.z = Random.Range(-50f, 50f);

            particleArray[i].lifetime = Random.Range(0f, 25f);
        }
        //创建并设置 buffer 数据
        particleBuffer = new ComputeBuffer(particleCount, SIZE_PARTICLE);
        particleBuffer.SetData(particleArray);

        mComputeShaderKernelID = computeShader.FindKernel("CSMain");

        //设置当前 kernel 的 buffer
        computeShader.SetBuffer(mComputeShaderKernelID, "particleBuffer", particleBuffer);
        //shader 共享此 buffer
        material.SetBuffer("particleBuffer", particleBuffer);
    }

    void OnDestroy()
    {
        if (particleBuffer != null)
            particleBuffer.Release();
    }

    void Update()
    {
        computeShader.SetFloat("deltaTime", Time.deltaTime);
        //执行运算
        computeShader.Dispatch(mComputeShaderKernelID, mWarpCount, 1, 1);
    }

    void OnRenderObject()
    {
        //激活给定的 pass 以进行渲染
        material.SetPass(0);
        //GPU 画图
        Graphics.DrawProceduralNow(MeshTopology.Points, 1, particleCount);
    }
}

引用资料

1. [代码资料] 【Github】 lukakostic ComputeShaderExamples