UE5 compute shader 原子加

核心链路：

复制代码

蓝图调用 ReadRTPercent
-> C++ 把 RenderTarget 注册进 RDG
-> Dispatch Global Compute Shader
-> Shader 统计命中像素数到 CountBuffer
-> GPU Copy 到 Readback
-> C++ 异步读回 uint32
-> 算 Percent01 = HitPixelCount / TotalPixelCount
-> 回调给蓝图

Source/MyProject/Public/RTPercentCompute.h

复制代码

#pragma once

#include "CoreMinimal.h"
#include "Modules/ModuleManager.h"

class FRTPercentComputeModule : public IModuleInterface
{
public:
    virtual void StartupModule() override;
    virtual void ShutdownModule() override;
};

Source/MyProject/Private/RTPercentCompute.cpp

复制代码

#include "RTPercentCompute.h"

#include "Interfaces/IPluginManager.h"
#include "Misc/Paths.h"
#include "ShaderCore.h"

void FRTPercentComputeModule::StartupModule()
{
    static const FString VirtualShaderDir = TEXT("/RTPercentCompute");

    TSharedPtr<IPlugin> Plugin = IPluginManager::Get().FindPlugin(TEXT("RTPercentCompute"));
    if (!Plugin.IsValid())
    {
        return;
    }

    const FString ShaderDir = FPaths::Combine(Plugin->GetBaseDir(), TEXT("Shaders"));
    if (!AllShaderSourceDirectoryMappings().Contains(VirtualShaderDir))
    {
        AddShaderSourceDirectoryMapping(VirtualShaderDir, ShaderDir);
    }
}

void FRTPercentComputeModule::ShutdownModule()
{
}

IMPLEMENT_MODULE(FRTPercentComputeModule, RTPercentCompute)

这是模块启动逻辑。核心作用：把插件真实 shader 文件夹映射到虚拟 shader 路径。

关键代码：

复制代码

static const FString VirtualShaderDir = TEXT("/RTPercentCompute");
const FString ShaderDir = FPaths::Combine(Plugin->GetBaseDir(), TEXT("Shaders"));
AddShaderSourceDirectoryMapping(VirtualShaderDir, ShaderDir);

真实路径是：

复制代码

Plugins/RTPercentCompute/Shaders

虚拟路径是：

复制代码

/RTPercentCompute

所以 C++ 里这句：

复制代码

"/RTPercentCompute/Private/RTPercentCS.usf"

实际会找到：

复制代码

Plugins/RTPercentCompute/Shaders/Private/RTPercentCS.usf

这里还加了：

复制代码

AllShaderSourceDirectoryMappings().Contains(VirtualShaderDir)

这是为了防止 Live Coding 或热重载时重复注册 /RTPercentCompute，否则 UE 会 check 崩掉。

Source/MyProject/Public/RTPercentComputeBPLibrary.h

复制代码

#pragma once

#include "CoreMinimal.h"
#include "Kismet/BlueprintFunctionLibrary.h"
#include "Engine/TextureRenderTarget2D.h"
#include "RTPercentComputeBPLibrary.generated.h"

UENUM(BlueprintType)
enum class ERTPercentChannel : uint8
{
    Red     UMETA(DisplayName = "Red"),
    Green   UMETA(DisplayName = "Green"),
    Blue    UMETA(DisplayName = "Blue"),
    Alpha   UMETA(DisplayName = "Alpha")
};

DECLARE_DYNAMIC_DELEGATE_ThreeParams(
    FRTPercentComputeCallback,
    float, Percent01,
    int32, HitPixelCount,
    int32, TotalPixelCount
);

UCLASS()
class RTPERCENTCOMPUTE_API URTPercentComputeBPLibrary : public UBlueprintFunctionLibrary
{
    GENERATED_BODY()

public:
    UFUNCTION(BlueprintCallable, Category = "RT Percent Compute")
    static void ReadRTPercent(
        UTextureRenderTarget2D* RenderTarget,
        float Threshold,
        ERTPercentChannel Channel,
        FRTPercentComputeCallback OnCompleted
    );
};

