Files
UnrealEngine/Engine/Source/Runtime/Apple/MetalRHI/Private/MetalRayTracing.cpp
2025-05-18 13:04:45 +08:00

777 lines
30 KiB
C++

// Copyright Epic Games, Inc. All Rights Reserved.
/*==============================================================================
MetalRayTracing.cpp: MetalRT Implementation
==============================================================================*/
#include "MetalRayTracing.h"
#if METAL_RHI_RAYTRACING
#include "MetalRHIContext.h"
#include "MetalShaderTypes.h"
#include "BuiltInRayTracingShaders.h"
#include "RayTracingValidationShaders.h"
#include "RayTracingBuiltInResources.h"
static int32 GMetalRayTracingAllowCompaction = 1;
static FAutoConsoleVariableRef CVarMetalRayTracingAllowCompaction(
TEXT("r.Metal.RayTracing.AllowCompaction"),
GMetalRayTracingAllowCompaction,
TEXT("Whether to automatically perform compaction for static acceleration structures to save GPU memory. (default = 1)\n"),
ECVF_ReadOnly
);
static int32 GRayTracingDebugForceBuildMode = 0;
static FAutoConsoleVariableRef CVarMetalRayTracingDebugForceFastTrace(
TEXT("r.Metal.RayTracing.DebugForceBuildMode"),
GRayTracingDebugForceBuildMode,
TEXT("Forces specific acceleration structure build mode (not runtime-tweakable).\n")
TEXT("0: Use build mode requested by high-level code (Default)\n")
TEXT("1: Force fast build mode\n")
TEXT("2: Force fast trace mode\n"),
ECVF_ReadOnly
);
static int32 GMetalRayTracingMaxBatchedCompaction = 64;
static FAutoConsoleVariableRef CVarMetalRayTracingMaxBatchedCompaction(
TEXT("r.Metal.RayTracing.MaxBatchedCompaction"),
GMetalRayTracingMaxBatchedCompaction,
TEXT("Maximum of amount of compaction requests and rebuilds per frame. (default = 64)\n"),
ECVF_ReadOnly
);
static ERayTracingAccelerationStructureFlags GetRayTracingAccelerationStructureBuildFlags(const FRayTracingGeometryInitializer& Initializer)
{
ERayTracingAccelerationStructureFlags BuildFlags = ERayTracingAccelerationStructureFlags::None;
if (Initializer.bFastBuild)
{
BuildFlags = ERayTracingAccelerationStructureFlags::FastBuild;
}
else
{
BuildFlags = ERayTracingAccelerationStructureFlags::FastTrace;
}
if (Initializer.bAllowUpdate)
{
EnumAddFlags(BuildFlags, ERayTracingAccelerationStructureFlags::AllowUpdate);
}
if (!Initializer.bFastBuild && !Initializer.bAllowUpdate && Initializer.bAllowCompaction && GMetalRayTracingAllowCompaction)
{
EnumAddFlags(BuildFlags, ERayTracingAccelerationStructureFlags::AllowCompaction);
}
if (GRayTracingDebugForceBuildMode == 1)
{
EnumAddFlags(BuildFlags, ERayTracingAccelerationStructureFlags::FastBuild);
EnumRemoveFlags(BuildFlags, ERayTracingAccelerationStructureFlags::FastTrace);
}
else if (GRayTracingDebugForceBuildMode == 2)
{
EnumAddFlags(BuildFlags, ERayTracingAccelerationStructureFlags::FastTrace);
EnumRemoveFlags(BuildFlags, ERayTracingAccelerationStructureFlags::FastBuild);
}
return BuildFlags;
}
static bool ShouldCompactAfterBuild(ERayTracingAccelerationStructureFlags BuildFlags)
{
return EnumHasAllFlags(BuildFlags, ERayTracingAccelerationStructureFlags::AllowCompaction | ERayTracingAccelerationStructureFlags::FastTrace)
&& !EnumHasAnyFlags(BuildFlags, ERayTracingAccelerationStructureFlags::AllowUpdate);
}
// Manages all the pending BLAS compaction requests
class FMetalRayTracingCompactionRequestHandler
{
public:
UE_NONCOPYABLE(FMetalRayTracingCompactionRequestHandler)
FMetalRayTracingCompactionRequestHandler(FMetalDevice& DeviceContext);
~FMetalRayTracingCompactionRequestHandler();
void RequestCompact(FMetalRayTracingGeometry* InRTGeometry);
bool ReleaseRequest(FMetalRayTracingGeometry* InRTGeometry);
void Update(FMetalRHICommandContext& Context);
private:
/** Enqueued requests (waiting on size request submit). */
TArray<FMetalRayTracingGeometry*> PendingRequests;
/** Enqueued compaction requests (submitted compaction size requests; waiting on readback and actual compaction). */
TQueue<FMetalRayTracingGeometry*> ActiveRequests;
/** Num of active requests enqueued. */
uint32 NumActiveRequests;
/** Buffer used for compacted size readback. */
FMetalBuffer CompactedStructureSizeBuffer;
/** Size entry allocated in the CompactedStructureSizeBuffer (in element count). */
uint32_t SizeBufferMaxCapacity;
/** CompactedStructureSizeBuffer write index. Wrap once equal to SizeBufferMaxCapacity. */
uint32_t WriteIndex;
};
FMetalRayTracingCompactionRequestHandler::FMetalRayTracingCompactionRequestHandler(FMetalDevice& Device)
: SizeBufferMaxCapacity(GMetalRayTracingMaxBatchedCompaction)
, WriteIndex(0u)
{
PendingRequests.Reserve(GMetalRayTracingMaxBatchedCompaction);
CompactedStructureSizeBuffer = FMetalBuffer(Device.GetDevice().NewBuffer(GMetalRayTracingMaxBatchedCompaction * sizeof(uint32), MTL::ResourceStorageModeShared));
check(CompactedStructureSizeBuffer);
NumActiveRequests = 0;
}
FMetalRayTracingCompactionRequestHandler::~FMetalRayTracingCompactionRequestHandler()
{
check(PendingRequests.IsEmpty());
SafeReleaseMetalBuffer(CompactedStructureSizeBuffer);
}
void FMetalRayTracingCompactionRequestHandler::RequestCompact(FMetalRayTracingGeometry* InRTGeometry)
{
check(InRTGeometry->GetAccelerationStructureRead()
&& InRTGeometry->GetAccelerationStructureRead()->IsAccelerationStructure()
&& InRTGeometry->GetAccelerationStructureRead()->AccelerationStructureHandle.GetPtr());
ERayTracingAccelerationStructureFlags GeometryBuildFlags = GetRayTracingAccelerationStructureBuildFlags(InRTGeometry->Initializer);
check(EnumHasAllFlags(GeometryBuildFlags, ERayTracingAccelerationStructureFlags::AllowCompaction) &&
EnumHasAllFlags(GeometryBuildFlags, ERayTracingAccelerationStructureFlags::FastTrace) &&
!EnumHasAnyFlags(GeometryBuildFlags, ERayTracingAccelerationStructureFlags::AllowUpdate));
PendingRequests.Add(InRTGeometry);
}
bool FMetalRayTracingCompactionRequestHandler::ReleaseRequest(FMetalRayTracingGeometry* InRTGeometry)
{
PendingRequests.Remove(InRTGeometry);
return true;
}
void FMetalRayTracingCompactionRequestHandler::Update(FMetalRHICommandContext& Context)
{
// Early exit to avoid unecessary encoding breaks.
if (PendingRequests.IsEmpty() && ActiveRequests.IsEmpty())
{
return;
}
check(CompactedStructureSizeBuffer);
// Submit build commands.
MTL::Device* Device = Context.GetDevice();
FMetalRenderPass& RenderPass = DeviceContext->GetCurrentRenderPass();
FMetalCommandEncoder& Encoder = RenderPass.GetCurrentCommandEncoder();
Encoder.EndEncoding();
DeviceContext->GetCurrentState().SetStateDirty();
Encoder.BeginAccelerationStructureCommandEncoding();
MTL::AccelerationStructureCommandEncoder* CommandEncoder = Encoder.GetAccelerationStructureCommandEncoder();
check(CommandEncoder.GetPtr());
// Process pending requests.
for (FMetalRayTracingGeometry* Geometry : PendingRequests)
{
Geometry->CompactionSizeIndex = WriteIndex;
WriteIndex = (++WriteIndex % SizeBufferMaxCapacity);
CommandEncoder.WriteCompactedAccelerationStructureSize(Geometry->GetAccelerationStructureRead()->AccelerationStructureHandle, CompactedStructureSizeBuffer, Geometry->CompactionSizeIndex * sizeof(uint32_t));
ActiveRequests.Enqueue(Geometry);
NumActiveRequests++;
// enqueued enough requests for this update round
if (NumActiveRequests >= GMetalRayTracingMaxBatchedCompaction)
{
break;
}
}
PendingRequests.Empty();
// Try to readback active requests.
uint32_t* CompactedSizes = (uint32_t*)CompactedStructureSizeBuffer.GetContents();
// Process active requests.
while (!ActiveRequests.IsEmpty())
{
FMetalRayTracingGeometry* ActiveRequestsTail = *ActiveRequests.Peek();
check(ActiveRequestsTail);
if (!ActiveRequestsTail->bHasPendingCompactionRequests)
{
ActiveRequests.Pop();
continue;
}
// TODO: Assuming the compacted size may always be zero (e.g. corrupted BLAS/incorrect descriptor/API bug...), we should probably track the active readback index; and move forward if the index is stuck (after X number of frames)
// Break whenever we encounter the first incomplete request (i.e. size == 0).
uint32_t CompactedSize = CompactedSizes[ActiveRequestsTail->CompactionSizeIndex];
if (CompactedSize == 0u)
{
break;
}
MTL::AccelerationStructure* SrcBLAS = ActiveRequestsTail->GetAccelerationStructureRead()->AccelerationStructureHandle;
MTL::AccelerationStructure* CompactedBLAS = ActiveRequestsTail->GetAccelerationStructureWrite()->AccelerationStructureHandle;
check(CompactedBLAS);
CommandEncoder.CopyAndCompactAccelerationStructure(SrcBLAS, CompactedBLAS);
ActiveRequestsTail->NextAccelerationStructure();
ActiveRequests.Pop();
NumActiveRequests--;
}
Encoder.EndEncoding();
}
/** Fills a MTLPrimitiveAccelerationStructureDescriptor with infos provided by the UE5 geometry descriptor.
* This function assumes that GeometryDescriptors has already been allocated, and that you are responsible of its lifetime.
*/
static void FillPrimitiveAccelerationStructureDesc(MTL::PrimitiveAccelerationStructureDescriptor* AccelerationStructureDescriptor, const FRayTracingGeometryInitializer& Initializer, NSMutableArray<MTLAccelerationStructureGeometryDescriptor*>*& GeometryDescriptors)
{
// Populate Segments Descriptors.
FMetalRHIBuffer* IndexBuffer = ResourceCast(Initializer.IndexBuffer.GetReference());
int32 SegmentIndex = 0;
for (const FRayTracingGeometrySegment& Segment : Initializer.Segments)
{
check(Segment.NumPrimitives > 0);
// Vertex Buffer Infos
FMetalRHIBuffer* VertexBuffer = ResourceCast(Segment.VertexBuffer.GetReference());
check(VertexBuffer);
MTL::AccelerationStructureTriangleGeometryDescriptor GeometryDescriptor = MTL::AccelerationStructureTriangleGeometryDescriptor();
GeometryDescriptor.SetOpaque(Segment.bForceOpaque);
GeometryDescriptor.SetTriangleCount((Segment.bEnabled) ? Segment.NumPrimitives : 0);
GeometryDescriptor.SetAllowDuplicateIntersectionFunctionInvocation(Segment.bAllowDuplicateAnyHitShaderInvocation);
// Index Buffer Infos
if (IndexBuffer != nullptr)
{
const FMetalBuffer& IndexBufferRes = IndexBuffer->GetCurrentBuffer();
GeometryDescriptor.SetIndexType(IndexBuffer->GetIndexType());
GeometryDescriptor.SetIndexBuffer(IndexBufferRes);
GeometryDescriptor.SetIndexBufferOffset(IndexBufferRes.GetOffset() + Initializer.IndexBufferOffset);
}
const FMetalBuffer& VertexBufferRes = VertexBuffer->GetCurrentBuffer();
GeometryDescriptor.SetVertexBuffer(VertexBufferRes);
GeometryDescriptor.SetVertexBufferOffset(VertexBufferRes.GetOffset() + Segment.VertexBufferOffset);
GeometryDescriptor.SetVertexBufferStride(Segment.VertexBufferStride);
[GeometryDescriptors addObject: GeometryDescriptor];
SegmentIndex++;
}
// Populate Acceleration Structure Descriptor.
MTL::AccelerationStructureUsage Usage = MTL::AccelerationStructureUsage::None;
if (Initializer.bAllowUpdate)
Usage = MTL::AccelerationStructureUsageRefit;
else if (Initializer.bFastBuild)
Usage = MTL::AccelerationStructureUsagePreferFastBuild;
AccelerationStructureDescriptor.SetUsage(Usage);
AccelerationStructureDescriptor.SetGeometryDescriptors((__bridge NSArray*)GeometryDescriptors);
// Explicity retain the descriptor (will be re-used for refit and compaction).
[AccelerationStructureDescriptor retain];
}
static FRayTracingAccelerationStructureSize CalcRayTracingGeometrySize(FMetalDevice& Device, MTL::AccelerationStructureDescriptor* AccelerationStructureDescriptor)
{
MTL::AccelerationStructureSizes DescriptorSize = Device.AccelerationStructureSizesWithDescriptor(AccelerationStructureDescriptor);
FRayTracingAccelerationStructureSize SizeInfo = {};
SizeInfo.ResultSize = Align(DescriptorSize.accelerationStructureSize, GRHIRayTracingAccelerationStructureAlignment);
SizeInfo.BuildScratchSize = Align(DescriptorSize.buildScratchBufferSize, GRHIRayTracingScratchBufferAlignment);
SizeInfo.UpdateScratchSize = Align(DescriptorSize.refitScratchBufferSize, GRHIRayTracingScratchBufferAlignment);
return SizeInfo;
}
FRayTracingAccelerationStructureSize FMetalDynamicRHI::RHICalcRayTracingGeometrySize(const FRayTracingGeometryInitializer& Initializer)
{
MTL_SCOPED_AUTORELEASE_POOL;
MTL::PrimitiveAccelerationStructureDescriptor AccelerationStructureDescriptor = MTL::PrimitiveAccelerationStructureDescriptor();
NSMutableArray<MTLAccelerationStructureGeometryDescriptor*>* GeometryDescriptors = [[NSMutableArray<MTLAccelerationStructureGeometryDescriptor*> new] init];
FillPrimitiveAccelerationStructureDesc(AccelerationStructureDescriptor, Initializer, GeometryDescriptors);
[GeometryDescriptors release];
return CalcRayTracingGeometrySize(AccelerationStructureDescriptor);
}
FRayTracingAccelerationStructureSize FMetalDynamicRHI::RHICalcRayTracingSceneSize(const FRayTracingSceneInitializer& Initializer)
{
// TODO: Do we need to take in account the flags provided by the function call?
// TODO: Can we get away with the instance count only? (works on AS; what about AMD?)
MTL_SCOPED_AUTORELEASE_POOL;
MTL::InstanceAccelerationStructureDescriptor InstanceDescriptor = MTL::InstanceAccelerationStructureDescriptor();
InstanceDescriptor.SetInstanceCount(Initializer.MaxNumInstances);
return CalcRayTracingGeometrySize(InstanceDescriptor);
}
FMetalRayTracingGeometry::FMetalRayTracingGeometry(FRHICommandListBase& RHICmdList, const FRayTracingGeometryInitializer& InInitializer)
: FRHIRayTracingGeometry(InInitializer)
, bHasPendingCompactionRequests(false)
{
uint32 IndexBufferStride = 0;
if (Initializer.IndexBuffer)
{
// In case index buffer in initializer is not yet in valid state during streaming we assume the geometry is using UINT32 format.
IndexBufferStride = Initializer.IndexBuffer->GetSize() > 0
? Initializer.IndexBuffer->GetStride()
: 4;
}
checkf(!Initializer.IndexBuffer || (IndexBufferStride == 2 || IndexBufferStride == 4), TEXT("Index buffer must be 16 or 32 bit if in use."));
GeomArray = [NSMutableArray arrayWithCapacity:Initializer.Segments.Num()];
AccelerationStructureDescriptor = MTL::PrimitiveAccelerationStructureDescriptor();
RebuildDescriptors();
// NOTE: We do not use the RHI API in order to avoid re-filling another descriptor.
SizeInfo = CalcRayTracingGeometrySize(AccelerationStructureDescriptor);
AccelerationStructureIndex = 0;
// If this RayTracingGeometry going to be used as streaming destination
// we don't want to allocate its memory as it will be replaced later by streamed version
// but we still need correct SizeInfo as it is used to estimate its memory requirements outside of RHI.
if (Initializer.Type == ERayTracingGeometryInitializerType::StreamingDestination)
{
return;
}
FString DebugNameString = Initializer.DebugName.ToString();
const FRHIBufferCreateDesc CreateDesc =
FRHIBufferCreateDesc::Create(*DebugNameString, SizeInfo.ResultSize, 0, EBufferUsageFlags::AccelerationStructure)
.SetInitialState(ERHIAccess::BVHWrite);
for (uint32 i = 0; i < MaxNumAccelerationStructure; i++)
{
AccelerationStructure[i] = ResourceCast(RHICmdList.CreateBuffer(CreateDesc).GetReference());
check(AccelerationStructure[i]);
AccelerationStructure[i]->AccelerationStructureHandle.SetLabel(TCHAR_TO_ANSI(*DebugNameString));
}
}
FMetalRayTracingGeometry::~FMetalRayTracingGeometry()
{
ReleaseUnderlyingResource();
}
void FMetalRayTracingGeometry::ReleaseUnderlyingResource()
{
[GeomArray removeAllObjects];
[GeomArray release];
for (uint32 i = 0; i < MaxNumAccelerationStructure; i++)
{
AccelerationStructure[i].SafeRelease();
AccelerationStructure[i] = nullptr;
}
}
void FMetalRayTracingGeometry::Swap(FMetalRayTracingGeometry& Other)
{
::Swap(AccelerationStructureDescriptor, Other.AccelerationStructureDescriptor);
for (uint32 i = 0; i < MaxNumAccelerationStructure; i++)
{
::Swap(AccelerationStructure[i], Other.AccelerationStructure[i]);
}
::Swap(AccelerationStructureIndex, Other.AccelerationStructureIndex);
Initializer = Other.Initializer;
// HitGroup Parameters Update is handled by the Scene
}
void FMetalRayTracingGeometry::RemoveCompactionRequest()
{
if (bHasPendingCompactionRequests)
{
check(GetAccelerationStructureRead());
bool bRequestFound = GetMetalDeviceContext().GetRayTracingCompactionRequestHandler()->ReleaseRequest(this);
check(bRequestFound);
bHasPendingCompactionRequests = false;
}
}
void FMetalRayTracingGeometry::RebuildDescriptors()
{
[GeomArray removeAllObjects];
AccelerationStructureDescriptor = MTL::PrimitiveAccelerationStructureDescriptor();
FillPrimitiveAccelerationStructureDesc(AccelerationStructureDescriptor, Initializer, GeomArray);
}
FMetalRayTracingScene::FMetalRayTracingScene(FRayTracingSceneInitializer InInitializer)
: Initializer(MoveTemp(InInitializer))
{
MTL::InstanceAccelerationStructureDescriptor* InstanceDescriptor;
InstanceDescriptor.SetInstanceCount(Initializer.NumNativeInstances);
SizeInfo = CalcRayTracingGeometrySize(InstanceDescriptor);
MutableAccelerationStructures = [[NSMutableArray<id<MTLAccelerationStructure>> new] init];
}
FMetalRayTracingScene::~FMetalRayTracingScene()
{
AccelerationStructureBuffer.SafeRelease();
InstanceBufferSRV.SafeRelease();
[MutableAccelerationStructures removeAllObjects];
[MutableAccelerationStructures release];
}
void FMetalRayTracingScene::BindBuffer(FRHIBuffer* InBuffer, uint32 InBufferOffset)
{
check(IsInRHIThread() || !IsRunningRHIInSeparateThread());
check(SizeInfo.ResultSize + InBufferOffset <= InBuffer->GetSize());
AccelerationStructureBuffer = ResourceCast(InBuffer);
FMetalDeviceContext& Context = GetMetalDeviceContext();
MTL::Device* Device = Context.GetDevice();
{
checkf(ShaderResourceView == nullptr, TEXT("Binding multiple buffers is not currently supported."));
check(InBufferOffset % GRHIRayTracingAccelerationStructureAlignment == 0);
check(AccelerationStructureBuffer->IsAccelerationStructure());
FShaderResourceViewInitializer ViewInitializer(AccelerationStructureBuffer, InBufferOffset, 0);
ShaderResourceView = new FMetalShaderResourceView(ViewInitializer);
FString DebugNameString = Initializer.DebugName.ToString();
DebugNameString = (DebugNameString.IsEmpty()) ? TEXT("TLAS") : DebugNameString;
AccelerationStructureBuffer->AccelerationStructureHandle.SetLabel(TCHAR_TO_ANSI(*DebugNameString));
}
}
void FMetalRayTracingScene::BuildAccelerationStructure(
FMetalRHICommandContext& CommandContext,
FMetalRHIBuffer* InScratchBuffer, uint32 ScratchOffset,
FMetalRHIBuffer* InstanceBuffer, uint32 InstanceOffset)
{
check(AccelerationStructureBuffer.IsValid());
check(InstanceBuffer != nullptr);
FMetalBuffer CurInstanceBuffer = InstanceBuffer->GetCurrentBuffer();
check(CurInstanceBuffer);
uint32 InstanceBufferOffset = InstanceOffset + static_cast<uint32>(CurInstanceBuffer.GetOffset());
// Create SRV first (since we collect BLAS to map in BuildPerInstanceGeometryParameterBuffer()).
FShaderResourceViewInitializer ViewInitializer(InstanceBuffer, InstanceOffset, 0);
InstanceBufferSRV = new FMetalShaderResourceView(ViewInitializer);
TRefCountPtr<FMetalRHIBuffer> ScratchBuffer;
if (InScratchBuffer == nullptr)
{
const FRHIBufferCreateDesc CreateDesc =
FRHIBufferCreateDesc::CreateStructured(TEXT("BuildScratchTLAS"), SizeInfo.BuildScratchSize, 0)
.AddUsage(EBufferUsageFlags::UnorderedAccess)
.SetInitialState(ERHIAccess::UAVCompute);
ScratchBuffer = ResourceCast(RHICreateBuffer(CreateDesc).GetReference());
InScratchBuffer = ScratchBuffer.GetReference();
ScratchOffset = 0;
}
FMetalDeviceContext& Context = GetMetalDeviceContext();
MTL::Device* Device = Context.GetDevice();
// Reset current renderpass to kick off acceleration structures build
FMetalRenderPass& RenderPass = Context.GetCurrentRenderPass();
FMetalCommandEncoder& Encoder = RenderPass.GetCurrentCommandEncoder();
// Crappy workaround for inline raytracing: we must bind the TLAS instances descriptors to emulate a missing intrinsic (GetBindingRecordOffset).
// Since inline RT uses the regular compute pipeline, we must do the binding in the regular path (in FMetalRenderPass).
// Let's hope we'll be able to find a better way later...
RenderPass.SetRayTracingInstanceBufferSRV(InstanceBufferSRV);
Encoder.EndEncoding();
Context.GetCurrentState().SetStateDirty();
Encoder.BeginAccelerationStructureCommandEncoding();
MTL::AccelerationStructureCommandEncoder* CommandEncoder = Encoder.GetAccelerationStructureCommandEncoder();
check(CommandEncoder.GetPtr());
FMetalBuffer CurScratchBuffer = InScratchBuffer->GetCurrentBuffer();
check(CurScratchBuffer);
{
MTL::InstanceAccelerationStructureDescriptor* InstanceDescriptor = MTL::InstanceAccelerationStructureDescriptor();
InstanceDescriptor.SetInstanceCount(Initializer.NumNativeInstances);
InstanceDescriptor.SetInstanceDescriptorBuffer(CurInstanceBuffer);
InstanceDescriptor.SetInstanceDescriptorBufferOffset(InstanceBufferOffset);
InstanceDescriptor.SetInstancedAccelerationStructures((__bridge NSArray*)MutableAccelerationStructures);
InstanceDescriptor.SetInstanceDescriptorStride(GRHIRayTracingInstanceDescriptorSize);
InstanceDescriptor.SetInstanceDescriptorType(MTL::AccelerationStructureInstanceDescriptorType::UserID);
MTL::AccelerationStructure* AS = ResourceCast(ShaderResourceView->GetBuffer())->AccelerationStructureHandle;
CommandEncoder.BuildAccelerationStructure(AS, InstanceDescriptor, CurScratchBuffer, ScratchOffset);
}
Encoder.EndEncoding();
}
void FMetalRHICommandContext::RHIBuildAccelerationStructure(TConstArrayView<FRayTracingSceneBuildParams> Params)
{
for (const FRayTracingSceneBuildParams& SceneBuildParams : Params)
{
FMetalRayTracingScene* const Scene = ResourceCast(SceneBuildParams.Scene);
FMetalRHIBuffer* const ScratchBuffer = ResourceCast(SceneBuildParams.ScratchBuffer);
FMetalRHIBuffer* const InstanceBuffer = ResourceCast(SceneBuildParams.InstanceBuffer);
Scene->BuildAccelerationStructure(
*this,
ScratchBuffer, SceneBuildParams.ScratchBufferOffset,
InstanceBuffer, SceneBuildParams.InstanceBufferOffset);
}
}
void FMetalRHICommandContext::RHIBuildAccelerationStructures(TConstArrayView<FRayTracingGeometryBuildParams> Params, const FRHIBufferRange& ScratchBufferRange)
{
checkf(ScratchBufferRange.Buffer != nullptr, TEXT("BuildAccelerationStructures requires valid scratch buffer"));
// Update geometry vertex buffers
for (const FRayTracingGeometryBuildParams& P : Params)
{
FMetalRayTracingGeometry* const Geometry = ResourceCast(P.Geometry.GetReference());
if (P.Segments.Num())
{
checkf(P.Segments.Num() == Geometry->Initializer.Segments.Num(),
TEXT("If updated segments are provided, they must exactly match existing geometry segments. Only vertex buffer bindings may change."));
for (int32 i = 0; i < P.Segments.Num(); ++i)
{
checkf(P.Segments[i].MaxVertices <= Geometry->Initializer.Segments[i].MaxVertices,
TEXT("Maximum number of vertices in a segment (%u) must not be smaller than what was declared during FRHIRayTracingGeometry creation (%u), as this controls BLAS memory allocation."),
P.Segments[i].MaxVertices, Geometry->Initializer.Segments[i].MaxVertices
);
Geometry->Initializer.Segments[i].VertexBuffer = P.Segments[i].VertexBuffer;
Geometry->Initializer.Segments[i].VertexBufferElementType = P.Segments[i].VertexBufferElementType;
Geometry->Initializer.Segments[i].VertexBufferStride = P.Segments[i].VertexBufferStride;
Geometry->Initializer.Segments[i].VertexBufferOffset = P.Segments[i].VertexBufferOffset;
}
// We must update the descriptor if any segments have changed.
Geometry->RebuildDescriptors();
}
}
uint32 ScratchBufferSize = ScratchBufferRange.Size ? ScratchBufferRange.Size : ScratchBufferRange.Buffer->GetSize();
checkf(ScratchBufferSize + ScratchBufferRange.Offset <= ScratchBufferRange.Buffer->GetSize(),
TEXT("BLAS scratch buffer range size is %lld bytes with offset %lld, but the buffer only has %lld bytes. "),
ScratchBufferRange.Size, ScratchBufferRange.Offset, ScratchBufferRange.Buffer->GetSize());
const uint64 ScratchAlignment = GRHIRayTracingScratchBufferAlignment;
FMetalRHIBuffer* ScratchBuffer = ResourceCast(ScratchBufferRange.Buffer);
uint32 ScratchBufferOffset = static_cast<uint32>(ScratchBufferRange.Offset);
MTL::Device* Device = Context->GetDevice();
TArray<TTuple<FMetalRayTracingGeometry*, uint32>, TInlineAllocator<32>> GeometryToBuild;
TArray<TTuple<FMetalRayTracingGeometry*, uint32>, TInlineAllocator<32>> GeometryToRefit;
GeometryToBuild.Reserve(Params.Num());
GeometryToRefit.Reserve(Params.Num());
for (const FRayTracingGeometryBuildParams& P : Params)
{
FMetalRayTracingGeometry* const Geometry = ResourceCast(P.Geometry.GetReference());
const bool bIsUpdate = P.BuildMode == EAccelerationStructureBuildMode::Update;
uint64 ScratchBufferRequiredSize = bIsUpdate ? Geometry->SizeInfo.UpdateScratchSize : Geometry->SizeInfo.BuildScratchSize;
checkf(ScratchBufferRequiredSize + ScratchBufferOffset <= ScratchBufferSize,
TEXT("BLAS scratch buffer size is %ld bytes with offset %ld (%ld bytes available), but the build requires %lld bytes. "),
ScratchBufferSize, ScratchBufferOffset, ScratchBufferSize - ScratchBufferOffset, ScratchBufferRequiredSize);
if (!bIsUpdate)
{
GeometryToBuild.Add(TTuple<FMetalRayTracingGeometry*, uint32>(Geometry, ScratchBufferOffset));
}
else
{
GeometryToRefit.Add(TTuple<FMetalRayTracingGeometry*, uint32>(Geometry, ScratchBufferOffset));
}
ScratchBufferOffset = Align(ScratchBufferOffset + ScratchBufferRequiredSize, ScratchAlignment);
// TODO: Add a CVAR to toggle validation (e.g. r.Metal.RayTracingValidate).
//FRayTracingValidateGeometryBuildParamsCS::Dispatch(FRHICommandListExecutor::GetImmediateCommandList(), P);
}
FMetalBuffer ScratchBufferRes = ScratchBuffer->GetCurrentBufferOrNil();
check(ScratchBufferRes);
// Submit build commands.
FMetalRenderPass& RenderPass = Context->GetCurrentRenderPass();
FMetalCommandEncoder& Encoder = RenderPass.GetCurrentCommandEncoder();
Encoder.EndEncoding();
Context->GetCurrentState().SetStateDirty();
Encoder.BeginAccelerationStructureCommandEncoding();
MTL::AccelerationStructureCommandEncoder* CommandEncoder = Encoder.GetAccelerationStructureCommandEncoder();
check(CommandEncoder.GetPtr());
for (TTuple<FMetalRayTracingGeometry*, uint32>& BuildRequest : GeometryToBuild)
{
FMetalRayTracingGeometry* Geometry = BuildRequest.Key;
uint32 ScratchOffset = BuildRequest.Value;
CommandEncoder.BuildAccelerationStructure(
Geometry->GetAccelerationStructureRead()->AccelerationStructureHandle,
Geometry->AccelerationStructureDescriptor,
ScratchBufferRes,
ScratchOffset
);
}
for (TTuple<FMetalRayTracingGeometry*, uint32>& RefitRequest : GeometryToRefit)
{
FMetalRayTracingGeometry* Geometry = RefitRequest.Key;
uint32 ScratchOffset = RefitRequest.Value;
MTL::AccelerationStructure* SrcBLAS = Geometry->GetAccelerationStructureRead()->AccelerationStructureHandle;
MTL::AccelerationStructure* DstBLAS = Geometry->GetAccelerationStructureWrite()->AccelerationStructureHandle;
CommandEncoder.RefitAccelerationStructure(
SrcBLAS,
Geometry->AccelerationStructureDescriptor,
&DstBLAS,
ScratchBufferRes,
ScratchOffset
);
Geometry->NextAccelerationStructure();
}
Encoder.EndEncoding();
for (const FRayTracingGeometryBuildParams& P : Params)
{
FMetalRayTracingGeometry* const Geometry = ResourceCast(P.Geometry.GetReference());
const bool bIsUpdate = P.BuildMode == EAccelerationStructureBuildMode::Update;
if (!bIsUpdate)
{
ERayTracingAccelerationStructureFlags GeometryBuildFlags = GetRayTracingAccelerationStructureBuildFlags(Geometry->Initializer);
if (ShouldCompactAfterBuild(GeometryBuildFlags))
{
GetMetalDeviceContext().GetRayTracingCompactionRequestHandler()->RequestCompact(Geometry);
Geometry->bHasPendingCompactionRequests = true;
}
}
}
}
void FMetalRHICommandContext::RHIBindAccelerationStructureMemory(FRHIRayTracingScene* InScene, FRHIBuffer* InBuffer, uint32 InBufferOffset)
{
FMetalRayTracingScene* MetalScene = ResourceCast(InScene);
MetalScene->BindBuffer(InBuffer, InBufferOffset);
}
void FMetalRHICommandContext::RHIClearShaderBindingTable(FRHIShaderBindingTable* SBT)
{
// TODO:
}
void FMetalRHICommandContext::RHIRayTraceDispatch(FRHIRayTracingPipelineState* RayTracingPipelineState, FRHIRayTracingShader* RayGenShader,
FRHIRayTracingScene* SceneRHI,
const FRayTracingShaderBindings& GlobalResourceBindings,
uint32 Width, uint32 Height)
{
checkNoEntry();
}
void FMetalRHICommandContext::RHIRayTraceDispatchIndirect(FRHIRayTracingPipelineState* RayTracingPipelineState, FRHIRayTracingShader* RayGenShader,
FRHIRayTracingScene* SceneRHI,
const FRayTracingShaderBindings& GlobalResourceBindings,
FRHIBuffer* ArgumentBuffer, uint32 ArgumentOffset)
{
checkNoEntry();
}
void FMetalRHICommandContext::RHISetRayTracingBindings(
FRHIRayTracingScene* InScene, FRHIRayTracingPipelineState* InPipeline,
uint32 NumBindings, const FRayTracingLocalShaderBindings* Bindings,
ERayTracingBindingType BindingType)
{
checkNoEntry();
}
void FMetalRHICommandContext::RHISetBindingsOnShaderBindingTable(
FRHIShaderBindingTable* SBT, FRHIRayTracingPipelineState* InPipeline,
uint32 NumBindings, const FRayTracingLocalShaderBindings* Bindings,
ERayTracingBindingType BindingType)
{
checkNoEntry();
}
FRayTracingSceneRHIRef FMetalDynamicRHI::RHICreateRayTracingScene(FRayTracingSceneInitializer Initializer)
{
MTL_SCOPED_AUTORELEASE_POOL;
return new FMetalRayTracingScene(MoveTemp(Initializer));
}
FRayTracingGeometryRHIRef FMetalDynamicRHI::RHICreateRayTracingGeometry(FRHICommandListBase& RHICmdList, const FRayTracingGeometryInitializer& Initializer)
{
MTL_SCOPED_AUTORELEASE_POOL;
return new FMetalRayTracingGeometry(RHICmdList, Initializer);
}
FRayTracingPipelineStateRHIRef FMetalDynamicRHI::RHICreateRayTracingPipelineState(const FRayTracingPipelineStateInitializer& Initializer)
{
checkNoEntry();
return nullptr;
}
FShaderBindingTableRHIRef FMetalDynamicRHI::RHICreateShaderBindingTable(FRHICommandListBase& RHICmdList, const FRayTracingShaderBindingTableInitializer& Initializer)
{
checkNoEntry();
return nullptr;
}
void FMetalDevice::InitializeRayTracing()
{
// Explicitly request a pointer to the DeviceContext since the CompactionHandler
// is initialized before the global getter is setup.
RayTracingCompactionRequestHandler = new FMetalRayTracingCompactionRequestHandler(this);
}
void FMetalDevice::UpdateRayTracing()
{
RayTracingCompactionRequestHandler->Update(this);
}
void FMetalDevice::CleanUpRayTracing()
{
delete RayTracingCompactionRequestHandler;
}
#endif // METAL_RHI_RAYTRACING