虚幻引擎 自定义ComputeShader

深海迷航

引言

如何在虚幻引擎中编写Compute Shader，在虚幻中可能不是那么的直观。那么这篇文章就带领大家从头开始创建一个Compute Shader，并且使用这个Compute Shader来读取硬盘上的权重贴图，结合贴图进行输出。本文基于4.26.2，往上的版本应该都是支持的。本文基于我自己的理解和参考资料，欢迎大家指正和讨论。
一、什么是Compute Shader

A Compute Shader is a Shader Stage that is used entirely for computing arbitrary information. While it can do rendering, it is generally used for tasks not directly related to drawing triangles and pixels.

上面是OpenGL Wiki对Compute Shader的定义：Compute Shader - OpenGL Wiki。Compute Shader就是一种多用途的Shader，输入不是直接从三角形和像素来的，而是我们在Shader程序里规定Input是什么类型，之后我们在编写的时候就可以看到了。虚幻中有很多使用Compute Shader的地方，例如虚幻里的EditLayer，Water插件的Gerstner Waves，还有Niagara的GPU Particle等等。
如果不清楚Compute Shader，建议看Reference的第4篇，那是一篇非常详细的Compute Shader的基础介绍。
在虚幻引擎里面可以看到很多Compute Shader，我们直接在IDE中搜素 "SF_COMPUTE); "，就可以看到很多种类繁多的Compute Shader了。


我之前参考的Shader就是GpuSkinCacheComputeShader，大家也可以看虚幻内置的其他Compute Shader进行学习。

二、创建一个基础的Compute Shader

UE端定义

我们新建一个插件，我这里就叫Compute Shader。然后我们在插件里新建一个Shader路径，这个就是在插件同名的CPP文件里的StartupModule里创建。





对应的目录结构

我们尽量把Shader路径映射到插件的路径下面，这样给别人用的时候就很方便了。尽量不要映射到工程和引擎里面去。而且注意这个虚拟路径"/CustomShaders"是不能和别的虚拟路径重名的，不然引擎编译的时候会报错。
创建好之后我们新建一个头文件和一个cpp。我们直接进入主题，来看我们的Compute Shader的结构。


我们来一步步拆解这个Shader。那么首先我们看到第一行有一个


这个就是方便我们修改每一个组里面有多少个线程。具体含义我们可以看这张很经典的图来理解。



来源：https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/sv-dispatchthreadid


然后我们就正式来定义我们的Shader类，我们的Shader都是继承自FGlobalShader的。所以直接这么写就可以。之后我们就写上两个宏。


DECLARE_GLOBAL_SHADER用于定义我们的shader类型。SHADER_USE_PARAMETER_STRUCT则是告诉引擎我们需要使用的Parameter Struct是什么。下面我们就定义我们的shader参数。


可以看到这里我们定义了一个输入SRV(Shader Resource View)和一个输出UAV(Unordered Access View)。这个地方的名字需要和HLSL文件中的名字对应上。


最后我们将前面Define的每组线程数输入到Shader的文件中即可。通过修改Shader的编译环境的方式，我们就可以通过不修改HLSL Code的方式来修改线程数。


最后我们需要在类定义的外面写明我们需要创建这个Shader，并且需要告诉引擎我们的Shader文件的位置，主函数入口以及Shader 类型。
Shader文件

先在我们创建好的Shader路径里新建一个和我们最后一句宏里定义一样的Shader文件。首先我们先将我们刚才在Shader Parameter Struct里定义的参数写在Shader的头部，注意名字需要和我们在类里面定义的一致。


之后我们需要写一个和我们刚才在宏里定义的名字一样的函数。


这个函数里面做的事情其实就很简单了，大家一看就能看懂。就是把InputTexture的数据输出到OutputTexture上。我们可以看到我使用了一个DTid来访问Texture，其实这个DTid就是一个全局id，结合上面的图我们就能看懂了。注意我们需要在函数上面写上一组需要多少个线程，如果这里不是从#define里取得，那么我们就需要手动填写一下。例如我随便从虚幻引擎中找了一个CS。


现在我们就写好了一个Compute Shader了。接下来我们需要给Shader参数绑定资源，以及调用这个Shader。
三、调用Shader

我们新建一个Component来调用这个Shader。


首先我们定义一下这个Component是继承自ActorComponent。


根据我们刚才组织的参数结构，我们需要一个Input和一个Output。Input就是一个Texture2D了。


Output我们可以指定一个RenderTarget2D，或者一张Texture都行。这里我们指定一个RenderTarget2D。


这里我创建了一个SRVRef用于后面给参数绑定，你也可以不这么做。


Override一下Tick函数，用于Dispatch 我们的Compute Shader。


创建我们的Output，也就是这个UAVRHIRef。


这几个是一些功能函数。第一个函数是创建一个Transient的RenderTarget，是直接从Water插件里抄的。
BeginDestroy用于释放渲染资源，UpdateRenderData用于第四部分的案例更新渲染数据。
现在我们就来一一实现一下我们刚才定义的函数。


首先是构造函数，主要是要给RenderTarget赋值。这个RenderTarget也可以是引擎里创建的Asset，这里我是想省略这一步直接在代码里定义。为了在Editor中运行，我还设置了允许在Editor里面Tick的boolean值。
GetOrCreateTransientRenderTarget2D 这个函数在官方的Water插件中的WaterUtils里有，这里我就不贴代码了，因为函数实在是太长了。还是跟刚才一样，你可以直接从引擎里拖一个RenderTarget2D的资源，效果是一样的。注意无论何种方法，我们都不创建其他MipMap。所以请保证调用的时候不让RT自动产生MipMap。


