UE5 Compute Shader 入门

Compute Shader

添加Shader文件映射
创建Shader，继承自FGlobalShader
使用RDG创建资源

代码

TaSimpleColorCS.usf

#include "/Engine/Public/Platform.ush"
uint2 TextureSize;
RWTexture2D<float4> OutputTexture;

[numthreads(THREADGROUPSIZE_X, THREADGROUPSIZE_Y, THREADGROUPSIZE_Z)]
void MainCS(uint3 id : SV_DispatchThreadID)
{
	float sizeX, sizeY;
	OutputTexture.GetDimensions(sizeX, sizeY);
	if (id.x < sizeX && id.y < sizeY)
	{
		float2 uv = float2(id.x / sizeX, id.y / sizeY);
		OutputTexture[id.xy] = float4(uv, 0, 1);
	}
}

TaCompute.cpp

#include "TaCompute.h"

#include "Misc/FileHelper.h"
#include "HAL/FileManager.h"
#include "Modules/ModuleManager.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "Engine/Texture2D.h"
#include "Engine/TextureRenderTarget2D.h"
#include "GlobalShader.h"
#include "RenderGraphUtils.h"
#include "ShaderParameterStruct.h"
#include "RenderGraphBuilder.h"

// RDG 版本
class FComputeShaderRDG : public FGlobalShader
{
public:
	DECLARE_GLOBAL_SHADER(FComputeShaderRDG)
	SHADER_USE_PARAMETER_STRUCT(FComputeShaderRDG, FGlobalShader)

	BEGIN_SHADER_PARAMETER_STRUCT(FParameters,)
		SHADER_PARAMETER(FIntPoint, TextureSize)
		SHADER_PARAMETER_RDG_TEXTURE_UAV(RWTexture2D<float4>, OutputTexture) // RDG UAV
		END_SHADER_PARAMETER_STRUCT()

	static inline FIntVector ThreadGroupSize = FIntVector(16, 16, 1);

	static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters)
	{
		return RHISupportsComputeShaders(Parameters.Platform);
	}

	static void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters,
	                                         FShaderCompilerEnvironment& Environment)
	{
		Environment.SetDefine(TEXT("THREADGROUPSIZE_X"), ThreadGroupSize.X);
		Environment.SetDefine(TEXT("THREADGROUPSIZE_Y"), ThreadGroupSize.Y);
		Environment.SetDefine(TEXT("THREADGROUPSIZE_Z"), ThreadGroupSize.Z);
	}

	static void Execute(FRHICommandListImmediate& RHICmdList, UTextureRenderTarget2D* renderTarget)
	{
		int32 Width = renderTarget->SizeX;
		int32 Height = renderTarget->SizeY;
		if (Width < 1 || Height < 1) return;

		FRDGBuilder GraphBuilder(RHICmdList);
		const TShaderMapRef<FComputeShaderRDG> ComputeShader(GetGlobalShaderMap(GMaxRHIFeatureLevel));

		// create resource
		const FRDGTextureDesc& RenderTargetDesc = FRDGTextureDesc::Create2D(
			FIntPoint(Width, Height),
			renderTarget->GetFormat(),
			FClearValueBinding::Black,
			TexCreate_RenderTargetable |
			TexCreate_ShaderResource | TexCreate_UAV);
		FRDGTextureRef RDGRenderTarget = GraphBuilder.CreateTexture(RenderTargetDesc, TEXT("RDGRenderTarget"));
		FRDGTextureUAVDesc UAVDesc(RDGRenderTarget);

		// set params
		FParameters* PassParameters = GraphBuilder.AllocParameters<FParameters>();
		PassParameters->OutputTexture = GraphBuilder.CreateUAV(UAVDesc);

		uint32 GroupSizeX = FMath::DivideAndRoundUp(Width, ThreadGroupSize.X);
		uint32 GroupSizeY = FMath::DivideAndRoundUp(Height, ThreadGroupSize.Y);

		FComputeShaderUtils::AddPass(
			GraphBuilder,
			RDG_EVENT_NAME("ComputeShaderRDG"),
			ERDGPassFlags::Compute,
			ComputeShader,
			PassParameters,
			FIntVector(GroupSizeX, GroupSizeY, 1));

		TRefCountPtr<IPooledRenderTarget> PooledRenderTarget;
		GraphBuilder.QueueTextureExtraction(RDGRenderTarget, &PooledRenderTarget);

		GraphBuilder.Execute();

		RHICmdList.CopyTexture(PooledRenderTarget->GetRenderTargetItem().ShaderResourceTexture,
		                       renderTarget->GetRenderTargetResource()->TextureRHI, FRHICopyTextureInfo());
	}
};

// 旧版本
class FComputeShaderDemo : public FGlobalShader
{
public:
	DECLARE_GLOBAL_SHADER(FComputeShaderDemo)
	SHADER_USE_PARAMETER_STRUCT(FComputeShaderDemo, FGlobalShader)

