补充篇:Unity中Compute Shader的基本使用

补充篇:Unity中Compute Shader的基本使用

Compute Shader 可以充分利用GPU来帮助我们处理大规模的并行任务。虽然名字带Shader,但它可不光用于图形学,所以即便对渲染相关的知识不甚了解,也不妨碍学习它的用法。

基本流程

对任意 Project的文件夹右键Create/Shader/Compute Shader即可创建一个 Compute Shader,我们给它取名「MatrixCompute」,我们将其中的默认代码全部删除。

我们从一个故事开始:在17世纪,C国的现任国王拥有无上智慧,现面临外敌入侵,他从异国雇佣了1000个神奇士兵,以保得山河无恙。

这些神奇的士兵有这么一些能力:

  • 他们可以上天入地,因此可以结成各种阵型;
  • 他们会做同样的事情;
  • 他们喜欢结队而行;

现在,他开始思考战事:

指令(内核的主函数)

国王在一天中会下达很多指令,但鉴于士兵的特殊性,他会专门标出士兵们要执行的指令。

Compute Shader 也是如此,#pragma kernel XXX(作为内核主函数的函数名字) 来标记将在多线程中执行的函数,同时应当写一个名字相同的函数与之匹配,返回值要为 void ,入参暂时不管:

glsl
#pragma kernel OrdersToSoliders

void 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 的后面三个参数,就是以工作组为单位在组建好大阵型并运行。

csharp
public void Dispatch (int kernelIndex, int threadGroupsX, int threadGroupsY, int threadGroupsZ)

而确定具体单个线程的坐标可由带了 : SV_DispatchThreadIDuint3 类型参数获取,这种参数不需要我们传入,会自动提供。

总结一下:Compute Shader的划分与C国的排兵布阵很类似,每个线程就像一个士兵,numthreads(a, b, c) 规定了小队阵型,Dispatch 再以小队为单位构建了总体的大阵型,而 uint3 id : SV_DispatchThreadID 的这个 id 就是每个士兵在该阵型中独一无二的标识(其实参数名也可以是别的,只不过一定得后面跟着: SV_DispatchThreadID)。

现在为作为“指令”的函数,加上numthreads,规模暂定为(8, 8, 1),保存后报错就消失了;考虑到区分各线程还是挺有必要的,我们也传入 uint3 id : SV_DispatchThreadID

glsl
#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声明内存的,通过 SetDataGetData 的方式来写读数据(后面实践会讲到)。

有一个值得注意的是自定义类型Compute Shader中我们可以定义结构体,但注意其内部变量的顺序要与CPU的一致,不然读取出来的数据顺序会不一样,例如下面这样就是极有可能导致错误的,因为在数据搬运的过程中是不会看变量名的,只是按内存分布:

csharp
struct PersonCS
{
    int age;
    float weight;
    float height;
}
glsl
struct PersonCompute
{
    int age;
    float height;
    float weight;
}
//这类数据从GPU传回CPU的PersonCS类变量后,PersonCS类变量的weight和height逻辑上就反了

示例:用Compute Shader进行矩阵相乘

前提知识(如果你知道矩阵乘法,可以跳过)

矩阵是一个非常神奇的东西,它既可以参与神经网络的计算,也可以解方程、位姿变化……但今天我们只说说矩阵乘法。

矩阵从组成上看,就是一个数字组成的方块,类似这样:

[123456789]

上面这个矩阵,行数为3,列数也为3,但行数和列数可以不相等。在进行矩阵乘法时,必须满足这样的条件:前者的列数必须等于后者的行数,例如:

[123456][123]=[1432]

前者行数为2、列数为3;后者行数为3,列数为1,得到的结果矩阵行数为2,列数为1……等等,这是怎么算出来的,为什么结果矩阵是这样的行列数目呢?可以用下面这张图说明计算过程:

结果矩阵的行数就等于前者的行数,列数就等于后者的列数。而具体的数值,就是前者的行中的元素(从左到右)分别与后者的列中的元素(从上到下)相乘后再求和。如果像知道自己是否已经理解,可以试试算出这个的结果矩阵:

[122036][123456]=

揭晓答案:

Compute Shader 代码

可以发现,其实矩阵乘法中结果矩阵的每个元素的计算都不容易呢(如果计算的两个矩阵都很大的话),用普通方法实现的矩阵乘法,假设为 a 矩阵 乘 b 矩阵,得用三层循环,一层遍历a中各行的元素,一层遍历b中各列的元素,剩下一层设置得到的结果在新得到的矩阵中的位置:

csharp
//结果矩阵的行数就等于前者的行数,列数就等于后者的列数
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吧:

glsl
#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个数就好了。

glsl
#pragma kernel OrdersToSoliders
#pragma kernel MyMatrixFunc

RWStructuredBuffer<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中的索引,为此,还要再用「第几个」中的这个「几」减一,即:

[(a+1)1]+(b+1)1a+b

能理解这些的话,MyMatrixFunc 函数就不难写了:

glsl
[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的类:

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

public class GPUMatrixCompute : MonoBehaviour
{

}

首先肯定要一个Compute Shader类变量用于获取刚刚写的Compute Shader(public,方便编辑器页面拖入),然后还要有3个ComputeBuffer类变量,分别与之前「MatrixCompute」中的三个存储矩阵数据的RWStructuredBuffer<float>对应。没错,我们不能直接把普通的数据传入GPU,而是使用ComputeBufferStructuredBuffer(只读),这是GPU和CPU之间传递数据的主要方式,最后是同样3个记录矩阵维度的M、K、N:

csharp
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()得到指定类型的大小);ComputeBufferSetData,我们可以将容器里(不只是数组类型)的数据传入ComputeBuffer中:

csharp
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 + 赋值的变量

csharp
//——————————————————第一种——————————————————   
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,按照声明顺序,分别有自己的kernelIndexOrdersToSoliders是0,MyMatrixFunc是1,如果有更多,就依次往后2、3、4……

glsl
#pragma kernel OrdersToSoliders
#pragma kernel MyMatrixFunc
……

我们的确可以就这样传入SetBuffer(要用MyMatrixFunc,所以取1):

csharp
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中它们的顺序,岂不是会出错 (而且难以发现。所幸,还可以用名字来指定:

csharp
int martixKernel = computeShader.FindKernel("MyMatrixFunc");

整理下到目前为止的代码:

csharp
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函数指定内核主函数与工作组数的尺寸(大阵型),对于这次的矩阵计算,显然一组就绰绰有余了:

csharp
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],我们就可以:

    csharp
    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方法即可,但也要足够大小的容器来收纳:

csharp
float[] matrixOut = new float[M * N]; 
matrixOutBuffer.GetData(matrixOut);

最后的最后,还要记得使用完后,释放掉在GPU中Buffer占用的内存:

csharp
matrixABuffer.Release();
matrixBBuffer.Release();
matrixOutBuffer.Release();

最终代码如下,为了直观看到计算结果,将matrixOut数组全局公开变量,在编辑器中指定大小,也可以直接在编辑器中看到结果:

csharp
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之间的传输所消耗的时间远大于计算。使用时也应当避免频繁的读取与写入,即SetDataGetData的使用(SetBuffer倒是不要紧,它只涉及 GPU 资源的绑定,而不涉及数据传输)。

还有一点需要注意的便是numthreadsDispatch尺寸的设置都需与执行任务匹配,这样才能更好地利用 uint3 id: SV_DispatchThreadID,其实除了SV_DispatchThreadID,还有一些其他的系统值参数,也可以在需要时传入用于辅助函数,像SV_GroupIDSV_GroupIndex等,可参考 https://learn.microsoft.com/zh-cn/windows/win32/direct3dglsl/dx-graphics-glsl-semantics

这篇文章如果能帮到你,那再好不过。( ̄▽ ̄)


__EOF__

  • 本文作者: FoxTiger
  • 本文链接: https://www.cnblogs.com/OwlCat/p/18684161
  • 关于博主: 评论和私信会在第一时间回复。或者直接私信我。
  • 版权声明: 本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!
  • 声援博主: 如果您觉得文章对您有帮助,可以点击文章右下角推荐一下。
  • posted @   狐王驾虎  阅读(121)  评论(0编辑  收藏  举报
    相关博文:
    阅读排行:
    · [翻译] 为什么 Tracebit 用 C# 开发
    · 腾讯ima接入deepseek-r1,借用别人脑子用用成真了~
    · Deepseek官网太卡,教你白嫖阿里云的Deepseek-R1满血版
    · DeepSeek崛起:程序员“饭碗”被抢,还是职业进化新起点?
    · 深度对比:PostgreSQL 和 SQL Server 在统计信息维护中的关键差异
    点击右上角即可分享
    微信分享提示