Unity ComputeShader 使用GPU快速计算复杂问题

前言

遇到一个问题，需要大量的计算，在Unity中直接写会长时间的阻塞主线程，正好使用ComputeShader让GPU来帮我算。

项目

创建ComputeShader

右键Project空白的地方或者文件夹Create->Shader->Compute Shader

修改文件名为TestCompShader

编写CompturShader

修改代码，每一行注释写的很详细

//每个 #kernel告诉编译哪个函数 你可以有很多kernel
#pragma kernel CSMain

//由程序赋值的RenderTexture（Unity中必须给RenderTexture设置enableRandomWrite标志）
RWTexture2D<float4> Result;

//每个线程组的线程数量，为了满足每个线程都只有一个像素对应
//这里是 16 x 16 = 256 正好是传入图片的大小，最后的1是一维
[numthreads(16, 16, 1)]

//CSMain接受的参数为：(uint3 id : SV_DispatchThreadID)
void CSMain(uint3 id : SV_DispatchThreadID)
{
    //将像素的x坐标映射到0-1范围 也就是让图片从黑到红
    Result[id.xy] = float4((id.x / 255.0), 0.0, 0.0, 0.0);
}

创建Unity代码

在Unity中创建名为ComputeBufferExample的脚本
并写入下面的内容

using UnityEditor;
using UnityEngine;

public class ComputeBufferExample : MonoBehaviour
{
    /// <summary>
    /// Compute Shader
    /// </summary>
    public ComputeShader shader;

    private void Start()
    {
        RunShader();
    }

    /// <summary>
    /// 运行Compute Shader
    /// </summary>
    void RunShader()
    {
        //获取kernel
        int kernelHandle = shader.FindKernel("CSMain");

        //创建ComputeBuffer
        //RenderTexture tex = new RenderTexture(256, 256, 24);
        RenderTexture tex = new RenderTexture(256, 256, 24, RenderTextureFormat.ARGB32);
        //这里必须要设置为true，否则无法使用SetTexture
        tex.enableRandomWrite = true;
        tex.Create();

        //设置Compute Shader参数
        shader.SetTexture(kernelHandle, "Result", tex);
        //运行Compute Shader
        shader.Dispatch(kernelHandle, 16, 16, 1);

        //将结果保存到 Assets 路径下的 test.renderTexture
        AssetDatabase.CreateAsset(tex, "Assets/test.renderTexture");
    }
}

场景布置

创建空物体并挂载脚本

运行场景

运行，并在文件中找到保存的图片
Tip:停止运行就看不到了

参考

Unity | 浅谈 Compute Shader