Add ReadRTArea function

Source/AsyncReadRT/Private/AsyncReadRTAreaAction.cpp

@@ -0,0 +1,211 @@
+// Fill out your copyright notice in the Description page of Project Settings.
+
+
+#include "AsyncReadRTAreaAction.h"
+
+#include "AsyncReadRT.h"
+#include "RHICommandList.h"
+#include "Engine/TextureRenderTarget2D.h"
+#include "RenderGraphBuilder.h"
+#include "Async/Async.h"
+#include "CoreGlobals.h"
+#include "TextureResource.h"
+#include "Runtime/Launch/Resources/Version.h"
+#include "Engine/World.h"
+#include "TimerManager.h"
+#include "Chaos/ChaosPerfTest.h"
+
+UAsyncReadRTAreaAction* UAsyncReadRTAreaAction::AsyncReadRenderTargetArea(UObject* WorldContextObject, UTextureRenderTarget2D* TextureRenderTarget, int32 StartX, int32 StartY, int32 EndX, int32 EndY, bool bFlushRHI)
+{
+	UAsyncReadRTAreaAction* BlueprintNode = NewObject<UAsyncReadRTAreaAction>();
+	BlueprintNode->WorldContextObject = WorldContextObject;
+	BlueprintNode->RT = TextureRenderTarget;
+	BlueprintNode->StartX = StartX;
+	BlueprintNode->StartY = StartY;
+	BlueprintNode->EndX = EndX;
+	BlueprintNode->EndY = EndY;
+	BlueprintNode->bFlushRHI = bFlushRHI;
+	BlueprintNode->ReadRTData = MakeShared<FAsyncReadRTAreaData, ESPMode::ThreadSafe>();
+	BlueprintNode->ReadRTData->FinishedRead = false;
+	return BlueprintNode;
+}
+
+// Maps the texture and reads a single pixel.
+// If bFlushRHI is false, then it checks the render fence. If the render fence hasn't completed, then the function early exits
+static void PollRTRead(FRHICommandListImmediate& RHICmdList, 
+	TSharedPtr<FAsyncReadRTAreaData, ESPMode::ThreadSafe> ReadData,
+	TWeakObjectPtr<UAsyncReadRTAreaAction> ReadAction, bool bFlushRHI)
+{
+	SCOPED_NAMED_EVENT_TEXT("AsyncReadRTAction::AsyncReadRT::PollRTRead", FColor::Magenta);
+
+	check(IsInRenderingThread());
+
+	// If we didn't flush the RHI then make sure the previous rendering commands got done
+	if (!bFlushRHI)
+	{
+		// Return if we haven't finished the texture commands
+		if (!ReadData->TextureFence.IsValid() || !ReadData->TextureFence->Poll())
+		{
+			return;
+		}
+	}
+
+	SCOPED_NAMED_EVENT_TEXT("AsyncReadRTAction::AsyncReadRT::MapTexture", FColor::Magenta);
+	void* OutputBuffer = NULL;
+	int32 Width; int32 Height;
+
+	if (bFlushRHI)
+	{
+		// This flushes the command list
+		RHICmdList.MapStagingSurface(ReadData->Texture, ReadData->TextureFence, OutputBuffer, Width, Height);
+	}
+	else
+	{
+#if ENGINE_MAJOR_VERSION >= 5 && ENGINE_MINOR_VERSION >= 2
+		GDynamicRHI->RHIMapStagingSurface_RenderThread(RHICmdList, ReadData->Texture, INDEX_NONE, ReadData->TextureFence, OutputBuffer, Width, Height);
+#else
+		GDynamicRHI->RHIMapStagingSurface_RenderThread(RHICmdList, ReadData->Texture, ReadData->TextureFence, OutputBuffer, Width, Height);
+#endif
+	}
+
+	ReadData->PixelColors.Empty(Width * Height);
+
+	const EPixelFormat Format = ReadData->Texture->GetFormat();
+	for (int32 YIndex = 0; YIndex < Height; YIndex++)
+	{
+		for (int32 X = 0; X < Width; ++X)
+		{
+			const int32 PixelOffset = X + (YIndex * Width);
+
+			FLinearColor& OutColor = ReadData->PixelColors.AddDefaulted_GetRef();
+			switch (Format)
+			{
+			case EPixelFormat::PF_FloatRGBA:
+				{
+					FFloat16Color* OutputColor = reinterpret_cast<FFloat16Color*>(OutputBuffer) + PixelOffset;
+					OutColor.R = OutputColor->R;
+					OutColor.G = OutputColor->G;
+					OutColor.B = OutputColor->B;
+					OutColor.A = OutputColor->A;
+					break;
+				}
+			case EPixelFormat::PF_B8G8R8A8:
+				{
+					FColor* OutputColor = reinterpret_cast<FColor*>(OutputBuffer) + PixelOffset;
+					OutColor.R = OutputColor->R;
+					OutColor.G = OutputColor->G;
+					OutColor.B = OutputColor->B;
+					OutColor.A = OutputColor->A;
+					OutColor /= 255.f;
+					break;
+				}
+			default:
+				UE_LOG(LogTemp, Warning, TEXT("UAsyncReadRTAction: Unsupported RT format! Format: %d"), static_cast<int32>(Format)); // Unsupported, add a new switch statement.
+			}
+		}
+	}
+
+	RHICmdList.UnmapStagingSurface(ReadData->Texture);
+	ReadData->FinishedRead = true;
+}
+
+void UAsyncReadRTAreaAction::Activate()
+{
+	FTextureRenderTarget2DResource* TextureResource = (FTextureRenderTarget2DResource*)RT->GetResource();
+	check(TextureResource);
+	check(TextureResource->GetRenderTargetTexture());
+
+	StartX = FMath::Clamp(StartX, 0, RT->SizeX - 1);
+	EndX = FMath::Clamp(EndX, StartX, RT->SizeX - 1);
+	StartY = FMath::Clamp(StartY, 0, RT->SizeY - 1);
+	EndY = FMath::Clamp(EndY, StartY, RT->SizeY - 1);
+
+	StartFrame = GFrameCounter;
+
+	ENQUEUE_RENDER_COMMAND(FCopyRTAsync)([bFlushRHI = bFlushRHI, StartX = StartX, StartY = StartY, EndX=EndX, EndY=EndY, AsyncReadPtr = TWeakObjectPtr<UAsyncReadRTAreaAction>(this), TextureRHI = TextureResource->GetRenderTargetTexture(), ReadData = ReadRTData](FRHICommandListImmediate& RHICmdList)
+	{
+		check(IsInRenderingThread());
+		check(TextureRHI.IsValid());
+
+		FGPUFenceRHIRef Fence = RHICreateGPUFence(TEXT("AsyncRTReadback"));
+
+		SCOPED_NAMED_EVENT_TEXT("AsyncReadRTAction::AsyncReadRT", FColor::Magenta);
+
+#if ENGINE_MAJOR_VERSION >= 5 && ENGINE_MINOR_VERSION >= 1
+		FTextureRHIRef IORHITextureCPU;
+#else
+		FTexture2DRHIRef IORHITextureCPU;
+#endif
+		{
+			SCOPED_NAMED_EVENT_TEXT("AsyncReadRTAction::AsyncReadRT::CreateCopyTexture", FColor::Magenta);
+
+			int32 Width, Height;
+			Width = EndX-StartX+1;
+			Height = EndY-StartY+1;
+#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION > 2
+			FRHITextureCreateDesc TextureDesc = FRHITextureCreateDesc::Create2D(TEXT("AsyncRTReadback"), Width, Height, TextureRHI->GetFormat());
+			TextureDesc.AddFlags(ETextureCreateFlags::CPUReadback);
+			TextureDesc.InitialState = ERHIAccess::CopyDest;
+#if ENGINE_MINOR_VERSION >= 5
+			IORHITextureCPU = RHICreateTexture(TextureDesc);
+#elif ENGINE_MINOR_VERSION == 4
+			IORHITextureCPU = GDynamicRHI->RHICreateTexture(FRHICommandListExecutor::GetImmediateCommandList(), TextureDesc);
+#else // ENGINE_MINOR_VERSION
+			IORHITextureCPU = GDynamicRHI->RHICreateTexture(TextureDesc);
+#endif // ENGINE_MINOR_VERSION
+#else
+			FRHIResourceCreateInfo CreateInfo(TEXT("AsyncRTReadback"));
+			IORHITextureCPU = RHICreateTexture2D(1, 1, TextureRHI->GetFormat(), 1, 1, TexCreate_CPUReadback, ERHIAccess::CopyDest, CreateInfo);
+#endif
+
+			FRHICopyTextureInfo CopyTextureInfo;
+			CopyTextureInfo.Size = FIntVector(EndX-StartX+1, EndY-StartY+1, 1);
+			CopyTextureInfo.SourceMipIndex = 0;
+			CopyTextureInfo.DestMipIndex = 0;
+			CopyTextureInfo.SourcePosition = FIntVector(StartX, StartY, 0);
+			CopyTextureInfo.DestPosition = FIntVector(0, 0, 0);
+
+			RHICmdList.Transition(FRHITransitionInfo(TextureRHI, ERHIAccess::Unknown, ERHIAccess::CopySrc));
+			RHICmdList.CopyTexture(TextureRHI, IORHITextureCPU, CopyTextureInfo);
+
+			RHICmdList.Transition(FRHITransitionInfo(IORHITextureCPU, ERHIAccess::CopyDest, ERHIAccess::CopySrc));
+			RHICmdList.WriteGPUFence(Fence);
+		}
+		check(Fence.IsValid());
+
+		ReadData->Texture = IORHITextureCPU;
+		ReadData->TextureFence = Fence;
+
+		// If we flush the RHI then we can just go ahead and read the mapped texture asap
+		if (bFlushRHI)
+		{
+			PollRTRead(RHICmdList, ReadData, AsyncReadPtr, bFlushRHI);
+		}
+	});
+
+	WorldContextObject->GetWorld()->GetTimerManager().SetTimerForNextTick(this, &UAsyncReadRTAreaAction::OnNextFrame);
+}
+
+void UAsyncReadRTAreaAction::OnNextFrame()
+{
+	//const int32 FramesWaited = GFrameCounter - StartFrame;
+	//UE_LOG(LogTemp, Warning, TEXT("Frames waited: %d"), FramesWaited);
+
+	check(ReadRTData.IsValid());
+
+	if (ReadRTData->FinishedRead)
+	{
+		OnReadRenderTargetArea.Broadcast(ReadRTData->PixelColors);
+		SetReadyToDestroy();
+	}
+	else
+	{
+		ENQUEUE_RENDER_COMMAND(FReadRTAsync)([WeakThis = TWeakObjectPtr<UAsyncReadRTAreaAction>(this), ReadRTData = ReadRTData](FRHICommandListImmediate& RHICmdList)
+		{
+			PollRTRead(RHICmdList, ReadRTData, WeakThis, false);
+		});
+
+		WorldContextObject->GetWorld()->GetTimerManager().SetTimerForNextTick(this, &UAsyncReadRTAreaAction::OnNextFrame);
+	}
+}
+

