补充篇:Unity中Compute Shader的基本使用
Compute Shader 可以充分利用GPU来帮助我们处理大规模的并行任务。虽然名字带Shader
,但它可不光用于图形学,所以即便对渲染相关的知识不甚了解,也不妨碍学习它的用法。
基本流程
对任意 Project
的文件夹右键Create/Shader/Compute Shader
即可创建一个 Compute Shader
,我们给它取名「MatrixCompute」,我们将其中的默认代码全部删除。
我们从一个故事开始:在17世纪,C国的现任国王拥有无上智慧,现面临外敌入侵,他从异国雇佣了1000个神奇士兵,以保得山河无恙。
这些神奇的士兵有这么一些能力:
- 他们可以上天入地,因此可以结成各种阵型;
- 他们会做同样的事情;
- 他们喜欢结队而行;
现在,他开始思考战事:
指令(内核的主函数)
国王在一天中会下达很多指令,但鉴于士兵的特殊性,他会专门标出士兵们要执行的指令。
Compute Shader
也是如此,#pragma kernel XXX(作为内核主函数的函数名字)
来标记将在多线程中执行的函数,同时应当写一个名字相同的函数与之匹配,返回值要为 void
,入参暂时不管:
#pragma kernel OrdersToSolidersvoid OrdersToSoliders()
{}
保存后返回Unity,会发现一个报错,需要我们添加numthreads
,怎么做呢?我们继续看
阵型(线程组规模划分)
这次的外敌勾结了魔物,漫天凶禽、遍地走兽(和士兵)组成了块状阵列;在X尺寸中有20个单位,Y尺寸中有30个单位,Z尺寸中有5个单位,记作 (20, 30, 5)
,共 20 * 30 * 5 = 3000
个敌人,来势汹汹!
国王见这架势,将士兵按小队划分,每个小队在X尺寸中有4人、Y尺寸中有5人、Z尺寸中有1人,记作(4, 5, 1)
,共20人(下图用方块代替了);国王又以这样的小队为单位,组建了一个大阵型:在X尺寸中有5个小队单位、Y尺寸中有6个小队单位、Z尺寸中有5个小队单位,记作(5, 6, 5)
。
这样一来,大阵型就是 (4 * 5, 5 * 6, 1 * 5)
刚好等于 (20, 30, 5)
,只要出动就可以击退敌军了!但是……不是总共就1000个士兵吗!?没关系,国王自有调整的策略以补足阵型(当然,不会无中生有的)。
那如何在这次作战任务中找到具体的某个士兵呢?显然,可以通过在大阵容中(即在当前的(20, 30, 5)
范围下)士兵的位置来寻找。比如下图这个士兵,就可以用 (8, 2, 0)
来找到:
Compute Shader
多线程的调度也是如此,[numthreads(a,b,c)]
就是单个工作组的规模,调用Compute Shader
的函数 Dispatch
的后面三个参数,就是以工作组为单位在组建好大阵型并运行。
public void Dispatch (int kernelIndex, int threadGroupsX, int threadGroupsY, int threadGroupsZ)
而确定具体单个线程的坐标可由带了 : SV_DispatchThreadID
的 uint3
类型参数获取,这种参数不需要我们传入,会自动提供。
总结一下:Compute Shader
的划分与C国的排兵布阵很类似,每个线程就像一个士兵,numthreads(a, b, c)
规定了小队阵型,Dispatch
再以小队为单位构建了总体的大阵型,而 uint3 id : SV_DispatchThreadID
的这个 id
就是每个士兵在该阵型中独一无二的标识(其实参数名也可以是别的,只不过一定得后面跟着: SV_DispatchThreadID
)。
现在为作为“指令”的函数,加上numthreads
,规模暂定为(8, 8, 1)
,保存后报错就消失了;考虑到区分各线程还是挺有必要的,我们也传入 uint3 id : SV_DispatchThreadID
:
#pragma kernel OrdersToSoliders[numthreads(8, 8, 1)]
void OrdersToSoliders(uint3 id : SV_DispatchThreadID)
{}
翻译(CPU数据搬运到GPU)
战斗还未开始,准备后方支援部门的人却开始头疼了。毕竟是异国的士兵,很多用品的叫法与C国不同,无从得知C国是具备还是不具备,难以筹备。国王便开始让人做些“翻译”工作,比如士兵们称为“褡裢”的东西,可以用背包替代……
Compute Shader
也有许多与我们熟知的数据类型相似而叫法不同的东西,比如Vector3
其实与Compute Shader
中的float3
相似、List<T>
与RWStructuredBuffer<T>
相似……这里就不一一盘点了,只要知道有这么一个类似翻译
与搬运
的工作要在运行之前做就行。
ComputeBuffer
类变量就是用于在GPU声明内存的,通过 SetData
和 GetData
的方式来写读数据(后面实践会讲到)。
有一个值得注意的是自定义类型,Compute Shader
中我们可以定义结构体,但注意其内部变量的顺序要与CPU的一致,不然读取出来的数据顺序会不一样,例如下面这样就是极有可能导致错误的,因为在数据搬运的过程中是不会看变量名的,只是按内存分布:
struct PersonCS
{int age;float weight;float height;
}
struct PersonCompute
{int age;float height;float weight;
}
//这类数据从GPU传回CPU的PersonCS类变量后,PersonCS类变量的weight和height逻辑上就反了
示例:用Compute Shader进行矩阵相乘
前提知识(如果你知道矩阵乘法,可以跳过)
矩阵是一个非常神奇的东西,它既可以参与神经网络的计算,也可以解方程、位姿变化……但今天我们只说说矩阵乘法。
矩阵从组成上看,就是一个数字组成的方块,类似这样:
上面这个矩阵,行数为3,列数也为3,但行数和列数可以不相等。在进行矩阵乘法时,必须满足这样的条件:前者的列数必须等于后者的行数,例如:
前者行数为2、列数为3;后者行数为3,列数为1,得到的结果矩阵行数为2,列数为1……等等,这是怎么算出来的,为什么结果矩阵是这样的行列数目呢?可以用下面这张图说明计算过程:
结果矩阵的行数就等于前者的行数,列数就等于后者的列数。而具体的数值,就是前者的行中的元素(从左到右)分别与后者的列中的元素(从上到下)相乘后再求和。如果像知道自己是否已经理解,可以试试算出这个的结果矩阵:
揭晓答案:
Compute Shader 代码
可以发现,其实矩阵乘法中结果矩阵的每个元素的计算都不容易呢(如果计算的两个矩阵都很大的话),用普通方法实现的矩阵乘法,假设为 a 矩阵 乘 b 矩阵,得用三层循环,一层遍历a中各行的元素,一层遍历b中各列的元素,剩下一层设置得到的结果在新得到的矩阵中的位置:
//结果矩阵的行数就等于前者的行数,列数就等于后者的列数
Matrix result = new Matrix(a.Rows, b.Columns);for (int i = 0; i < a.Rows; i++)
{for (int j = 0; j < b.Columns; j++){for (int k = 0; k < a.Columns; k++){result[i, j] += a[i, k] * b[k, j];}}
}
return result;
那如果用了 Compute Shader
, 可以怎么优化呢?你应该想到了,可以 让每个线程各自计算结果矩阵中的一个元素并设置好位置,换句话说,就是让每个线程只计算上述三层循环中最内层的循环。
以刚才的矩阵为例,就可以用9个线程,分别计算每个元素的结果,一个线程算 1*1+2*4
、一个线程算 2*2+0*5
……这样分摊下来,计算所需时间就很少了。
下面就试试借助Compute Shader
吧。首先,定义好线程将要调用的主函数,就叫 MyMatrixFunc
吧:
#pragma kernel OrdersToSoliders
#pragma kernel MyMatrixFunc[numthreads(8, 8, 1)]
void MyMatrixFunc(uint3 id: SV_DispatchThreadID)
{}[numthreads(8, 8, 1)]
void OrdersToSoliders(uint3 id : SV_DispatchThreadID)
{}
PS:虽然OrdersToSoliders
没什么用,但也先留着 (后面用来当例子
我们还需要传入三个矩阵:矩阵a、矩阵b和用于输出的结果矩阵,即3个RWStructuredBuffer<float>
类变量。然而这类变量的存储是一维的,即使传入二维数组,它也还是以一维形式存储。所以我们再传入矩阵a与矩阵b的行列数,以方便定位计算出的结果的位置,但考虑到矩阵a的行数与矩阵b的列数得相同,那用3个数就好了。
#pragma kernel OrdersToSoliders
#pragma kernel MyMatrixFuncRWStructuredBuffer<float> matrixA;
RWStructuredBuffer<float> matrixB;
RWStructuredBuffer<float> matrixOut;// 矩阵维度:M为A的行数、K为A的列数/B的行数、N为B的列数
uint M, K, N;[numthreads(8, 8, 1)]
void MyMatrixFunc(uint3 id: SV_DispatchThreadID)
{}[numthreads(8, 8, 1)]
void OrdersToSoliders(uint3 id : SV_DispatchThreadID)
{}
借助 M、K、N
该怎么找到位置呢?先来看个例子:
如上图,想找到原本二维数组中 [1, 1]
的数转成在Buffer
中的位置,该怎么确定呢?
我们知道索引是从0开始的,所以 [1,1]
是在二维中的第二行第二个,这也意味着它之前首先肯定有一行数,每行元素的个数其实就是列数,又因为它在当前行的下标是1,那它之前还有一个数。所以,它应当排在Buffer
中的第四个位置。
总结一下就是:二维索引为[a, b]的数,在一维中应排在 第 [(a + 1) - 1] * 二维的列数 + (b + 1) 个。应该不难理解吧,(a + 1) 求的就是 在第几行,再 减一就是求之前有几行。
然而我们需要的是在Buffer
中的索引,为此,还要再用「第几个」中的这个「几」减一,即:
能理解这些的话,MyMatrixFunc
函数就不难写了:
[numthreads(8, 8, 1)]
void MyMatrixFunc(uint3 id: SV_DispatchThreadID)
{// 当前线程对应的结果矩阵Out的元素索引uint row = id.y; // Out的行uint col = id.x; // Out的列// 超出矩阵范围的不管if (row < M && col < N){// 计算Out[row,col]float a, b, sum = 0;for (uint k = 0; k < K; ++k){a = matrixA[row * K + k]; // A[row,k]b = matrixB[k * N + col]; // B[k,col]sum += a * b;}// 将结果写入Out[row,col]matrixOut[row * N + col] = sum;}
}
-
会有超出范围的情况吗(你可以看到在执行逻辑前先进行了一次判断)?
其实是有且常有的。就那之前的例子来说,最终结果只是一个规模为[3, 3]的矩阵,[8, 8, 1]可太足够了,但又没法只调用半个线程组或者四分之一的线程组,所以就会出现“浑水摸鱼的士兵”: -
为什么入参
id
能表示结果矩阵所在的行列?
这其实与numthreads
与后面调用时的Dispatch
密切相关,接下来就来看看C#
代码该如何调用Compute Shader
。
C#代码
先新建一个名为GPUMatrixCompute
的类:
using System.Collections;
using System.Collections.Generic;
using UnityEngine;public class GPUMatrixCompute : MonoBehaviour
{}
首先肯定要一个Compute Shader
类变量用于获取刚刚写的Compute Shader(public
,方便编辑器页面拖入),然后还要有3个ComputeBuffer
类变量,分别与之前「MatrixCompute」中的三个存储矩阵数据的RWStructuredBuffer<float>
对应。没错,我们不能直接把普通的数据传入GPU,而是使用ComputeBuffer
或StructuredBuffer
(只读),这是GPU和CPU之间传递数据的主要方式,最后是同样3个记录矩阵维度的M、K、N
:
using System.Collections;
using System.Collections.Generic;
using UnityEngine;public class GPUMatrixCompute : MonoBehaviour
{public ComputeShader computeShader;private ComputeBuffer matrixABuffer;private ComputeBuffer matrixBBuffer;private ComputeBuffer matrixOutBuffer;//用const限定,是为了方便初始化数组// 矩阵维度:M为A的行数、K为A的列数/B的行数、N为B的列数private const int M = 3, K = 2, N = 3;
}
在OnEnable
函数中,我们将这些进行初始化,ComputeBuffer
的初始化需要我们指定它的大小,我们可以把它想象成一个一维数组,它所需的两个参数,一个是数组长度,另一个是数组中每个元素的大小(可借助sizeof()
得到指定类型的大小);ComputeBuffer
的SetData
,我们可以将容器里(不只是数组类型)的数据传入ComputeBuffer
中:
private void OnEnable()
{float[,] myMatrixA = new float[M, K]{{1f, 2f},{2f, 0f},{3f, 6f},};float[,] myMatrixB = new float[K, N]{{1f, 2f, 3f},{4f, 5f, 6f},};matrixABuffer = new ComputeBuffer(M * K, sizeof(float));matrixABuffer.SetData(myMatrixA);matrixBBuffer = new ComputeBuffer(K * N, sizeof(float));matrixBBuffer.SetData(myMatrixB);matrixOutBuffer = new ComputeBuffer(M * N, sizeof(float));
}
在Start
函数中,将我们的这些变量传到Compute Shader
中吧。首先传入矩阵的3个维度,它们都是Int
变量(uint
也算),用SetInt
函数传入,它有两种传入方式,一种是Compute Shader
中对应参数的字符串名字 + 赋值的变量;另一种是一种是Compute Shader
中对应参数的ID + 赋值的变量
//——————————————————第一种——————————————————
computeShader.SetInt("M", M);
computeShader.SetInt("N", N);
computeShader.SetInt("K", K);//——————————————————第二种——————————————————
int M_ID = Shader.PropertyToID("M");
int N_ID = Shader.PropertyToID("N");
int K_ID = Shader.PropertyToID("K");computeShader.SetInt(M_ID, M);
computeShader.SetInt(N_ID, N);
computeShader.SetInt(K_ID, K);
通常来说,第二种方法会更好。如果后来在Compute Shader
中对应参数的名字修改了,那么第二种方法只要改Shader.PropertyToID()
传入的参数就行了,而第一种方法却需要修改所有用到的地方。
那ComputeBuffer
要怎么传入呢?有个SetBuffer
方法可以做到,但相比SetInt
,它还需要传入kernelIndex
,用于指定Buffer
被用在哪个内核主函数中。还记得吗?我们的MatrixCompute
中现在有两个kernel
,按照声明顺序,分别有自己的kernelIndex
,OrdersToSoliders
是0,MyMatrixFunc
是1,如果有更多,就依次往后2、3、4……
#pragma kernel OrdersToSoliders
#pragma kernel MyMatrixFunc
……
我们的确可以就这样传入SetBuffer
(要用MyMatrixFunc
,所以取1):
int martixKernel = 1;int matrixA_ID = Shader.PropertyToID("matrixA");
int matrixB_ID = Shader.PropertyToID("matrixB");
int matrixOut_ID = Shader.PropertyToID("matrixOut");computeShader.SetBuffer(martixKernel, matrixA_ID, matrixABuffer);
computeShader.SetBuffer(martixKernel, matrixB_ID, matrixBBuffer);
computeShader.SetBuffer(martixKernel, matrixOut_ID, matrixOutBuffer);
可万一“不小心”调整了Compute Shader
中它们的顺序,岂不是会出错 (而且难以发现。所幸,还可以用名字来指定:
int martixKernel = computeShader.FindKernel("MyMatrixFunc");
整理下到目前为止的代码:
public class GPUMatrixCompute : MonoBehaviour
{public ComputeShader computeShader;private ComputeBuffer matrixABuffer;private ComputeBuffer matrixBBuffer;private ComputeBuffer matrixOutBuffer;private const int M = 3, N = 3, K = 2;//运行时,Compute Shader中的各个变量名不会变,因此可以将获取的ID作静态只读修饰(不加当然也可以)private static readonly int matrixA_ID = Shader.PropertyToID("matrixA");private static readonly int matrixB_ID = Shader.PropertyToID("matrixB");private static readonly int matrixOut_ID = Shader.PropertyToID("matrixOut");private static readonly int M_ID = Shader.PropertyToID("M");private static readonly int N_ID = Shader.PropertyToID("N");private static readonly int K_ID = Shader.PropertyToID("K");……private void Start(){int martixKernel = computeShader.FindKernel("MyMatrixFunc");computeShader.SetInt(M_ID, M);computeShader.SetInt(N_ID, N);computeShader.SetInt(K_ID, K);computeShader.SetBuffer(martixKernel, matrixA_ID, matrixABuffer);computeShader.SetBuffer(martixKernel, matrixB_ID, matrixBBuffer);computeShader.SetBuffer(martixKernel, matrixOut_ID, matrixOutBuffer);}
}
终于要到调用了!在一切数据都写进Compute Shader
后,Dispatch
函数指定内核主函数与工作组数的尺寸(大阵型),对于这次的矩阵计算,显然一组就绰绰有余了:
private void Start()
{int martixKernel = computeShader.FindKernel("MyMatrixFunc");computeShader.SetInt(M_ID, M);computeShader.SetInt(N_ID, N);computeShader.SetInt(K_ID, K);computeShader.SetBuffer(martixKernel, matrixA_ID, matrixABuffer);computeShader.SetBuffer(martixKernel, matrixB_ID, matrixBBuffer);computeShader.SetBuffer(martixKernel, matrixOut_ID, matrixOutBuffer);computeShader.Dispatch(martixKernel, 1, 1, 1);
}
-
但遇见比较庞大的工作任务时又该怎么指定规模呢?
其实可以通过计算获得合适的尺寸,比如一个矩阵的输出规模是[64, 45]
,而我们每个线程组的大小还是[8, 8, 1]
,我们就可以:uint x, y, z; //获取指定内核函数的numthreads大小(即小队规模) computeShader.GetKernelThreadGroupSizes(martixKernel, out x, out y, out z);//假设输出矩阵的尺寸仍是[M,N] //向上取整,向上取整可能会有部分浪费,但向下取整就完成不了任务 int groupX = Mathf.CeilToInt(M / x); int groupY = Mathf.CeilToInt(N / y); //z尺寸就不用算了computeShader.Dispatch(martixKernel, x, y, z);
现在,计算完的结果矩阵,还是在matrixOutBuffer
中,怎么把结果拿回到数组中呢?用GetData
方法即可,但也要足够大小的容器来收纳:
float[] matrixOut = new float[M * N];
matrixOutBuffer.GetData(matrixOut);
最后的最后,还要记得使用完后,释放掉在GPU中Buffer
占用的内存:
matrixABuffer.Release();
matrixBBuffer.Release();
matrixOutBuffer.Release();
最终代码如下,为了直观看到计算结果,将matrixOut
数组全局公开变量,在编辑器中指定大小,也可以直接在编辑器中看到结果:
using UnityEngine;public class GPUMatrixCompute : MonoBehaviour
{public ComputeShader computeShader;public float[] matrixOut;private ComputeBuffer matrixABuffer;private ComputeBuffer matrixBBuffer;private ComputeBuffer matrixOutBuffer;private const int M = 3, N = 3, K = 2;private static readonly int matrixA_ID = Shader.PropertyToID("matrixA");private static readonly int matrixB_ID = Shader.PropertyToID("matrixB");private static readonly int matrixOut_ID = Shader.PropertyToID("matrixOut");private static readonly int M_ID = Shader.PropertyToID("M");private static readonly int N_ID = Shader.PropertyToID("N");private static readonly int K_ID = Shader.PropertyToID("K");private void OnEnable() {float[,] myMatrixA = new float[M, K]{{1f, 2f},{2f, 0f},{3f, 6f},};float[,] myMatrixB = new float[K, N]{{1f, 2f, 3f},{4f, 5f, 6f},};matrixABuffer = new ComputeBuffer(M * K, sizeof(float));matrixABuffer.SetData(myMatrixA);matrixBBuffer = new ComputeBuffer(K * N, sizeof(float));matrixBBuffer.SetData(myMatrixB);matrixOutBuffer = new ComputeBuffer(M * N, sizeof(float));}private void Start(){int martixKernel = computeShader.FindKernel("MyMatrixFunc");computeShader.SetInt(M_ID, M);computeShader.SetInt(N_ID, N);computeShader.SetInt(K_ID, K);computeShader.SetBuffer(martixKernel, matrixA_ID, matrixABuffer);computeShader.SetBuffer(martixKernel, matrixB_ID, matrixBBuffer);computeShader.SetBuffer(martixKernel, matrixOut_ID, matrixOutBuffer);computeShader.Dispatch(martixKernel, 1, 1, 1);matrixOut = new float[M * N];matrixOutBuffer.GetData(matrixOut);matrixABuffer.Release();matrixBBuffer.Release();matrixOutBuffer.Release();}
}
尾声
现在回到编辑器中,将脚本拖入任意物体中,并把ComputeShader
也赋值上,点击运行就能看到MatrixOut
所显示的输出结果了:
其实对于小规模的这类运算,用Compute Shader
并不划算,因为这时数据在CPU与GPU之间的传输所消耗的时间远大于计算。使用时也应当避免频繁的读取与写入,即SetData
和GetData
的使用(SetBuffer
倒是不要紧,它只涉及 GPU 资源的绑定,而不涉及数据传输)。
还有一点需要注意的便是numthreads
和Dispatch
尺寸的设置都需与执行任务匹配,这样才能更好地利用 uint3 id: SV_DispatchThreadID
,其实除了SV_DispatchThreadID
,还有一些其他的系统值参数,也可以在需要时传入用于辅助函数,像SV_GroupID
、SV_GroupIndex
等,可参考 https://learn.microsoft.com/zh-cn/windows/win32/direct3dhlsl/dx-graphics-hlsl-semantics
这篇文章如果能帮到你,那再好不过。( ̄▽ ̄)