PBR Renderer: implemented OIT layers rendering

Author: TheMostDiligent

Hydrogent/interface/HnRenderPassState.hpp

@@ -84,6 +84,15 @@ class HnRenderPassState final : public pxr::HdRenderPassState
         return m_FrameAttribsSRB;
     }
 
+    void SetRWOITLayersSRB(IShaderResourceBinding* pSRB)
+    {
+        m_RWOITLayersSRB = pSRB;
+    }
+    IShaderResourceBinding* GetRWOITLayersSRB() const
+    {
+        return m_RWOITLayersSRB;
+    }
+
     void SetFrontFaceCCW(bool FrontFaceCCW)
     {
         m_FrontFaceCCW = FrontFaceCCW;
@@ -144,6 +153,7 @@ class HnRenderPassState final : public pxr::HdRenderPassState
     TEXTURE_FORMAT                                 m_DepthFormat      = TEX_FORMAT_UNKNOWN;
 
     IShaderResourceBinding* m_FrameAttribsSRB = nullptr;
+    IShaderResourceBinding* m_RWOITLayersSRB  = nullptr;
 
     std::array<ITextureView*, MAX_RENDER_TARGETS> m_RTVs        = {};
     ITextureView*                                 m_DSV         = nullptr;

Hydrogent/interface/Tasks/HnBeginOITPassTask.hpp

@@ -80,6 +80,7 @@ class HnBeginOITPassTask final : public HnTask
     HnRenderPassState m_RenderPassState;
 
     RefCntAutoPtr<IShaderResourceBinding> m_ClearLayersSRB;
+    RefCntAutoPtr<IShaderResourceBinding> m_RWLayersSRB;
 
     Uint32 m_BoundOITResourcesVersion = ~0u;
 };

Hydrogent/src/HnRenderPass.cpp

@@ -114,7 +114,8 @@ struct HnRenderPass::RenderState
 
     IDeviceContext* const pCtx;
 
-    const USD_Renderer::ALPHA_MODE AlphaMode;
+    const USD_Renderer::RenderPassType Type;
+    const USD_Renderer::ALPHA_MODE     AlphaMode;
 
     const Uint32 ConstantBufferOffsetAlignment;
     const bool   NativeMultiDrawSupported;
@@ -129,6 +130,7 @@ struct HnRenderPass::RenderState
         USDRenderer{*RenderDelegate.GetUSDRenderer()},
         RendererSettings{USDRenderer.GetSettings()},
         pCtx{RenderDelegate.GetDeviceContext()},
+        Type{RenderPass.m_Params.Type},
         AlphaMode{MaterialTagToPbrAlphaMode(RenderPass.m_MaterialTag)},
         ConstantBufferOffsetAlignment{RenderDelegate.GetDevice()->GetAdapterInfo().Buffer.ConstantBufferOffsetAlignment},
         NativeMultiDrawSupported{RenderDelegate.GetDevice()->GetDeviceInfo().Features.NativeMultiDraw == DEVICE_FEATURE_STATE_ENABLED}
@@ -170,6 +172,13 @@ struct HnRenderPass::RenderState
             pCtx->CommitShaderResources(pFrameSRB, RESOURCE_STATE_TRANSITION_MODE_VERIFY);
         }
 
+        if (Type == PBR_Renderer::RenderPassType::OITLayers && pRWOITLayersSRB == nullptr)
+        {
+            pRWOITLayersSRB = RPState.GetRWOITLayersSRB();
+            VERIFY_EXPR(pRWOITLayersSRB != nullptr);
+            pCtx->CommitShaderResources(pRWOITLayersSRB, RESOURCE_STATE_TRANSITION_MODE_VERIFY);
+        }
+
         pCtx->CommitShaderResources(pNewSRB, RESOURCE_STATE_TRANSITION_MODE_VERIFY);
         pMaterialSRB = pNewSRB;
     }
@@ -221,9 +230,10 @@ struct HnRenderPass::RenderState
     }
 
 private:
-    IPipelineState*         pPSO         = nullptr;
-    IShaderResourceBinding* pMaterialSRB = nullptr;
-    IShaderResourceBinding* pFrameSRB    = nullptr;
+    IPipelineState*         pPSO            = nullptr;
+    IShaderResourceBinding* pMaterialSRB    = nullptr;
+    IShaderResourceBinding* pFrameSRB       = nullptr;
+    IShaderResourceBinding* pRWOITLayersSRB = nullptr;
 
     IBuffer* pIndexBuffer = nullptr;
 

Hydrogent/src/Tasks/HnBeginOITPassTask.cpp

@@ -83,6 +83,7 @@ void HnBeginOITPassTask::Prepare(pxr::HdTaskContext* TaskCtx,
         {
             FrameTargets->OIT = {};
             m_ClearLayersSRB.Release();
+            m_RWLayersSRB.Release();
         }
     }
 
@@ -192,11 +193,16 @@ void HnBeginOITPassTask::Execute(pxr::HdTaskContext* TaskCtx)
     {
         Renderer.CreateClearOITLayersSRB(RenderDelegate->GetFrameAttribsCB(), m_FrameTargets->OIT.Layers, &m_ClearLayersSRB);
     }
+    if (!m_RWLayersSRB)
+    {
+        Renderer.CreateRWOITLayersSRB(m_FrameTargets->OIT.Layers, &m_RWLayersSRB);
+    }
     const TextureDesc& OITTailDesc = m_FrameTargets->OIT.Tail->GetDesc();
     Renderer.ClearOITLayers(pCtx, m_ClearLayersSRB, OITTailDesc.Width, OITTailDesc.Height);
 
     IShaderResourceBinding* pFrameAttribsSRB = RenderDelegate->GetFrameAttribsSRB(HnRenderDelegate::FrameAttribsSRBType::OITLayers);
     m_RenderPassState.SetFrameAttribsSRB(pFrameAttribsSRB);
+    m_RenderPassState.SetRWOITLayersSRB(m_RWLayersSRB);
     m_RenderPassState.Commit(pCtx);
 }
 

PBR/interface/PBR_Renderer.hpp

@@ -796,6 +796,7 @@ class PBR_Renderer
     static OITResources CreateOITResources(IRenderDevice* pDevice, Uint32 Width, Uint32 Height, Uint32 LayerCount);
     OITResources        CreateOITResources(Uint32 Width, Uint32 Height) const;
     void                CreateClearOITLayersSRB(IBuffer* pFrameAttribs, IBuffer* OITLayers, IShaderResourceBinding** ppSRB) const;
+    void                CreateRWOITLayersSRB(IBuffer* OITLayers, IShaderResourceBinding** ppSRB) const;
     void                ClearOITLayers(IDeviceContext* pCtx, IShaderResourceBinding* pSRB, Uint32 Width, Uint32 Height) const;
 
 protected:
@@ -932,8 +933,9 @@ class PBR_Renderer
 
     std::unordered_map<GraphicsPipelineDesc, PsoHashMapType> m_PSOs;
 
-    static constexpr Uint32       ClearOITLayersThreadGroupSize = 16;
-    RefCntAutoPtr<IPipelineState> m_ClearOITLayersPSO;
+    static constexpr Uint32                   ClearOITLayersThreadGroupSize = 16;
+    RefCntAutoPtr<IPipelineState>             m_ClearOITLayersPSO;
+    RefCntAutoPtr<IPipelineResourceSignature> m_RWOITLayersSignature;
 
     std::unique_ptr<StaticShaderTextureIdsArrayType> m_StaticShaderTextureIds;
 };

PBR/src/PBR_Renderer.cpp

@@ -1250,6 +1250,17 @@ void PBR_Renderer::CreateSignature()
     }
 
     CreateCustomSignature(std::move(SignatureDesc));
+
+    if (m_Settings.OITLayerCount > 0)
+    {
+        PipelineResourceSignatureDescX OITLayersSignDesc{"RW OIT Layers"};
+        OITLayersSignDesc
+            .SetBindingIndex(static_cast<Uint8>(m_ResourceSignatures.size()))
+            .SetUseCombinedTextureSamplers(m_Device.GetDeviceInfo().IsGLDevice())
+            .AddResource(SHADER_TYPE_PIXEL, "g_rwOITLayers", SHADER_RESOURCE_TYPE_BUFFER_UAV, SHADER_RESOURCE_VARIABLE_TYPE_MUTABLE);
+        m_RWOITLayersSignature = m_Device.CreatePipelineResourceSignature(OITLayersSignDesc);
+        VERIFY_EXPR(m_RWOITLayersSignature);
+    }
 }
 
 void PBR_Renderer::CreateCustomSignature(PipelineResourceSignatureDescX&& SignatureDesc)