这是蓝图接口层。它定义了蓝图能看到什么。

通道枚举：

复制代码

enum class ERTPercentChannel : uint8
{
    Red,
    Green,
    Blue,
    Alpha
};

意思是你可以选择统计 RT 的哪个通道，比如 mask 常用 Red 或 Alpha。

回调：

复制代码

DECLARE_DYNAMIC_DELEGATE_ThreeParams(
    FRTPercentComputeCallback,
    float, Percent01,
    int32, HitPixelCount,
    int32, TotalPixelCount
);

返回三个值：

Percent01：百分比，范围 0~1。

HitPixelCount：满足阈值的像素数。

TotalPixelCount：总像素数，比如 256*256 = 65536。

Source/MyProject/Private/RTPercentComputeBPLibrary.cpp

复制代码

#include "RTPercentComputeBPLibrary.h"

#include "Async/Async.h"
#include "Containers/Ticker.h"
#include "GlobalShader.h"
#include "RenderGraphBuilder.h"
#include "RenderGraphUtils.h"
#include "RenderTargetPool.h"
#include "RHI.h"
#include "RHICommandList.h"
#include "RHIGPUReadback.h"
#include "ShaderParameterStruct.h"
#include "TextureResource.h"

class FRTPercentCS : public FGlobalShader
{
public:
    DECLARE_GLOBAL_SHADER(FRTPercentCS);
    SHADER_USE_PARAMETER_STRUCT(FRTPercentCS, FGlobalShader);

    BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
        SHADER_PARAMETER_RDG_TEXTURE(Texture2D, InputTexture)
        SHADER_PARAMETER_RDG_BUFFER_UAV(RWStructuredBuffer<uint>, CountBuffer)
        SHADER_PARAMETER(FVector4f, Settings)
    END_SHADER_PARAMETER_STRUCT()

    static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters)
    {
        return IsFeatureLevelSupported(Parameters.Platform, ERHIFeatureLevel::SM5);
    }
};

IMPLEMENT_GLOBAL_SHADER(FRTPercentCS, "/RTPercentCompute/Private/RTPercentCS.usf", "MainCS", SF_Compute);

static uint32 GetRTPercentChannelIndex(ERTPercentChannel Channel)
{
    switch (Channel)
    {
    case ERTPercentChannel::Green:
        return 1;
    case ERTPercentChannel::Blue:
        return 2;
    case ERTPercentChannel::Alpha:
        return 3;
    case ERTPercentChannel::Red:
    default:
        return 0;
    }
}