Source/AsyncReadRT/Public/AsyncReadRTAreaAction.h

@@ -0,0 +1,63 @@
+// Fill out your copyright notice in the Description page of Project Settings.
+
+#pragma once
+
+#include "CoreMinimal.h"
+#include "Kismet/BlueprintAsyncActionBase.h"
+#include "RHIResources.h"
+#include "AsyncReadRTAreaAction.generated.h"
+
+DECLARE_DYNAMIC_MULTICAST_DELEGATE_OneParam(FAsyncReadRTAreaOutputPin, const TArray<FLinearColor>&, Color);
+
+class UTextureRenderTarget2D;
+
+struct FAsyncReadRTAreaData
+{
+	FGPUFenceRHIRef TextureFence;
+#if ENGINE_MAJOR_VERSION >= 5 && ENGINE_MINOR_VERSION >= 1
+	FTextureRHIRef Texture;
+#else
+	FTexture2DRHIRef Texture;
+#endif
+	TAtomic<bool> FinishedRead;
+	TArray<FLinearColor> PixelColors;
+};
+
+/**
+ * 
+ */
+UCLASS()
+class ASYNCREADRT_API UAsyncReadRTAreaAction : public UBlueprintAsyncActionBase
+{
+	GENERATED_BODY()
+public:
+	UFUNCTION(BlueprintCallable, meta = (BlueprintInternalUseOnly = "true", WorldContext = "WorldContextObject"), Category = "Rendering")
+	static UAsyncReadRTAreaAction* AsyncReadRenderTargetArea(UObject* WorldContextObject, UTextureRenderTarget2D* TextureRenderTarget, int32 StartX, int32 StartY, int32 EndX, int32 EndY, bool bFlushRHI);
+
+	// UBlueprintAsyncActionBase interface
+	virtual void Activate() override;
+	//~UBlueprintAsyncActionBase interface
+
+	UPROPERTY()
+	UObject* WorldContextObject;
+
+	UPROPERTY()
+	UTextureRenderTarget2D* RT;
+
+	int32 StartX;
+	int32 StartY;
+	int32 EndX;
+	int32 EndY;
+	bool bFlushRHI;
+
+	UPROPERTY(BlueprintAssignable)
+	FAsyncReadRTAreaOutputPin OnReadRenderTargetArea;
+
+	TSharedPtr<FAsyncReadRTAreaData, ESPMode::ThreadSafe> ReadRTData;
+
+protected:
+	UFUNCTION()
+	void OnNextFrame();
+
+	uint64 StartFrame;
+};