Add GPU sorting for gaussian splat rendering

engine/src/Rendering/Passes/BuiltinGaussianSplatPass.cpp

@@ -158,6 +158,27 @@ RHI::RHIResourceView* ResolveWorkingSetView(
     }
 }
 
+bool TransitionWorkingSetBuffer(
+    RHI::RHICommandList* commandList,
+    Internal::BuiltinGaussianSplatPassResources::CachedBufferView& bufferView,
+    RHI::ResourceStates targetState) {
+    if (commandList == nullptr || bufferView.currentState == targetState) {
+        return true;
+    }
+
+    RHI::RHIResourceView* resourceView =
+        targetState == RHI::ResourceStates::UnorderedAccess
+            ? bufferView.unorderedAccessView
+            : bufferView.shaderResourceView;
+    if (resourceView == nullptr) {
+        return false;
+    }
+
+    commandList->TransitionBarrier(resourceView, bufferView.currentState, targetState);
+    bufferView.currentState = targetState;
+    return true;
+}
+
 template <typename BindFn>
 void BindDescriptorSetRanges(
     Core::uint32 firstDescriptorSet,
@@ -309,6 +330,13 @@ bool BuiltinGaussianSplatPass::Execute(const RenderPassContext& context) {
             return false;
         }
 
+        if (!SortVisibleGaussianSplat(
+                context.renderContext,
+                context.sceneData,
+                visibleGaussianSplat)) {
+            return false;
+        }
+
         if (!DrawVisibleGaussianSplat(
                 context.renderContext,
                 context.surface,
@@ -479,6 +507,26 @@ BuiltinGaussianSplatPass::ResolvedShaderPass BuiltinGaussianSplatPass::ResolvePr
     return resolved;
 }
 
+BuiltinGaussianSplatPass::ResolvedShaderPass BuiltinGaussianSplatPass::ResolveBitonicSortShaderPass(
+    const RenderSceneData& sceneData) const {
+    ResolvedShaderPass resolved = {};
+    if (!m_builtinGaussianSplatUtilitiesShader.IsValid()) {
+        return resolved;
+    }
+
+    const Resources::Shader* shader = m_builtinGaussianSplatUtilitiesShader.Get();
+    const Resources::ShaderBackend backend = ::XCEngine::Rendering::Internal::ToShaderBackend(m_backendType);
+    const Containers::String passName("GaussianSplatBitonicSort");
+    if (const Resources::ShaderPass* shaderPass =
+            FindCompatibleComputePass(*shader, passName, sceneData.globalShaderKeywords, backend)) {
+        resolved.shader = shader;
+        resolved.pass = shaderPass;
+        resolved.passName = passName;
+    }
+
+    return resolved;
+}
+
 BuiltinGaussianSplatPass::PassResourceLayout* BuiltinGaussianSplatPass::GetOrCreatePassResourceLayout(
     const RenderContext& context,
     const ResolvedShaderPass& resolvedShaderPass,
@@ -555,6 +603,12 @@ BuiltinGaussianSplatPass::PassResourceLayout* BuiltinGaussianSplatPass::GetOrCre
             return failLayout(
                 "BuiltinGaussianSplatPass prepare-order pass requires sort distance, order, position, other, color, and view-data gaussian splat buffer bindings");
         }
+    } else if (usage == PassLayoutUsage::BitonicSort) {
+        if (!passLayout.gaussianSplatSortDistanceBuffer.IsValid() ||
+            !passLayout.gaussianSplatOrderBuffer.IsValid()) {
+            return failLayout(
+                "BuiltinGaussianSplatPass bitonic-sort pass requires sort distance and order gaussian splat buffer bindings");
+        }
     }
 
     std::vector<RHI::DescriptorSetLayoutDesc> nativeSetLayouts(passLayout.setLayouts.size());
@@ -640,8 +694,9 @@ RHI::RHIPipelineState* BuiltinGaussianSplatPass::GetOrCreatePipelineState(
 
 RHI::RHIPipelineState* BuiltinGaussianSplatPass::GetOrCreateComputePipelineState(
     const RenderContext& context,
-    const RenderSceneData& sceneData) {
-    const ResolvedShaderPass resolvedShaderPass = ResolvePrepareOrderShaderPass(sceneData);
+    const ResolvedShaderPass& resolvedShaderPass,
+    PassLayoutUsage usage,
+    const Resources::ShaderKeywordSet& keywordSet) {
     if (resolvedShaderPass.shader == nullptr || resolvedShaderPass.pass == nullptr) {
         return nullptr;
     }
@@ -649,7 +704,7 @@ RHI::RHIPipelineState* BuiltinGaussianSplatPass::GetOrCreateComputePipelineState
     PassResourceLayout* passLayout = GetOrCreatePassResourceLayout(
         context,
         resolvedShaderPass,
-        PassLayoutUsage::PrepareOrder);
+        usage);
     if (passLayout == nullptr || passLayout->pipelineLayout == nullptr) {
         return nullptr;
     }
@@ -658,7 +713,7 @@ RHI::RHIPipelineState* BuiltinGaussianSplatPass::GetOrCreateComputePipelineState
     pipelineKey.shader = resolvedShaderPass.shader;
     pipelineKey.passName = resolvedShaderPass.passName;
     pipelineKey.keywordSignature =
-        ::XCEngine::Rendering::Internal::BuildShaderKeywordSignature(sceneData.globalShaderKeywords);
+        ::XCEngine::Rendering::Internal::BuildShaderKeywordSignature(keywordSet);
 
     const auto existing = m_computePipelineStates.find(pipelineKey);
     if (existing != m_computePipelineStates.end()) {
@@ -672,7 +727,7 @@ RHI::RHIPipelineState* BuiltinGaussianSplatPass::GetOrCreateComputePipelineState
             *resolvedShaderPass.shader,
             *resolvedShaderPass.pass,
             resolvedShaderPass.passName,
-            sceneData.globalShaderKeywords);
+            keywordSet);
     RHI::RHIPipelineState* pipelineState = context.device->CreateComputePipelineState(pipelineDesc);
     if (pipelineState == nullptr || !pipelineState->IsValid()) {
         const Containers::String error =
@@ -985,35 +1040,20 @@ bool BuiltinGaussianSplatPass::PrepareVisibleGaussianSplat(
         context,
         resolvedShaderPass,
         PassLayoutUsage::PrepareOrder);
-    RHI::RHIPipelineState* pipelineState = GetOrCreateComputePipelineState(context, sceneData);
+    RHI::RHIPipelineState* pipelineState = GetOrCreateComputePipelineState(
+        context,
+        resolvedShaderPass,
+        PassLayoutUsage::PrepareOrder,
+        sceneData.globalShaderKeywords);
     if (passLayout == nullptr || pipelineState == nullptr) {
         return fail("BuiltinGaussianSplatPass prepare-order failed: compute pipeline setup was not created");
     }
 
     RHI::RHICommandList* commandList = context.commandList;
-
-    if (workingSet->sortDistances.currentState != RHI::ResourceStates::UnorderedAccess) {
-        commandList->TransitionBarrier(
-            workingSet->sortDistances.unorderedAccessView,
-            workingSet->sortDistances.currentState,
-            RHI::ResourceStates::UnorderedAccess);
-        workingSet->sortDistances.currentState = RHI::ResourceStates::UnorderedAccess;
-    }
-
-    if (workingSet->orderIndices.currentState != RHI::ResourceStates::UnorderedAccess) {
-        commandList->TransitionBarrier(
-            workingSet->orderIndices.unorderedAccessView,
-            workingSet->orderIndices.currentState,
-            RHI::ResourceStates::UnorderedAccess);
-        workingSet->orderIndices.currentState = RHI::ResourceStates::UnorderedAccess;
-    }
-
-    if (workingSet->viewData.currentState != RHI::ResourceStates::UnorderedAccess) {
-        commandList->TransitionBarrier(
-            workingSet->viewData.unorderedAccessView,
-            workingSet->viewData.currentState,
-            RHI::ResourceStates::UnorderedAccess);
-        workingSet->viewData.currentState = RHI::ResourceStates::UnorderedAccess;
+    if (!TransitionWorkingSetBuffer(commandList, workingSet->sortDistances, RHI::ResourceStates::UnorderedAccess) ||
+        !TransitionWorkingSetBuffer(commandList, workingSet->orderIndices, RHI::ResourceStates::UnorderedAccess) ||
+        !TransitionWorkingSetBuffer(commandList, workingSet->viewData, RHI::ResourceStates::UnorderedAccess)) {
+        return fail("BuiltinGaussianSplatPass prepare-order failed: resource transition failed");
     }
 
     commandList->SetPipelineState(pipelineState);
@@ -1029,7 +1069,11 @@ bool BuiltinGaussianSplatPass::PrepareVisibleGaussianSplat(
             static_cast<float>(sceneData.cameraData.viewportHeight),
             sceneData.cameraData.viewportWidth > 0u ? 1.0f / static_cast<float>(sceneData.cameraData.viewportWidth) : 0.0f,
             sceneData.cameraData.viewportHeight > 0u ? 1.0f / static_cast<float>(sceneData.cameraData.viewportHeight) : 0.0f),
-        Math::Vector4(static_cast<float>(cachedGaussianSplat->splatCount), 0.0f, 0.0f, 0.0f)
+        Math::Vector4(
+            static_cast<float>(cachedGaussianSplat->splatCount),
+            static_cast<float>(workingSet->sortCapacity),
+            0.0f,
+            0.0f)
     };
 
     if (passLayout->descriptorSetCount > 0u) {
@@ -1112,18 +1156,184 @@ bool BuiltinGaussianSplatPass::PrepareVisibleGaussianSplat(
             });
     }
 
-    commandList->Dispatch((cachedGaussianSplat->splatCount + 63u) / 64u, 1u, 1u);
+    commandList->Dispatch((workingSet->sortCapacity + 63u) / 64u, 1u, 1u);
+
+    if (!TransitionWorkingSetBuffer(commandList, workingSet->viewData, RHI::ResourceStates::NonPixelShaderResource)) {
+        return fail("BuiltinGaussianSplatPass prepare-order failed: view-data transition failed");
+    }
+    return true;
+}
+
+bool BuiltinGaussianSplatPass::SortVisibleGaussianSplat(
+    const RenderContext& context,
+    const RenderSceneData& sceneData,
+    const VisibleGaussianSplatItem& visibleGaussianSplat) {
+    auto fail = [](const char* message) -> bool {
+        Debug::Logger::Get().Error(Debug::LogCategory::Rendering, message);
+        return false;
+    };
+
+    if (visibleGaussianSplat.gameObject == nullptr ||
+        visibleGaussianSplat.gaussianSplat == nullptr ||
+        !visibleGaussianSplat.gaussianSplat->IsValid() ||
+        m_passResources == nullptr) {
+        return fail("BuiltinGaussianSplatPass bitonic-sort failed: invalid visible gaussian splat item");
+    }
+
+    const RenderResourceCache::CachedGaussianSplat* cachedGaussianSplat =
+        m_resourceCache.GetOrCreateGaussianSplat(m_device, visibleGaussianSplat.gaussianSplat);
+    if (cachedGaussianSplat == nullptr) {
+        return fail("BuiltinGaussianSplatPass bitonic-sort failed: gaussian splat GPU cache is missing");
+    }
+
+    Internal::BuiltinGaussianSplatPassResources::WorkingSet* workingSet = nullptr;
+    if (!m_passResources->EnsureWorkingSet(m_device, visibleGaussianSplat, workingSet) ||
+        workingSet == nullptr ||
+        workingSet->sortDistances.shaderResourceView == nullptr ||
+        workingSet->sortDistances.unorderedAccessView == nullptr ||
+        workingSet->orderIndices.shaderResourceView == nullptr ||
+        workingSet->orderIndices.unorderedAccessView == nullptr) {
+        return fail("BuiltinGaussianSplatPass bitonic-sort failed: working set allocation is incomplete");
+    }
+
+    RHI::RHICommandList* commandList = context.commandList;
+    if (workingSet->sortCapacity <= 1u) {
+        return TransitionWorkingSetBuffer(
+            commandList,
+            workingSet->orderIndices,
+            RHI::ResourceStates::NonPixelShaderResource);
+    }
+
+    const ResolvedShaderPass resolvedShaderPass = ResolveBitonicSortShaderPass(sceneData);
+    if (resolvedShaderPass.shader == nullptr || resolvedShaderPass.pass == nullptr) {
+        return fail("BuiltinGaussianSplatPass bitonic-sort failed: utilities shader pass was not resolved");
+    }
+
+    PassLayoutKey passLayoutKey = {};
+    passLayoutKey.shader = resolvedShaderPass.shader;
+    passLayoutKey.passName = resolvedShaderPass.passName;
+
+    PassResourceLayout* passLayout = GetOrCreatePassResourceLayout(
+        context,
+        resolvedShaderPass,
+        PassLayoutUsage::BitonicSort);
+    RHI::RHIPipelineState* pipelineState = GetOrCreateComputePipelineState(
+        context,
+        resolvedShaderPass,
+        PassLayoutUsage::BitonicSort,
+        sceneData.globalShaderKeywords);
+    if (passLayout == nullptr || pipelineState == nullptr) {
+        return fail("BuiltinGaussianSplatPass bitonic-sort failed: compute pipeline setup was not created");
+    }
+
+    commandList->SetPipelineState(pipelineState);
+
+    for (Core::uint32 levelMask = 2u; levelMask <= workingSet->sortCapacity; levelMask <<= 1u) {
+        for (Core::uint32 partnerMask = levelMask >> 1u; partnerMask > 0u; partnerMask >>= 1u) {
+            if (!TransitionWorkingSetBuffer(commandList, workingSet->sortDistances, RHI::ResourceStates::UnorderedAccess) ||
+                !TransitionWorkingSetBuffer(commandList, workingSet->orderIndices, RHI::ResourceStates::UnorderedAccess)) {
+                return fail("BuiltinGaussianSplatPass bitonic-sort failed: resource transition to UAV failed");
+            }
+
+            const PerObjectConstants perObjectConstants = {
+                sceneData.cameraData.projection,
+                sceneData.cameraData.view,
+                visibleGaussianSplat.localToWorld.Transpose(),
+                Math::Vector4(sceneData.cameraData.worldRight, 0.0f),
+                Math::Vector4(sceneData.cameraData.worldUp, 0.0f),
+                Math::Vector4(
+                    static_cast<float>(sceneData.cameraData.viewportWidth),
+                    static_cast<float>(sceneData.cameraData.viewportHeight),
+                    sceneData.cameraData.viewportWidth > 0u ? 1.0f / static_cast<float>(sceneData.cameraData.viewportWidth) : 0.0f,
+                    sceneData.cameraData.viewportHeight > 0u ? 1.0f / static_cast<float>(sceneData.cameraData.viewportHeight) : 0.0f),
+                Math::Vector4(
+                    static_cast<float>(cachedGaussianSplat->splatCount),
+                    static_cast<float>(workingSet->sortCapacity),
+                    static_cast<float>(partnerMask),
+                    static_cast<float>(levelMask))
+            };
+
+            if (passLayout->descriptorSetCount > 0u) {
+                std::vector<RHI::RHIDescriptorSet*> descriptorSets(passLayout->descriptorSetCount, nullptr);
+                for (Core::uint32 descriptorOffset = 0u; descriptorOffset < passLayout->descriptorSetCount; ++descriptorOffset) {
+                    const Core::uint32 setIndex = passLayout->firstDescriptorSet + descriptorOffset;
+                    if (setIndex >= passLayout->setLayouts.size()) {
+                        return fail("BuiltinGaussianSplatPass bitonic-sort failed: descriptor set index overflow");
+                    }
+
+                    const BuiltinPassSetLayoutMetadata& setLayout = passLayout->setLayouts[setIndex];
+                    if (setLayout.layout.bindingCount == 0u) {
+                        continue;
+                    }
+
+                    if (!(setLayout.usesPerObject ||
+                          setLayout.usesGaussianSplatSortDistanceBuffer ||
+                          setLayout.usesGaussianSplatOrderBuffer)) {
+                        return fail("BuiltinGaussianSplatPass bitonic-sort failed: unexpected descriptor set layout");
+                    }
+
+                    const Core::uint64 objectId =
+                        setLayout.usesPerObject ? visibleGaussianSplat.gameObject->GetID() : 0u;
+                    const Resources::GaussianSplat* gaussianSplatKey =
+                        (setLayout.usesGaussianSplatSortDistanceBuffer ||
+                         setLayout.usesGaussianSplatOrderBuffer)
+                            ? visibleGaussianSplat.gaussianSplat
+                            : nullptr;
+
+                    CachedDescriptorSet* cachedDescriptorSet = GetOrCreateDynamicDescriptorSet(
+                        passLayoutKey,
+                        *passLayout,
+                        setLayout,
+                        setIndex,
+                        objectId,
+                        visibleGaussianSplat.gaussianSplatRenderer,
+                        nullptr,
+                        gaussianSplatKey,
+                        MaterialConstantPayloadView(),
+                        *cachedGaussianSplat,
+                        workingSet->sortDistances.unorderedAccessView,
+                        workingSet->orderIndices.unorderedAccessView,
+                        nullptr);
+                    if (cachedDescriptorSet == nullptr || cachedDescriptorSet->descriptorSet.set == nullptr) {
+                        return fail("BuiltinGaussianSplatPass bitonic-sort failed: dynamic descriptor set resolution failed");
+                    }
+
+                    RHI::RHIDescriptorSet* descriptorSet = cachedDescriptorSet->descriptorSet.set;
+                    if (setLayout.usesPerObject) {
+                        if (!passLayout->perObject.IsValid() || passLayout->perObject.set != setIndex) {
+                            return fail("BuiltinGaussianSplatPass bitonic-sort failed: per-object binding is invalid");
+                        }
+
+                        descriptorSet->WriteConstant(
+                            passLayout->perObject.binding,
+                            &perObjectConstants,
+                            sizeof(perObjectConstants));
+                    }
+
+                    descriptorSets[descriptorOffset] = descriptorSet;
+                }
+
+                BindDescriptorSetRanges(
+                    passLayout->firstDescriptorSet,
+                    descriptorSets,
+                    [commandList, passLayout](Core::uint32 firstSet, Core::uint32 count, RHI::RHIDescriptorSet** sets) {
+                        commandList->SetComputeDescriptorSets(
+                            firstSet,
+                            count,
+                            sets,
+                            passLayout->pipelineLayout);
+                    });
+            }
+
+            commandList->Dispatch((workingSet->sortCapacity + 255u) / 256u, 1u, 1u);
+
+            if (!TransitionWorkingSetBuffer(commandList, workingSet->sortDistances, RHI::ResourceStates::NonPixelShaderResource) ||
+                !TransitionWorkingSetBuffer(commandList, workingSet->orderIndices, RHI::ResourceStates::NonPixelShaderResource)) {
+                return fail("BuiltinGaussianSplatPass bitonic-sort failed: resource transition to SRV failed");
+            }
+        }
+    }
 
-    commandList->TransitionBarrier(
-        workingSet->orderIndices.shaderResourceView,
-        RHI::ResourceStates::UnorderedAccess,
-        RHI::ResourceStates::NonPixelShaderResource);
-    workingSet->orderIndices.currentState = RHI::ResourceStates::NonPixelShaderResource;
-    commandList->TransitionBarrier(
-        workingSet->viewData.shaderResourceView,
-        RHI::ResourceStates::UnorderedAccess,
-        RHI::ResourceStates::NonPixelShaderResource);
-    workingSet->viewData.currentState = RHI::ResourceStates::NonPixelShaderResource;
     return true;
 }
 
@@ -1206,7 +1416,11 @@ bool BuiltinGaussianSplatPass::DrawVisibleGaussianSplat(
             static_cast<float>(sceneData.cameraData.viewportHeight),
             sceneData.cameraData.viewportWidth > 0u ? 1.0f / static_cast<float>(sceneData.cameraData.viewportWidth) : 0.0f,
             sceneData.cameraData.viewportHeight > 0u ? 1.0f / static_cast<float>(sceneData.cameraData.viewportHeight) : 0.0f),
-        Math::Vector4(static_cast<float>(cachedGaussianSplat->splatCount), 0.0f, 0.0f, 0.0f)
+        Math::Vector4(
+            static_cast<float>(cachedGaussianSplat->splatCount),
+            static_cast<float>(workingSet->sortCapacity),
+            0.0f,
+            0.0f)
     };
 
     if (passLayout->descriptorSetCount > 0u) {