void URTPercentComputeBPLibrary::ReadRTPercent(
    UTextureRenderTarget2D* RenderTarget,
    float Threshold,
    ERTPercentChannel Channel,
    FRTPercentComputeCallback OnCompleted
)
{
    if (!RenderTarget || RenderTarget->SizeX <= 0 || RenderTarget->SizeY <= 0)
    {
        OnCompleted.ExecuteIfBound(0.0f, 0, 0);
        return;
    }

    FTextureRenderTargetResource* RTResource = RenderTarget->GameThread_GetRenderTargetResource();
    if (!RTResource)
    {
        OnCompleted.ExecuteIfBound(0.0f, 0, 0);
        return;
    }

    const int32 Width = RenderTarget->SizeX;
    const int32 Height = RenderTarget->SizeY;
    const int32 TotalPixelCount = Width * Height;
    const float SafeThreshold = FMath::Clamp(Threshold, 0.0f, 1.0f);
    const uint32 ChannelIndex = GetRTPercentChannelIndex(Channel);

    TSharedRef<FRHIGPUBufferReadback, ESPMode::ThreadSafe> Readback =
        MakeShared<FRHIGPUBufferReadback, ESPMode::ThreadSafe>(TEXT("RTPercentReadback"));

    ENQUEUE_RENDER_COMMAND(RTPercentComputeCommand)(
        [RTResource, Readback, Width, Height, SafeThreshold, ChannelIndex](FRHICommandListImmediate& RHICmdList)
        {
            FTextureRHIRef TextureRHI = RTResource->GetRenderTargetTexture();
            if (!TextureRHI.IsValid())
            {
                return;
            }

            FRDGBuilder GraphBuilder(RHICmdList);

            FRDGTextureRef InputTexture = RegisterExternalTexture(GraphBuilder, TextureRHI, TEXT("RTPercentInputTexture"));

            FRDGBufferRef CountBuffer = GraphBuilder.CreateBuffer(
                FRDGBufferDesc::CreateStructuredDesc(sizeof(uint32), 1),
                TEXT("RTPercentCountBuffer")
            );

            FRDGBufferUAVRef CountUAV = GraphBuilder.CreateUAV(CountBuffer);
            AddClearUAVPass(GraphBuilder, CountUAV, 0u);

            FRTPercentCS::FParameters* PassParameters =
                GraphBuilder.AllocParameters<FRTPercentCS::FParameters>();

            PassParameters->InputTexture = InputTexture;
            PassParameters->CountBuffer = CountUAV;
            PassParameters->Settings = FVector4f(
                static_cast<float>(Width),
                static_cast<float>(Height),
                SafeThreshold,
                static_cast<float>(ChannelIndex)
            );

            TShaderMapRef<FRTPercentCS> ComputeShader(GetGlobalShaderMap(GMaxRHIFeatureLevel));

            const FIntVector GroupCount(
                FMath::DivideAndRoundUp(Width, 16),
                FMath::DivideAndRoundUp(Height, 16),
                1
            );

            FComputeShaderUtils::AddPass(
                GraphBuilder,
                RDG_EVENT_NAME("RTPercentCompute"),
                ComputeShader,
                PassParameters,
                GroupCount
            );

            AddEnqueueCopyPass(GraphBuilder, &Readback.Get(), CountBuffer, sizeof(uint32));

            GraphBuilder.Execute();
        }
        );

    FTSTicker::GetCoreTicker().AddTicker(
        FTickerDelegate::CreateLambda(
            [Readback, OnCompleted, TotalPixelCount](float DeltaTime) mutable
            {
                if (!Readback->IsReady())
                {
                    return true;
                }

                ENQUEUE_RENDER_COMMAND(RTPercentReadbackCommand)(
                    [Readback, OnCompleted, TotalPixelCount](FRHICommandListImmediate& RHICmdList) mutable
                    {
                        uint32 HitPixelCount = 0;

                        if (Readback->GetGPUSizeBytes() >= sizeof(uint32))
                        {
                            void* Data = Readback->Lock(sizeof(uint32));
                            if (Data)
                            {
                                HitPixelCount = *static_cast<uint32*>(Data);
                            }
                            Readback->Unlock();
                        }

                        AsyncTask(
                            ENamedThreads::GameThread,
                            [OnCompleted, HitPixelCount, TotalPixelCount]() mutable
                            {
                                const float Percent01 = TotalPixelCount > 0
                                    ? static_cast<float>(HitPixelCount) / static_cast<float>(TotalPixelCount)
                                    : 0.0f;

                                OnCompleted.ExecuteIfBound(
                                    Percent01,
                                    static_cast<int32>(HitPixelCount),
                                    TotalPixelCount
                                );
                            }
                        );
                    }
                );
                return false;
            }
        )
    );
}

进入渲染线程：

复制代码

ENQUEUE_RENDER_COMMAND(RTPercentComputeCommand)(...)

因为 RDG/RHI 资源不能随便在 GameThread 操作。

在 RenderThread 里：

复制代码

FRDGTextureRef InputTexture = RegisterExternalTexture(GraphBuilder, TextureRHI, TEXT("RTPercentInputTexture"));

把外部 RT 注册成 RDG 可用 texture。

然后创建统计 buffer：

复制代码

FRDGBufferRef CountBuffer = GraphBuilder.CreateBuffer(
    FRDGBufferDesc::CreateStructuredDesc(sizeof(uint32), 1),
    TEXT("RTPercentCountBuffer")
);

这里只有一个 uint32，用来存命中像素数。

清零：

复制代码

AddClearUAVPass(GraphBuilder, CountUAV, 0u);