	BEGIN_SHADER_PARAMETER_STRUCT(FParameters,)
		SHADER_PARAMETER(FIntPoint, TextureSize)
		SHADER_PARAMETER_UAV(RWTexture2D<float4>, OutputTexture) // Old UAV
		END_SHADER_PARAMETER_STRUCT()

	static inline FIntVector ThreadGroupSize = FIntVector(16, 16, 1);

	static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters)
	{
		return RHISupportsComputeShaders(Parameters.Platform);
	}

	static void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters,
	                                         FShaderCompilerEnvironment& Environment)
	{
		Environment.SetDefine(TEXT("THREADGROUPSIZE_X"), ThreadGroupSize.X);
		Environment.SetDefine(TEXT("THREADGROUPSIZE_Y"), ThreadGroupSize.Y);
		Environment.SetDefine(TEXT("THREADGROUPSIZE_Z"), ThreadGroupSize.Z);
	}

	static void Execute(FRHICommandListImmediate& RHICmdList, UTextureRenderTarget2D* renderTarget)
	{
		check(IsInRenderingThread());
		if (!renderTarget) return;

		const TShaderMapRef<FComputeShaderDemo> computeShader(GetGlobalShaderMap(GMaxRHIFeatureLevel));
		RHICmdList.SetComputeShader(computeShader.GetComputeShader());

		int32 Width = renderTarget->SizeX;
		int32 Height = renderTarget->SizeY;
		if (Width < 1 || Height < 1) return;

		// create gpu resource
		FRHIResourceCreateInfo TextureInfo(TEXT("TextureInfo"));
		const FTexture2DRHIRef GPUOutTexture = RHICreateTexture2D(
			Width, Height, renderTarget->GetFormat(), 1, 1,
			TexCreate_ShaderResource | ETextureCreateFlags::UAV, TextureInfo);
		FUnorderedAccessViewRHIRef GPUOutTextureUAV = RHICreateUnorderedAccessView(GPUOutTexture);

		FParameters PassParameters;
		PassParameters.OutputTexture = GPUOutTextureUAV;
		PassParameters.TextureSize = FIntPoint(Width, Height);

		uint32 GroupSizeX = FMath::DivideAndRoundUp(Width, ThreadGroupSize.X);
		uint32 GroupSizeY = FMath::DivideAndRoundUp(Height, ThreadGroupSize.Y);
		FComputeShaderUtils::Dispatch(RHICmdList, computeShader, PassParameters, FIntVector(GroupSizeX, GroupSizeY, 1));

		RHICmdList.CopyTexture(GPUOutTexture,
		                       renderTarget->GetRenderTargetResource()->TextureRHI,
		                       FRHICopyTextureInfo());

		GPUOutTextureUAV.SafeRelease();
	}
};

IMPLEMENT_GLOBAL_SHADER(FComputeShaderRDG, "/TaTools/Private/TaSimpleColorCS.usf", "MainCS", SF_Compute);
IMPLEMENT_GLOBAL_SHADER(FComputeShaderDemo, "/TaTools/Private/TaSimpleColorCS.usf", "MainCS", SF_Compute);

// 蓝图调用
void UTaCompute::RunComputeShader(UTextureRenderTarget2D* Texture)
{
	ENQUEUE_RENDER_COMMAND(CommandList)(
		[Texture](FRHICommandListImmediate& RHICmdList)
		{
			FComputeShaderRDG::Execute(RHICmdList, Texture);
		}
	);
}

TaCompute.h

#pragma once
#include "CoreMinimal.h"
#include "RHICommandList.h"
#include "Engine/Texture.h"
#include "TaCompute.generated.h"

UCLASS()
class UTaCompute : public UObject
{
	GENERATED_BODY()
public:
	UFUNCTION(BlueprintCallable, Category="Compute Shader")
	static void RunComputeShader(UTextureRenderTarget2D* Texture);
};

Module.cpp 添加Shader路径映射

void FTaShaderModule::StartupModule()
{
    FString PluginShaderDir = FPaths::Combine(IPluginManager::Get().FindPlugin(TEXT("TaTools"))->GetBaseDir(), TEXT("Shaders"));
    AddShaderSourceDirectoryMapping(TEXT("/TaTools"), PluginShaderDir);
}

参考

Unreal

#Shader #Unreal #UE5 #C++

UE5 Compute Shader 入门

https://automask.github.io/wild/2022/11/20/lab/S_Unreal_ShaderCompute/

作者

Kyle Zhou

发布于

2022年11月20日

许可协议

UE5 保存纹理上一篇

Unity Marching Square 网格生成下一篇

UE5 Compute Shader 入门

Compute Shader

相关API

FRHICommandListImmediate

FRDGBuilder

Buffer

代码

参考