【Unity】计算着色器编写和应用
前言
想想过去学习DirectX12时，添加常量缓冲区、上传\读取GPU资源，都需要申请缓冲区、设置命令，其中要考虑方方面面的内容，相比起来Unity要方便多了，现在试着实现一下计算着色器的内容。
如果说计算着色器有什么好处，其一是GPU通用计算，以GPU数量巨多的核心，做并行运算好处多多；其二是，计算着色器所处GPU，对后台缓冲区进行处理，无需频繁进行CPU与GPU的传输开销，例如高斯模糊等后处理操作，可以在GPU渲染，在GPU后处理，在GPU呈现，省下了每帧CPU与GPU切换操作。
实验一：向量加减
DirectX12的第一个实例也是这个，向GPU传递向量，计算后传送会CPU，然后程序写入文件，这个写法我是参照Unity中的ComputeShader入门。
在之前的书中有过描述，我们可以将线程组想象成三维Cube集合，可以分派（Dispatch）XYZ个线程组，而每个线程组内部含有UVW个线程。
蓝色是很多线程组，每个线程组包含很多红色的线程
想起来似乎有些烧脑，但我们平时派遣的模式还是比较规整的，基本是一维和二维。
我们现在需要一个任务，计算两个长度为16的向量之和，我们希望申请2x2x1=4个线程组，每个线程组内包含2x2x1=4个线程，总共4x4=16个线程来完成这个任务。
我们直接来看整个compute shader：
//VectorAdd.compute
#pragma kernel CSMain
#define thread_group_x 2
#define thread_group_y 2
#define thread_x 2
#define thread_y 2
RWStructuredBuffer<float3> Result;
RWStructuredBuffer<float3> preVertices;
RWStructuredBuffer<float3> nextVertices;
[numthreads(2,2,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
int index = id.x
+ (id.y * thread_x * thread_group_x)
+ (id.z * thread_group_x * thread_group_y * thread_x * thread_y);
Result[index] = preVertices[index] + nextVertices[index];
}
#pragma kernel CSMain就是定义运算核函数的入口点。
下面四个define只是定义常量，只是为了程序好看，说明我有2x2的线程组以及每个线程组包含2x2的线程。
RWStructuredBuffer为可读写（Read and Write）结构化缓冲区，是模版类型，模版参数可以为float3这种内置类型，也可以是自定义的struct，在这里，pre和next为输入的两个向量，Result为输出结果。
[numthreads(2,2,1)]是定义线程组内的线程分布，至于线程组的分布，是在C#程序代码中分派的，一会儿到后面看。
CSMain为我们定义的入口函数，其中可包含一些添加系统语义标签的参数，分别为：SV_GroupThreadID(int3)\SV_GroupIndex(int)\SV_GroupID(int3)\SV_DispatchThreadID(int3)，这四种参数的作用后面再看，他们可以任意组合，类型int、uint无所谓。
现在先说一下SV_DispatchThreadID的用处，它的意思是表示当前线程在所有分发线程组中的ID。看图：
深红色立方体，在组内的UVW坐标为（1,0,1）[注：下标从0开始]，而在全部线程坐标的XYZ中，他的坐标是（9,0,1）。
这里的XYZ坐标就是SV_DispatchThreadID，而UVW坐标是上面的SV_GroupThreadID。
回过来看程序，index后面那一大串，其实就是把三维坐标转换成一维坐标，例如深红色立方体，它在图中的index就应该是49。
我们可以再看其他参数，SV_GroupIndex，它的用法是，当前线程在组内的索引值，例如图中的深红色方块是5。
SV_GroupID为当前线程组在总共方程组中的ID，图中的深红色方块，以及它所在线程组的全部线程的值都为(4,0,0)
总之，程序将从0~15索引的数值相加，并存入了Result中，接下来我们来到C#中写程序。
新建一个脚本如下所示：
//VectorAdd.cs
public class VectorAdd : MonoBehaviour
{
public ComputeShader vectorAdd;//计算着色器
private ComputeBuffer preBuffer;//GPU缓冲区1
private ComputeBuffer nextBuffer;//GPU缓冲区2
private ComputeBuffer resultBuffer;//GPU缓冲区3
public Vector3[] array1;//CPU内存1
public Vector3[] array2;//CPU内存2
public Vector3[] resultArr;//CPU内存3
public int length = 16;
private int kernel;//计算着色器句柄
void Start()
{
//初始化CPU内存
array1 = new Vector3[length];
array2 = new Vector3[length];
resultArr = new Vector3[length];
// 填入数字
for(int i = 0; i < length; i++)
{
array1[i] = new Vector3(i, i, i);
array2[i] = Vector3.one;
}
//初始化GPU Buffer空间
InitBuffers();
//得到句柄，根据变量名称和缓冲区引用设置缓冲区
kernel = vectorAdd.FindKernel("CSMain");
vectorAdd.SetBuffer(kernel, "preVertices", preBuffer);
vectorAdd.SetBuffer(kernel, "nextVertices", nextBuffer);
vectorAdd.SetBuffer(kernel, "Result", resultBuffer);
}
void InitBuffers()
{
//传入数据个数、步长
preBuffer = new ComputeBuffer(array1.Length, 12);
preBuffer.SetData(array1);
nextBuffer = new ComputeBuffer(array2.Length, 12);
nextBuffer.SetData(array2);
resultBuffer = new ComputeBuffer(resultArr.Length, 12);
}
void Update()
{
//如果按下A键，分派线程组，将得到的数据放回CPU的resultArr中
if (Input.GetKeyDown(KeyCode.A))
{
vectorAdd.Dispatch(kernel, 2, 2, 1);
resultBuffer.GetData(resultArr);
}
}
private void OnDestroy()
{
resultBuffer.Release();
preBuffer.Release();
nextBuffer.Release();
}
}
完成了向量加法。
实验二：纹理操作
我们的目的当然不止是做向量加法，现在试着操控一下纹理。
在这一节，我希望能传入一张纹理，一个颜色，然后将颜色与纹理叠加，并实时传出纹理，将纹理塞入某个贴在Plane上的材质球中，以供实时显示。
既然要输出，那就先建立一个普通的shader，只需要根据纹理采样显示即可：
Shader "Custom/outputShader"
{
Properties
{
_MainTex("Output Tex", 2D) = "white" {}
}
SubShader
{
Tags { "RenderType" = "Opaque" }
LOD 200
Pass{
CGPROGRAM
#pragma vertex vs
#pragma fragment ps
sampler2D _MainTex;
struct a2v {
float4 vertex : POSITION;
float2 texcoord : TEXCOORD0;
};
struct v2f {
float4 vertex : SV_POSITION;
float2 texcoord : TEXCOORD1;
};
v2f vs(a2v i) {
v2f o;
o.vertex = UnityObjectToClipPos(i.vertex);
o.texcoord = i.texcoord;
return o;
}
fixed4 ps(v2f i) : SV_Target
{
return tex2D(_MainTex, i.texcoord);
}
ENDCG
}
}
FallBack "Diffuse"
}
新建一个材质球，将材质球的shader设为当前的shader。
新建一个compute shader，一个C#脚本，计算着色器代码如下：
#pragma kernel CSMain
cbuffer cbColor
{
float4 mutiColor;
};
Texture2D Input;
RWTexture2D<float4> gOutput;
[numthreads(8,8,1)]
void CSMain (uint3 groupThreadID :SV_GroupThreadID, uint3 dispatchThreadID : SV_DispatchThreadID)
{
gOutput[dispatchThreadID.xy] = Input[dispatchThreadID.xy] * mutiColor;
}
一个只读的输入材质，一个输出的读写材质。
cbuffer是常量缓冲区，是所有核心的通用数据，学过DirectX或OpenGL的一定都不陌生。
结果就是将输入的值乘颜色分量后交给输出纹理。
接下来是C#代码：
public class TextureAdd : MonoBehaviour
{
public ComputeShader shader;
public GameObject outputMaterial;
public Texture2D inputTex;
public Color mutiColor;
private RenderTexture rt;
private int kernel;
void Start()
{
kernel = shader.FindKernel("CSMain");
rt = new RenderTexture(inputTex.width, inputTex.height, 4);
rt.enableRandomWrite = true;
rt.Create();
shader.SetTexture(kernel, "Input", inputTex);
shader.SetTexture(kernel, "gOutput", rt);
}
void Update()
{
shader.SetVector("mutiColor", new Vector4(mutiColor.r, mutiColor.g, mutiColor.b, mutiColor.a));
shader.Dispatch(kernel, inputTex.width / 8, inputTex.height / 8, 1);
outputMaterial.GetComponent<MeshRenderer>().sharedMaterial.SetTexture("_MainTex", rt);
}
private void OnDestroy()
{
rt.Release();
}
}
在这里，我们用RenderTexture新建了一个纹理，并在Start中设置为可读写并创建。
每帧Update更新，都将mutiColor提交给计算着色器计算颜色，并将混合的纹理返回，然后将其塞入材质的纹理中显示。
我们挂上脚本，填入计算着色器，输出材质为Plane本身，纹理随意添加一张，点击运行，会发现材质球已经被填入了纹理，此时可以动态改变混合颜色，视窗中的纹理颜色也会跟着实时变化。
实验三：高斯模糊
高斯模糊的原理以及DirectX版的compute shader实现可参照：DirectX11 使用计算着色器实现高斯模糊。实际上，就是用一个总和为1，分布为高斯分布的矩阵（卷积核）去卷积图像。
就是类似上面的操作，不过卷积核一般要和为一（右上角），大于一则图片变量，反之变暗。很多图像处理都用到了卷积操作，不过高斯模糊有一点比较特别，他可以用一维向量，横竖分别对图像做点乘(dot)操作，效果和矩阵卷积一样。
这样的好处是减少了计算次数，原来如果卷积核大小为3x3，那么需要算9次乘法，而横竖分别dot只需要6次，如果卷积核更大，则省下的计算次数会更多。
所以我们准备大小为N的线程组若干个，两个kernel分别对图像做横竖模糊：
图中假设我们的图片大小为8x7，权重向量（代替卷积核的一维向量）大小为2*半径+1，每个线程组包含7个线程，水平模糊需要7*2=14个线程组，垂直模糊需要1*8个线程组。水平模糊时，像素大小不是线程组的整数倍，超出部分也要会进行计算，对超出部分取纹理不是一个好的行为，我记得在DirectX中，这样的行为会返回0；但我们最好用代码注意这部分因为我们希望超出边界的颜色取值为边界颜色（Clamp模式）。
另外一点，当我们关注一个线程组内所有的线程计算时我们思考一个情况：当前线程需要获取自己的像素颜色，自己前两个颜色，后两个颜色；当每个线程都这样，就会出现下面这种状况：
每个线程都用了好几倍的空间，并且取得好多次颜色，这样再空间、时间上的效率都不是最优的，那么和不采取下面这种策略：
申请一块线程组内的公共空间，大小为线程组大小+两个半径长度；首先每个线程都把自己的颜色填入公共空间的中间区域，同时边缘半径数量的线程要计算出索引，同时要将索引控制在图像可取颜色范围，然后将颜色填入公共空间。这样再空间与时间上都能得到很大提升。
好了，来看看代码，其实直接把DirectX范例教程的hlsl贴过来再改一改好使了：
#pragma kernel HorzBlurCS
#pragma kernel VertBlurCS
cbuffer cbSettings
{
//最大模糊半径为5
float gWeights[11];
};
static const int gBlurRadius = 5;
Texture2D gInput;
RWTexture2D<float4> gOutput;
#define N 256
#define CacheSize (N + 2*gBlurRadius)
groupshared float4 gCache[CacheSize];//公共空间
//水平模糊
[numthreads(N, 1, 1)]
void HorzBlurCS(int3 groupThreadID : SV_GroupThreadID,
int3 dispatchThreadID : SV_DispatchThreadID)
{
//当前线程是否在线程组中处于左边界，是的话要负责填充公共空间像素
if (groupThreadID.x < gBlurRadius)
{
//找到左侧像素的索引并限制大于0
int x = max(dispatchThreadID.x - gBlurRadius, 0);
gCache[groupThreadID.x] = gInput[int2(x, dispatchThreadID.y)];
}
//是否处于右边界
if (groupThreadID.x >= N - gBlurRadius)
{
//找到右侧像素索引并限制小于图片长度
int x = min(dispatchThreadID.x + gBlurRadius, gInput.Length.x - 1);
gCache[groupThreadID.x + 2 * gBlurRadius] = gInput[int2(x, dispatchThreadID.y)];
}
//将自己的颜色填充进索引，如果超出图像范围，则用边界的颜色
gCache[groupThreadID.x + gBlurRadius] = gInput[min(dispatchThreadID.xy, gInput.Length.xy - 1)];
// 等待线程组内所有线程同步
GroupMemoryBarrierWithGroupSync();
//
// 现在模糊每个像素
//
float4 blurColor = float4(0, 0, 0, 0);
[unroll]
for (int i = -gBlurRadius; i <= gBlurRadius; ++i)
{
int k = groupThreadID.x + gBlurRadius + i;
blurColor += gWeights[i + gBlurRadius] * gCache[k];
}
gOutput[dispatchThreadID.xy] = blurColor;
}
//垂直模糊
[numthreads(1, N, 1)]
void VertBlurCS(int3 groupThreadID : SV_GroupThreadID,
int3 dispatchThreadID : SV_DispatchThreadID)
{
if (groupThreadID.y < gBlurRadius)
{
int y = max(dispatchThreadID.y - gBlurRadius, 0);
gCache[groupThreadID.y] = gInput[int2(dispatchThreadID.x, y)];
}
if (groupThreadID.y >= N - gBlurRadius)
{
int y = min(dispatchThreadID.y + gBlurRadius, gInput.Length.y - 1);
gCache[groupThreadID.y + 2 * gBlurRadius] = gInput[int2(dispatchThreadID.x, y)];
}
gCache[groupThreadID.y + gBlurRadius] = gInput[min(dispatchThreadID.xy, gInput.Length.xy - 1)];
GroupMemoryBarrierWithGroupSync();
float4 blurColor = float4(0, 0, 0, 0);
[unroll]
for (int i = -gBlurRadius; i <= gBlurRadius; ++i)
{
int k = groupThreadID.y + gBlurRadius + i;
blurColor += gWeights[i + gBlurRadius] * gCache[k];
}
gOutput[dispatchThreadID.xy] = blurColor;
}
groupshared就是申请共有空间的关键词，代码只要看懂了上面的图就能看懂。
接下来是脚本的编写，我还是碰到一些坑。
public class GaussianBlur : MonoBehaviour
{
public ComputeShader shader;
private RenderTexture tmpTex;
private RenderTexture outTex;
private int kernel_1;
private int kernel_2;
private float[] gWeights;
void Start()
{
kernel_1 = shader.FindKernel("HorzBlurCS");
kernel_2 = shader.FindKernel("VertBlurCS");
InitWeight(2.5f);
shader.SetFloats("gWeights", gWeights);
}
private void InitWeight(float sigma)
{
float twoSigma2 = 2.0f * sigma * sigma;
int blurRadius = (int)Mathf.Ceil(2.0f * sigma);
gWeights = new float[(blurRadius * 2 + 1) * 4];
float weightSum = 0.0f;
for (int i = -blurRadius; i <= blurRadius; ++i)
{
gWeights[(i + blurRadius)*4] = Mathf.Exp(-i * i / twoSigma2);
weightSum += gWeights[(i + blurRadius)*4];
}
for (int i = 0; i < gWeights.Length; i+=4)
{
gWeights[i] /= weightSum;
}
}
void OnRenderImage(RenderTexture src, RenderTexture dest)
{
if(tmpTex == null)
{
tmpTex = new RenderTexture(src.width, src.height, 4);
tmpTex.enableRandomWrite = true;
tmpTex.Create();
outTex = new RenderTexture(src.width, src.height, 4);
outTex.enableRandomWrite = true;
outTex.Create();
}else if(tmpTex.width != src.width || tmpTex.height != src.height)
{
tmpTex.Release();
outTex.Release();
tmpTex = new RenderTexture(src.width, src.height, 4);
tmpTex.enableRandomWrite = true;
tmpTex.Create();
outTex = new RenderTexture(src.width, src.height, 4);
outTex.enableRandomWrite = true;
outTex.Create();
}
shader.SetTexture(kernel_1, "gInput", src);
shader.SetTexture(kernel_1, "gOutput", tmpTex);
shader.SetTexture(kernel_2, "gInput", tmpTex);
shader.SetTexture(kernel_2, "gOutput", outTex);
shader.Dispatch(kernel_1, (int)Mathf.Ceil(src.width / 256.0f), src.height, 1);
shader.Dispatch(kernel_2, src.width, (int)Mathf.Ceil(src.height / 256.0f), 1);
Graphics.Blit(outTex, dest);
}
private void OnDestroy()
{
if(tmpTex != null)
{
tmpTex.Release();
outTex.Release();
}
}
}
我希望这个脚本能挂在相机上，然后给个compute shader就能工作，所以只暴露出一个shader（甚至不暴露更好），为什么要两个RenderTexture一会儿说，分别对水平垂直模糊的两个方法，所以需要两个kernel，权重数组gWeights；
在Start中初始化kernel，然后计算权重，并塞给shader中常量缓冲区的gWeights。
计算方法平平无奇，不过申请空间大小时(blurRadius * 2 + 1) * 4可能会让人奇怪，假如半径为5，那么权重向量大小为11就好了，为什么要再乘4。
是因为我后面用的shader.SetFloats方法，我之前大小为11，模糊后发现图像变得很暗，CPU端权重向量没问题，所以用读写结构化缓冲区传出，发现只得到了第0、4、8的值，也就是说他给我4个数一传，加起来才等于1，才传了三个数字，肯定会变暗，而且还横竖模糊了两次，不认真看还以为是黑屏。后来看Unity文档以及一个博客ComputeShader.SetFloats()发现，hlsl中，数组元素都会存储在float4分量中，也就是说，看起来float[11]大小为4x11，其实是16x11，怪不得DirectX用的是这种排列方式：
cbuffer cbSettings : register(b0)
{
// We cannot have an array entry in a constant buffer that gets mapped onto
// root constants, so list each element.
int gBlurRadius;
// Support up to 11 blur weights.
float w0;
float w1;
float w2;
float w3;
float w4;
float w5;
float w6;
float w7;
float w8;
float w9;
float w10;
};
我们为了省事，直接将数组扩大四倍，并分别存入0、4、8、12……这样的四整数倍索引位置中，直接SetFloats就好。
我们希望能在每次绘制，将图片给我们，然后经过后处理并返还，void OnRenderImage(RenderTexture src, RenderTexture dest)方法正是这样的函数。
我希望能申请一个中间纹理tmpTex，将src水平模糊后给tmpTex，将tmpTex垂直模糊后给dest，似乎完美，但得到的结果是黑屏，然后发现dest本身是null，没有空间，因此不能给它存储值，所以又申请了一个outTex，垂直模糊后交给它并用Graphics.Blit(outTex, dest);将纹理交给dest。
这样，将脚本挂在Camera上，附加上刚刚写好的计算着色器，运行：
你成为鳌拜的弟弟熬夜后看见的图
只要反复启用/禁用脚本，很容易就能发现区别。
……
什么？封面？只要加了N个后处理，然后加上这个模糊，就能把下面的场景变成封面的样子：
这玩意做恐怖游戏一定很强。
溜了溜了。
