Unreal indirect draw

IndirectActor.cpp

#include "IndirectActor.h"
#include "RenderResource.h"
#include "LocalVertexFactory.h"

#include "GlobalShader.h"
#include "RHI.h"
#include "RHIUtilities.h"
#include "RenderGraphBuilder.h"
#include "RenderGraphUtils.h"
#include "HLSLTypeAliases.h"
#include "SceneViewExtension.h"
#include "EngineUtils.h"

class FPopulateVertexAndIndirectBufferCS : public FGlobalShader
{
public:
	static constexpr uint32 ThreadGroupSize = 64;

	DECLARE_GLOBAL_SHADER(FPopulateVertexAndIndirectBufferCS);
	SHADER_USE_PARAMETER_STRUCT(FPopulateVertexAndIndirectBufferCS, FGlobalShader);

	BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
		SHADER_PARAMETER_UAV(RWBuffer<float>, OutInstanceVertices)
		SHADER_PARAMETER_UAV(RWBuffer<uint>, OutIndirectArgs)
	END_SHADER_PARAMETER_STRUCT()
};

IMPLEMENT_GLOBAL_SHADER(FPopulateVertexAndIndirectBufferCS, "/Project/Private/PopulateVertexAndIndirectBuffer.usf", "MainCS", SF_Compute);

void AddPopulateVertexPass(FRDGBuilder& GraphBuilder,
	FRHIUnorderedAccessView* PositionsUAV,
	FRHIUnorderedAccessView* IndirectArgsBufferUAV
)
{
	TShaderMapRef<FPopulateVertexAndIndirectBufferCS> ComputeShader(GetGlobalShaderMap(GMaxRHIFeatureLevel));

	FPopulateVertexAndIndirectBufferCS::FParameters* PassParameters = GraphBuilder.AllocParameters<FPopulateVertexAndIndirectBufferCS::FParameters>();
	PassParameters->OutInstanceVertices = PositionsUAV;
	PassParameters->OutIndirectArgs = IndirectArgsBufferUAV;

	FComputeShaderUtils::AddPass(
		GraphBuilder,
		RDG_EVENT_NAME("PopulateVertexAndIndirectBuffer"),
		ERDGPassFlags::Compute | ERDGPassFlags::NeverCull,
		ComputeShader,
		PassParameters,
		FIntVector(1, 1, 1)
	);
}

class FPositionUAVVertexBuffer : public FVertexBuffer
{
public:
	virtual const TCHAR* GetName() const
	{
		return TEXT("TestPostionBuffer");
	}

	FPositionUAVVertexBuffer(int32 InNumVertices)
		: FVertexBuffer()
		, NumVertices(InNumVertices)
	{
	}

	virtual void InitRHI(FRHICommandListBase& RHICmdList) override
	{
		TRACE_CPUPROFILER_EVENT_SCOPE(FPositionUAVVertexBuffer::InitRHI);

		const uint32 Size = NumVertices * sizeof(FVector3f);
		const FRHIBufferCreateDesc CreateDesc =
			FRHIBufferCreateDesc::CreateVertex(GetName(), Size)
			.AddUsage(EBufferUsageFlags::Static)
			.AddUsage(EBufferUsageFlags::ShaderResource)
			.AddUsage(EBufferUsageFlags::UnorderedAccess)
			.SetInitialState(ERHIAccess::VertexOrIndexBuffer)
			.SetInitActionNone();

		VertexBufferRHI = RHICmdList.CreateBuffer(CreateDesc);

		if (VertexBufferRHI)
		{
			SRV = RHICmdList.CreateShaderResourceView(
				VertexBufferRHI,
				FRHIViewDesc::CreateBufferSRV()
				.SetType(FRHIViewDesc::EBufferType::Typed)
				.SetFormat(PF_R32_FLOAT));

			UAV = RHICmdList.CreateUnorderedAccessView(
				VertexBufferRHI,
				FRHIViewDesc::CreateBufferUAV()
				.SetType(FRHIViewDesc::EBufferType::Typed)
				.SetFormat(PF_R32_FLOAT)
			);
		}
	}

	virtual void ReleaseRHI() override
	{
		UAV.SafeRelease();
		SRV.SafeRelease();
		FVertexBuffer::ReleaseRHI();
	}