同时我们将RT的bCanCreateUAV设置成true。如果是你抄的Water的代码，那么请在InitResource之前设置。
BeginDestroy这个就是在渲染线程里释放渲染资源。


UpdateRenderData我们先空着，来看TickComponent部分。
首先，当我们没有RenderTarget，也就是我们外部的输出的RT时，我们就直接创建一个Transient RT。然后我们调用UpdateRenderData来更新渲染数据。
接下来就是在渲染线程中绑定参数和调用了。我们先从GlobalShader的索引中找到我们定义的FMyComputeShader，也就是之前我们写的Compute Shader。


接下来就是给我们的ShaderParameter赋值，我们把RT对应的UAV绑在Output上，把TextureSRV绑在Input上。


最后就是Dispatch我们的Compute Shader了。


Dispatch的时候我们可以看到除去前三个参数我们一目了然，最后一个参数是组数量，即GroupCount。这里就是一共需要Issue多少个线程组，方便我们索引。例如我们现在RT的大小是1024，那么我们的GroupCount就是（32， 32， 1）。如果是512，就是（16, 16，1）。
然后我们来填充我们的UpdateRenderData函数。


里面的结构很简单，当没有InputTexture的时候，我们Load一个黑色的1x1的贴图来防止程序崩溃。然后我们创建RT的UAV和Texture的SRV。
那么至此我们已经完成了ComputeShader的搭建了，只要运行引擎，然后新建一个空Actor，给好参数，就能看到结果了。注意我们RT是1024，没有在代码里做任何逻辑的话，只有1024*1024的贴图才能占满整个RT，其他尺寸都会不匹配。起码没有值的地方都会为0。


这里你会看到这种情况。当我们设置texture时，原本这个Texture的Mip 0级是1024*1024或者更大，但是进来我们的Compute Shader里就会变得很小很小。这周我都在查这个问题，那么最终定位问题在Texture Streaming上。简单来说，Texture Streaming就是在你需要用这个Texture的时候系统才会把这个贴图加载进GPU 内存里。然而我们拷贝是没有经过虚幻引擎之手的，所以虚幻引擎没有主动帮我们加载这个贴图，而我们又要去拿这个贴图，这时虚幻会给我们一个最大的Mip层级的贴图来让程序不崩溃。这时我们就需要主动让虚幻加载一下这个贴图。具体做法如下。


在我们的UpdateRenderData里加入上图的代码，这些代码可以在虚幻引擎加载StreamableRenderAsset的地方找到，直接搜索函数名即可。这样我们就手动让虚幻加载了我们贴图的Mip 0。这样我们再运行我们的程序，就能得到正确的结果了。



运行之后看到结果



RenderTarget内容和InputTexture一致

那么我们想要修改Shader，也是不用重启引擎的。只是修改Shader逻辑的话，直接编辑我们的usf文件。编辑完成之后直接按下"Ctrl+Shift+."（英文句号），引擎就会重新编译所有修改过的usf文件，里面自然包含我们的Compute Shader。



Recompile Shaders

四、应用案例

接下来我们继续修改我们的代码，在Component里回到我们的UpdateRenderData。我随便在我的电脑上找了一张图，那么我们就用这张图作为权重贴图来将我们的InputTexture进行混合输出。
我们先在Component类的定义里加上两个新的RHIRef。一个是新权重图的SRV，一个是StructuredBuffer。


在Shader的定义里我们添加新的权重的参数。注意StructuredBuffer必须声明为StructuredBuffer，如果是Buffer或者类型错误，后面d3d会报错。


接下来我们来完善我们的UpdateRenderData函数，其中我们需要将我们磁盘上的贴图读取成一个数组。


我们重点关注这里，前面读取图片的代码我是直接搜索ImageWrapperModule这个模块获得的。


注意SBResource是uint32，uint8 d3d同样会报错，这里我就使用uint32了。这个部分是创建StructuredBuffer的部分，时间原因我的方法就比较暴力了。
在Tick里调用我们的Update方法这个就不说了。接下来最后一步就是修改我们的ComputeShader。首先我们需要在头部将我们的StructuredBuffer添加进去，之后我们在输出的地方采样我们的Weight贴图。注意采样的时候需要正确获取角标，这里因为我们定义了32个线程组，每个线程组里又有32个线程，所以一共正好和我们磁盘上的贴图尺寸一致，如果贴图尺寸不一致则需要特殊处理一下。


运行之后可以看到我们的结果被正确计算了。
五、Reference

1.https://docs.unrealengine.com/4.27/en-US/ProgrammingAndScripting/Rendering/ShaderDevelopment/AddingGlobalShaders/
2.https://medium.com/realities-io/using-compute-shaders-in-unreal-engine-4-f64bac65a907
这篇强推，基本流程都是从这里学的，需要魔法上网。
3.YivanLee：虚幻4渲染编程(Shader篇)【第七卷：虚幻4中的ComputeShader】
也是强推。
4.Compute Shader 简介
一篇详细的ComputeShader的介绍。
六、更新日志

08/07/2022 更新UAV Output和Texture Streaming修改。修改都在第三部分，欢迎大家来看最新的代码。
08/19/2022 更新了第四部分的结果图。

出售水晶

可以控制buffer生成图形嘛？

北方刀客

生成三角形？

如风如梦

请问完整源码有吗？