@@ -1861,6 +1872,12 @@ void PBR_Renderer::CreatePSO(PsoHashMapType&             PsoHashMap,
         Macros.Add("USE_GL_POINT_SIZE", "1");
     }
 
+    const Uint32 OITLayerCount = (Key.GetType() == RenderPassType::OITLayers) ||
+            (Key.GetType() == RenderPassType::Main && Key.GetAlphaMode() == ALPHA_MODE_BLEND) ?
+        m_Settings.OITLayerCount :
+        0;
+    Macros.Add("NUM_OIT_LAYERS", static_cast<int>(OITLayerCount));
+
     const bool UseCombinedSamplers = m_Device.GetDeviceInfo().IsGLDevice();
 
     const SHADER_COMPILE_FLAGS ShaderCompileFlags =
@@ -1970,14 +1987,24 @@ void PBR_Renderer::CreatePSO(PsoHashMapType&             PsoHashMap,
         pPS = m_Device.CreateShader(ShaderCI);
     }
 
-    GraphicsPipeline             = GraphicsDesc;
+    GraphicsPipeline = GraphicsDesc;
+    if (Key.GetType() == RenderPassType::OITLayers)
+    {
+        GraphicsPipeline.DepthStencilDesc.DepthWriteEnable = False;
+    }
+
     GraphicsPipeline.InputLayout = InputLayout;
 
     IPipelineResourceSignature* ppSignatures[MAX_RESOURCE_SIGNATURES];
     for (size_t i = 0; i < m_ResourceSignatures.size(); ++i)
         ppSignatures[i] = m_ResourceSignatures[i];
     PSOCreateInfo.ppResourceSignatures    = ppSignatures;
     PSOCreateInfo.ResourceSignaturesCount = static_cast<Uint32>(m_ResourceSignatures.size());
+    if (Key.GetType() == RenderPassType::OITLayers)
+    {
+        ppSignatures[PSOCreateInfo.ResourceSignaturesCount] = m_RWOITLayersSignature;
+        ++PSOCreateInfo.ResourceSignaturesCount;
+    }
 
     PSOCreateInfo.pVS = pVS;
     PSOCreateInfo.pPS = pPS;
@@ -2193,8 +2220,21 @@ void PBR_Renderer::CreateClearOITLayersSRB(IBuffer* pFrameAttribs, IBuffer* OITL
 
     m_ClearOITLayersPSO->CreateShaderResourceBinding(ppSRB, true);
     VERIFY_EXPR(*ppSRB);
-    (*ppSRB)->GetVariableByName(SHADER_TYPE_COMPUTE, "cbFrameAttribs")->Set(pFrameAttribs);
-    (*ppSRB)->GetVariableByName(SHADER_TYPE_COMPUTE, "g_rwOITLayers")->Set(OITLayers->GetDefaultView(BUFFER_VIEW_UNORDERED_ACCESS));
+    ShaderResourceVariableX{*ppSRB, SHADER_TYPE_COMPUTE, "cbFrameAttribs"}.Set(pFrameAttribs);
+    ShaderResourceVariableX{*ppSRB, SHADER_TYPE_COMPUTE, "g_rwOITLayers"}.Set(OITLayers->GetDefaultView(BUFFER_VIEW_UNORDERED_ACCESS));
+}
+
+void PBR_Renderer::CreateRWOITLayersSRB(IBuffer* OITLayers, IShaderResourceBinding** ppSRB) const
+{
+    if (!m_RWOITLayersSignature)
+    {
+        LOG_ERROR_MESSAGE("RW OIT Layers signature is not initialized");
+        return;
+    }
+
+    m_RWOITLayersSignature->CreateShaderResourceBinding(ppSRB, true);
+    VERIFY_EXPR(*ppSRB);
+    ShaderResourceVariableX{*ppSRB, SHADER_TYPE_PIXEL, "g_rwOITLayers"}.Set(OITLayers->GetDefaultView(BUFFER_VIEW_UNORDERED_ACCESS));
 }
 
 void PBR_Renderer::ClearOITLayers(IDeviceContext* pCtx, IShaderResourceBinding* pSRB, Uint32 Width, Uint32 Height) const