	FRHIShaderResourceView* GetSRV() const { return SRV; }
	FRHIUnorderedAccessView* GetUAV() const { return UAV; }
private:
	int32 NumVertices;
	FShaderResourceViewRHIRef SRV;
	FUnorderedAccessViewRHIRef UAV;
};

class FIndirectSceneProxy final : public FPrimitiveSceneProxy
{
public:
	FLocalVertexFactory* VertexFactory;
	FPositionUAVVertexBuffer* PositionBuffer;
	TRefCountPtr<FRHIBuffer> IndirectArgsBuffer;
	FUnorderedAccessViewRHIRef IndirectArgsBufferUAV;

public:
	FIndirectSceneProxy(UIndirectComponent* Owner) : FPrimitiveSceneProxy(Owner)
		, VertexFactory{ nullptr }
		, PositionBuffer{ nullptr }
		, IndirectArgsBuffer{ nullptr }
		, IndirectArgsBufferUAV{}
	{
		bVFRequiresPrimitiveUniformBuffer = false;
	}

	virtual SIZE_T GetTypeHash() const override
	{
		static size_t UniquePointer;
		return reinterpret_cast<size_t>(&UniquePointer);
	}

	virtual uint32 GetMemoryFootprint(void) const override
	{
		return(sizeof(*this) + GetAllocatedSize());
	}

	virtual FPrimitiveViewRelevance GetViewRelevance(const FSceneView* View) const override
	{
		FPrimitiveViewRelevance Result;
		Result.bDrawRelevance = IsShown(View);
		Result.bShadowRelevance = IsShadowCast(View);
		Result.bDynamicRelevance = true;
		Result.bStaticRelevance = false;
		Result.bRenderInMainPass = ShouldRenderInMainPass();
		Result.bRenderInDepthPass = ShouldRenderInDepthPass();
		Result.bVelocityRelevance = DrawsVelocity();
		return Result;
	}

	virtual void CreateRenderThreadResources(FRHICommandListBase& RHICmdList) override
	{
		check(VertexFactory == nullptr);

		PositionBuffer = new FPositionUAVVertexBuffer(3);
		PositionBuffer->InitResource(RHICmdList);

		VertexFactory = new FLocalVertexFactory(GetScene().GetFeatureLevel(), "TestIndirectVertexFactory");
		FLocalVertexFactory::FDataType NewData;
		NewData.PositionComponent = FVertexStreamComponent(PositionBuffer, 0, sizeof(FVector3f), VET_Float3);
		if (RHISupportsManualVertexFetch(GMaxRHIShaderPlatform))
		{
			NewData.PositionComponentSRV = PositionBuffer->GetSRV();
		}

		NewData.ColorComponent = FVertexStreamComponent(&GNullColorVertexBuffer, 0, 0, VET_Color, EVertexStreamUsage::ManualFetch);
		NewData.TangentBasisComponents[0] = FVertexStreamComponent(&GNullColorVertexBuffer, 0, 0, VET_PackedNormal, EVertexStreamUsage::ManualFetch);
		NewData.TangentBasisComponents[1] = FVertexStreamComponent(&GNullColorVertexBuffer, 0, 0, VET_PackedNormal, EVertexStreamUsage::ManualFetch);

		if (RHISupportsManualVertexFetch(GMaxRHIShaderPlatform))
		{
			NewData.ColorComponentsSRV = GNullColorVertexBuffer.VertexBufferSRV;
			NewData.TangentsSRV = GNullColorVertexBuffer.VertexBufferSRV;
			NewData.TextureCoordinatesSRV = GNullColorVertexBuffer.VertexBufferSRV;
		}


		VertexFactory->SetData(RHICmdList, NewData);
		VertexFactory->InitResource(RHICmdList);

		FRHIBufferCreateDesc IndirectBufferCreateDesc =
			FRHIBufferCreateDesc::CreateVertex<uint32>(
				TEXT("CustomIndirectArgs"), uint32(sizeof(FRHIDrawIndirectParameters) / sizeof(uint32))
			)
			.AddUsage(EBufferUsageFlags::UnorderedAccess | EBufferUsageFlags::DrawIndirect)
			.SetInitialState(ERHIAccess::IndirectArgs);
		IndirectArgsBuffer = RHICmdList.CreateBuffer(IndirectBufferCreateDesc);
		IndirectArgsBufferUAV = RHICmdList.CreateUnorderedAccessView(
			IndirectArgsBuffer,
			FRHIViewDesc::CreateBufferUAV()
			.SetTypeFromBuffer(IndirectArgsBuffer)
			.SetFormat(PF_R32_UINT)
			.SetNumElements(uint32(sizeof(FRHIDrawIndirectParameters) / sizeof(uint32)))
		);
	}