不清零的话，上一次数据会污染结果。

Dispatch：

复制代码

FComputeShaderUtils::AddPass(..., ComputeShader, PassParameters, GroupCount);

GroupCount 是：

复制代码

DivideAndRoundUp(Width, 16)
DivideAndRoundUp(Height, 16)

因为 .usf 里每组是 16x16 线程。

最后 GPU readback：

复制代码

AddEnqueueCopyPass(GraphBuilder, &Readback.Get(), CountBuffer, sizeof(uint32));

把 GPU 上的 CountBuffer 拷到 CPU 可读的 staging buffer。

这个就是性能最敏感的地方，因为它涉及 GPU -> CPU。

后半段：

复制代码

FTSTicker::GetCoreTicker().AddTicker(...)

每帧检查 readback 是否 ready。ready 后再 enqueue 到 RenderThread 里：

复制代码

Readback->Lock(sizeof(uint32));

读出 HitPixelCount，然后：

复制代码

AsyncTask(ENamedThreads::GameThread, ...)

回到 GameThread 执行蓝图回调。这个设计是为了避免在 GameThread 直接 Lock RHI staging buffer 崩溃。

Shaders/Private/RTColoredPixelPercent.usf

复制代码

#include "/Engine/Public/Platform.ush"

Texture2D<float4> InputTexture;
RWStructuredBuffer<uint> CountBuffer;

// x = Width
// y = Height
// z = Threshold
// w = Channel: 0 R, 1 G, 2 B, 3 A
float4 Settings;

groupshared uint GroupHitCount;

[numthreads(16, 16, 1)]
void MainCS(
    uint3 DispatchThreadId : SV_DispatchThreadID,
    uint GroupIndex : SV_GroupIndex
)
{
    if (GroupIndex == 0)
    {
        GroupHitCount = 0;
    }

    GroupMemoryBarrierWithGroupSync();

    uint Width = (uint) Settings.x;
    uint Height = (uint) Settings.y;
    float Threshold = Settings.z;
    uint Channel = (uint) (Settings.w + 0.5);

    uint2 Pixel = DispatchThreadId.xy;

    if (Pixel.x < Width && Pixel.y < Height)
    {
        float4 Color = InputTexture.Load(int3(Pixel, 0));

        float Value = Color.r;

        if (Channel == 1)
        {
            Value = Color.g;
        }
        else if (Channel == 2)
        {
            Value = Color.b;
        }
        else if (Channel == 3)
        {
            Value = Color.a;
        }

        if (Value >= Threshold)
        {
            InterlockedAdd(GroupHitCount, 1);
        }
    }

    GroupMemoryBarrierWithGroupSync();

    if (GroupIndex == 0)
    {
        InterlockedAdd(CountBuffer[0], GroupHitCount);
    }
}

groupshared uint GroupHitCount;

每个 thread group 内部先统计到 GroupHitCount，不是所有像素都直接写全局 buffer。

组内累加：

复制代码

InterlockedAdd(GroupHitCount, 1);

同步：

复制代码

GroupMemoryBarrierWithGroupSync();

最后每组只做一次全局 atomic：

复制代码

if (GroupIndex == 0)
{
    InterlockedAdd(CountBuffer[0], GroupHitCount);
}

这就是为什么用了 groupshared：

如果 256x256 全白，直接全局 atomic 会有 65536 次；现在 16x16 一组，只会有 16*16 = 256 次全局 atomic，压力小很多。

Source/MyProject/MyProject.Build.cs

复制代码

using UnrealBuildTool;

public class RTPercentCompute : ModuleRules
{
    public RTPercentCompute(ReadOnlyTargetRules Target) : base(Target)
    {
        PCHUsage = PCHUsageMode.UseExplicitOrSharedPCHs;

        PublicDependencyModuleNames.AddRange(new string[]
        {
            "Core",
            "CoreUObject",
            "Engine"
        });

        PrivateDependencyModuleNames.AddRange(new string[]
        {
            "Projects",
            "RenderCore",
            "RHI",
            "Renderer"
        });
    }
}