Shaders/Common/public/OIT.fxh

@@ -21,3 +21,6 @@ uint GetOITLayerDataOffset(uint2 PixelCoord, uint2 ScreenSize, uint NumLayers)
 {
     return (PixelCoord.y * ScreenSize.x + PixelCoord.x) * NumLayers;
 }
+
+// Opacity is stored as 8-bit value. Any value less than 1/255 is flushed to 0.
+#define OIT_OPACITY_THRESHOLD (1.0/255.0)

Shaders/PBR/private/UpdateOITLayers.psh

@@ -5,6 +5,7 @@
 #include "VSOutputStruct.generated"
 
 #include "PBR_Textures.fxh"
+#include "OIT.fxh"
 
 cbuffer cbFrameAttribs 
 {
@@ -31,17 +32,72 @@ cbuffer cbMaterialAttribs
     PBRMaterialShaderInfo g_Material;
 }
 
+// WebGPU does not support the earlydepthstencil attribute, so we have
+// to manually perform the depth test in the shader.
+#if USE_MANUAL_DEPTH_TEST
+Texture2D<float> g_DepthBuffer;
+#endif
+
+RWStructuredBuffer<uint> g_rwOITLayers;
+
 struct PSOutput
 {
     float4 Color : SV_Target0;
 };
 
+// By default, early depth stencil will be disabled for this shader because it writes to a UAV.
+// Force it to be enabled.
+[earlydepthstencil]
 PSOutput main(in VSOutput VSOut,
               in bool     IsFrontFace : SV_IsFrontFace)
 {
+    float D = VSOut.ClipPos.z;
+#if USE_MANUAL_DEPTH_TEST
+    if (D >= g_DepthBuffer.Load(int3(VSOut.ClipPos.xy, 0)))
+    {
+        discard;
+    }
+#endif
+ 
     float4 BaseColor = GetBaseColor(VSOut, g_Material, g_Frame.Renderer.MipBias, PRIMITIVE.FallbackColor);
 
+    float A = BaseColor.a;
+    uint Layer = 0xFFFFFFFFu;
+    if (A > OIT_OPACITY_THRESHOLD)
+    {
+        float T = 1.0 - A; // Transmittance
+        Layer = PackOITLayer(D, T);
+        uint Offset = GetOITLayerDataOffset(uint2(VSOut.ClipPos.xy), uint2(g_Frame.Camera.f4ViewportSize.xy), uint(NUM_OIT_LAYERS));
+        for (uint i = 0; i < uint(NUM_OIT_LAYERS); ++i)
+        {
+            uint OrigLayer;
+            InterlockedMin(g_rwOITLayers[Offset + i], Layer, OrigLayer);
+            if (OrigLayer == 0xFFFFFFFFu || // Empty space
+                OrigLayer == Layer)         // Layer matches another one exactly
+            {
+                // Do not update tail transmittance
+                Layer = 0xFFFFFFFFu;
+                break;
+            }
+            // Layer > OrigLayer: we did not insert the layer
+            // Layer < OrigLayer: we inserted the layer, so now we need to insert the original layer
+            Layer = max(Layer, OrigLayer);
+        }
+    }
+    
     PSOutput PSOut;
-    PSOut.Color = BaseColor;
+ 
+    // RGB Blend: Src * 1 + Dst * 1
+    // A   Blend: Src * 0 + Dst * SrcA
+    if (Layer == 0xFFFFFFFFu)
+    {
+        PSOut.Color = float4(0.0, 0.0, 0.0, 1.0);
+    }
+    else
+    {
+        float TailT = GetOITLayerTransmittance(Layer);
+        PSOut.Color = float4(1.0 / 255.0, 0.0, 0.0, TailT);
+    }
+    
     return PSOut;
 }