	virtual void DestroyRenderThreadResources() override
	{
		if (VertexFactory != nullptr)
		{
			VertexFactory->ReleaseResource();
			delete VertexFactory;
			VertexFactory = nullptr;
		}
		if (PositionBuffer != nullptr)
		{
			PositionBuffer->ReleaseResource();
			delete PositionBuffer;
			PositionBuffer = nullptr;
		}
		if (IndirectArgsBuffer)
		{
			IndirectArgsBuffer.SafeRelease();
			IndirectArgsBuffer = nullptr;
		}
	}

	virtual void GetDynamicMeshElements(const TArray<const FSceneView*>& Views, const FSceneViewFamily& ViewFamily, uint32 VisibilityMap, class FMeshElementCollector& Collector) const
	{
		for (int32 ViewIndex = 0; ViewIndex < Views.Num(); ViewIndex++)
		{
			if (FMeshBatch* MeshBatch = CreateMeshBatch(Collector))
			{
				Collector.AddMesh(ViewIndex, *MeshBatch);
			}
		}
	}

	FMeshBatch* CreateMeshBatch(class FMeshElementCollector& Collector) const
	{
		UMaterial* DefaultMaterial = LoadObject<UMaterial>(nullptr, TEXT("/Script/Engine.Material'/Engine/EngineMaterials/WorldGridMaterial.WorldGridMaterial'"));
		FMaterialRenderProxy* MaterialRenderProxy = DefaultMaterial->GetRenderProxy();

		if (MaterialRenderProxy == nullptr)
		{
			return nullptr;
		}

		if (!VertexFactory || !PositionBuffer || !IndirectArgsBuffer)
		{
			return nullptr;
		}

		FMeshBatch& MeshBatch = Collector.AllocateMesh();

		MeshBatch.CastShadow = true;
		MeshBatch.bUseForDepthPass = true;
		MeshBatch.SegmentIndex = 0;

		FMeshBatchElement& MeshBatchElement = MeshBatch.Elements[0];
		MeshBatch.VertexFactory = VertexFactory;
		MeshBatch.MaterialRenderProxy = MaterialRenderProxy;

		MeshBatchElement.NumPrimitives = 0;
		MeshBatchElement.IndirectArgsBuffer = IndirectArgsBuffer;
		MeshBatchElement.IndirectArgsOffset = 0;

		check(MeshBatchElement.IndirectArgsBuffer);

		return &MeshBatch;
	}
};

UIndirectComponent::UIndirectComponent(FObjectInitializer const& Initializer)
	: UPrimitiveComponent(Initializer)
{
}

FPrimitiveSceneProxy* UIndirectComponent::CreateSceneProxy()
{
	return new FIndirectSceneProxy(this);
}


FBoxSphereBounds UIndirectComponent::CalcBounds(const FTransform& BoundTransform) const
{
	// Some fake very big bounds
	FBox Box(FVector(0,0,0), FVector(1000, 1000, 1000));
	return FBoxSphereBounds(Box);
}

AIndirectActor::AIndirectActor(FObjectInitializer const& Initializer)
	: AActor(Initializer)
{
	IndirectComponent = CreateDefaultSubobject<UIndirectComponent>(TEXT("IndirectComponent"));
	SetRootComponent(IndirectComponent);

	// Do not do this in production code
	if (!IsTemplate())
	{
		auto Subsystem = GetWorld()->GetSubsystem<UIndirectActorSubsystem>();
		Subsystem->TheActor = this;
	}
}


class FIndirectPopulateSceneViewExtension : public FWorldSceneViewExtension
{
public:
	FIndirectPopulateSceneViewExtension(const FAutoRegister& AutoReg, UWorld* InWorld, UIndirectActorSubsystem* System)
		: FWorldSceneViewExtension(AutoReg, InWorld)
		, OwnerSystem(System)
	{
	}
	void Invalidate()
	{
		OwnerSystem = nullptr;
	}

	//~ Begin ISceneViewExtension interface
	virtual void  PreRenderViewFamily_RenderThread(FRDGBuilder& GraphBuilder, FSceneViewFamily& InViewFamily) override
	{
		if (IsValid(OwnerSystem))
		{
			if (AIndirectActor* Actor = OwnerSystem->TheActor.Get())
			{
				UIndirectComponent* Comp = Actor->IndirectComponent;
				FIndirectSceneProxy* SceneProxy = static_cast<FIndirectSceneProxy*>(Comp->GetSceneProxy());
				if (SceneProxy && SceneProxy->PositionBuffer && SceneProxy->IndirectArgsBuffer)
				{
					AddPopulateVertexPass(GraphBuilder, SceneProxy->PositionBuffer->GetUAV(), SceneProxy->IndirectArgsBufferUAV);
				}
			}
		}
	}
	//~ End ISceneViewExtension interface

public:
	UIndirectActorSubsystem* OwnerSystem;
};

void UIndirectActorSubsystem::Initialize(FSubsystemCollectionBase& Collection)
{
	SceneViewExtension = FSceneViewExtensions::NewExtension<FIndirectPopulateSceneViewExtension>(GetWorld(), this);
	Super::Initialize(Collection);
}

void UIndirectActorSubsystem::Deinitialize()
{
	if (SceneViewExtension)
	{
		// Prevent this SVE from being gathered, in case it is kept alive by a strong reference somewhere else.
		{
			SceneViewExtension->IsActiveThisFrameFunctions.Empty();

			FSceneViewExtensionIsActiveFunctor IsActiveFunctor;

			IsActiveFunctor.IsActiveFunction = [](const ISceneViewExtension* SceneViewExtension, const FSceneViewExtensionContext& Context)
				{
					return TOptional<bool>(false);
				};

			SceneViewExtension->IsActiveThisFrameFunctions.Add(IsActiveFunctor);
		}

		ENQUEUE_RENDER_COMMAND(ReleaseSVE)([this](FRHICommandListImmediate& RHICmdList)
			{
				// Prevent this SVE from being gathered, in case it is kept alive by a strong reference somewhere else.
				{
					SceneViewExtension->IsActiveThisFrameFunctions.Empty();

					FSceneViewExtensionIsActiveFunctor IsActiveFunctor;

					IsActiveFunctor.IsActiveFunction = [](const ISceneViewExtension* SceneViewExtension, const FSceneViewExtensionContext& Context)
						{
							return TOptional<bool>(false);
						};

					SceneViewExtension->IsActiveThisFrameFunctions.Add(IsActiveFunctor);
				}

				SceneViewExtension->Invalidate();
				SceneViewExtension.Reset();
				SceneViewExtension = nullptr;
			});
	}

	// Finish all rendering commands first
	FlushRenderingCommands();

	Super::Deinitialize();
}

IndirectActor.h

#pragma once

#include "CoreMinimal.h"
#include "IndirectActor.generated.h"

UCLASS()
class UIndirectComponent : public UPrimitiveComponent
{
	GENERATED_UCLASS_BODY()
	virtual FPrimitiveSceneProxy* CreateSceneProxy() override;
	virtual FBoxSphereBounds CalcBounds(const FTransform& BoundTransform) const override;
};


UCLASS(Placeable)
class AIndirectActor : public AActor
{
	GENERATED_UCLASS_BODY()

public:
	UPROPERTY()
	UIndirectComponent* IndirectComponent;
};

class FIndirectPopulateSceneViewExtension;

UCLASS()
class UIndirectActorSubsystem : public UWorldSubsystem
{
	GENERATED_BODY()
public:
	TSharedPtr<FIndirectPopulateSceneViewExtension, ESPMode::ThreadSafe> SceneViewExtension;
	virtual void Initialize(FSubsystemCollectionBase& Collection) override;
	virtual void Deinitialize() override;

	// For demo
	TWeakObjectPtr<AIndirectActor> TheActor;
};