Implement transient buffer system for global uniforms and add constant buffer support

Author: hypernewbie

src/vrhi_backend.cpp

@@ -34,7 +34,7 @@ std::unordered_map< uint64_t, const vhVertexLayoutDef* > vhCmdBackendState::s_la
 std::vector< nvrhi::VertexAttributeDesc > vhCmdBackendState::s_attributes;
 
 // --------------------------------------------------------------------------
-// Implementation
+// Backend :: Utils & Helpers
 // --------------------------------------------------------------------------
 
 static const char* vhResourceTypeToString( nvrhi::ResourceType type )
@@ -91,6 +91,42 @@ bool vhCmdBackendState::BE_Util_ShaderStageMatches( uint64_t flags, bool useComp
     return false;
 }
 
+int64_t vhCmdBackendState::BE_Util_WriteGlobalUniform( const vhState& state, vhBackendTransientBuffer& tbuf, uint64_t& lastHash )
+{
+    vhGlobalUniform u;
+    vhWriteStateToGlobalUniform( state, u );
+    uint64_t hash = vhHashGlobalUniform( u );
+    
+    if ( hash == lastHash && tbuf.offset >= sizeof( vhGlobalUniform ) )
+        return tbuf.offset - sizeof( vhGlobalUniform );
+
+    if ( ( tbuf.offset % VRHI_CBUF_ALIGN ) != 0 )
+        return -1;
+
+    int64_t offset = tbuf.Alloc( sizeof( vhGlobalUniform ) );
+    if ( offset < 0 )
+        return -1;
+
+    {
+        std::lock_guard<std::mutex> lock( g_nvRHIStateMutex );
+        uint8_t* ptr = tbuf.Map_DeviceStateLocked();
+        if ( !ptr )
+        {
+            VRHI_ERR("BE_Util_WriteGlobalUniform: Failed to map global uniform buffer!\n");
+            return -1;
+        }
+        memcpy( ptr + offset, &u, sizeof( u ) );
+        // Unmap is deferred until Frame/Flush boundary
+    }
+
+    lastHash = hash;
+    return offset;
+}
+
+// --------------------------------------------------------------------------
+// Backend :: Complex BE Low Level NVRHI Device Functions
+// --------------------------------------------------------------------------
+
 void vhCmdBackendState::BE_UpdateTexture( vhBackendTexture& btex, const vhMem* data, glm::ivec4 arrayMipUpdateRange )
 {
     if ( !btex.handle || !data || !data->size() ) return;
@@ -637,6 +673,47 @@ bool vhCmdBackendState::BE_PreSubmitCommon_FindResource(
 
     switch ( item.type )
     {
+        case nvrhi::ResourceType::ConstantBuffer:
+        case nvrhi::ResourceType::VolatileConstantBuffer:
+        {
+            if ( item.slot == 0 )
+            {
+                // TODO: Implement global uniform buffer binding.
+                return false;
+            }
+
+            auto it = stageTable.bufferTable.find( item.slot );
+            if ( it == stageTable.bufferTable.end() )
+            {
+                if ( state.debugFlags & VRHI_STATE_DEBUG_LOG_BINDING_MISMATCH ) VRHI_ERR( "FindResource: ConstantBuffer not found in cache at slot %d\n", item.slot );
+                return false;
+            }
+            const auto result = &it->second;
+            assert( result );
+            if ( !result->handle )
+            {
+                if ( state.debugFlags & VRHI_STATE_DEBUG_LOG_BINDING_MISMATCH ) VRHI_ERR( "FindResource: ConstantBuffer found in cache at slot %d but null handle.\n", item.slot );
+                return false;
+            }
+            if ( !result->handle->getDesc().isConstantBuffer )
+            {
+                if ( state.debugFlags & VRHI_STATE_DEBUG_LOG_BINDING_MISMATCH ) VRHI_ERR( "FindResource: ConstantBuffer found in cache at slot %d but NOT a ConstantBuffer.\n", item.slot );
+                return false;
+            }
+
+            uint64_t size = result->binding->byteSize ? result->binding->byteSize : result->handle->getDesc().byteSize;
+            nvrhi::BufferRange range( result->binding->byteOffset, size );
+            outItem = nvrhi::BindingSetItem::ConstantBuffer( item.slot, result->handle, range );
+            if ( result->handle->getDesc().isVolatile && item.type != nvrhi::ResourceType::VolatileConstantBuffer )
+            {
+                 if ( state.debugFlags & VRHI_STATE_DEBUG_LOG_BINDING_MISMATCH ) VRHI_ERR( "FindResource: Volatile Buffer bound to Static ConstantBuffer slot %d. This may be unsafe!\n", item.slot );
+                return false;
+            }
+            outItem.type = item.type;
+            if ( state.debugFlags & VRHI_STATE_DEBUG_LOG_ALL_BINDINGS ) VRHI_LOG( "FindResource: ConstantBuffer found in cache at slot %d\n", item.slot );
+            return true;
+        }
+
         case nvrhi::ResourceType::Texture_SRV:
         case nvrhi::ResourceType::Texture_UAV:
         {
@@ -1029,29 +1106,17 @@ void vhCmdBackendState::init()
 {
     std::lock_guard< std::mutex > lock( backendMutex );
 
-    if ( !m_globalUniformBuffer )
+    if ( !m_globalUniformBuffer.handle[0] )
     {
         // Called from vhInit which already holds g_nvRHIStateMutex lock, and before RHI thread even starts.
         // So we don't need to lock g_nvRHIStateMutex here.
-    
+
         nvrhi::BufferDesc desc;
-        desc.setByteSize( sizeof( vhGlobalUniform ) );
+        desc.setByteSize( 16 * 1024 * sizeof( vhGlobalUniform ) ); // 16MB
         desc.setIsConstantBuffer( true );
+        desc.setCpuAccess( nvrhi::CpuAccessMode::Write );
         desc.setDebugName( "GlobalUniforms" );
-        nvrhi::BufferHandle bhandle = g_vhDevice->createBuffer( desc );
-        if ( !bhandle ) 
-        {
-            VRHI_ERR( "vhCmdBackendState::init() : Failed to create Global Uniform Buffer!\n" );
-            assert( !"Failed to create Global Uniform Buffer." );
-            return;
-        }
-
-        m_globalUniformBuffer = std::make_unique< vhBackendBuffer >();
-        m_globalUniformBuffer->handle = bhandle;
-        m_globalUniformBuffer->name = "GlobalUniforms";
-        m_globalUniformBuffer->desc = desc;
-        m_globalUniformBuffer->stride = sizeof( vhGlobalUniform );
-        m_globalUniformBuffer->flags = VRHI_BUFFER_NONE;
+        m_globalUniformBuffer.Init_DeviceStateLocked( desc );
     }
 }
 
@@ -1060,7 +1125,7 @@ void vhCmdBackendState::shutdown()
     std::lock_guard< std::mutex > lock( backendMutex );
     std::lock_guard< std::mutex > lock2( g_nvRHIStateMutex );
 
-    m_globalUniformBuffer.reset();
+    m_globalUniformBuffer.Shutdown_DeviceStateLocked();
 
     backendTextures.clear();
     backendBuffers.clear();
@@ -1821,17 +1886,22 @@ void vhCmdBackendState::Handle_vhCmdSetStateAttachments( VIDL_vhCmdSetStateAttac
 void vhCmdBackendState::Handle_vhFlushInternal( VIDL_vhFlushInternal* cmd )
 {
     BE_CmdRAII cmdRAII( cmd );
+
+    // TODO: Flush all transient buffer maps here.
+    {
+        std::lock_guard<std::mutex> lock( g_nvRHIStateMutex );
+        m_globalUniformBuffer.Unmap_DeviceStateLocked();
+    }
+
+
     // Free all cmd memory allocations, because hitting this flush means all previous commands have been processed.
     {
         std::lock_guard<std::mutex> lock( g_vhMemListMutex );
         g_vhMemList.clear();
     }
 
-    if ( cmd->waitForGPU )
-    {
-        // This uses g_nvRHIStateMutex then gives it up, we need to avoid double-locking.
-        vhCmdListFlushAll();
-    }
+    // This uses g_nvRHIStateMutex then gives it up, we need to avoid double-locking.
+    vhCmdListFlushAll();
 
     {
         std::lock_guard< std::mutex > lock( g_nvRHIStateMutex );
@@ -1840,6 +1910,8 @@ void vhCmdBackendState::Handle_vhFlushInternal( VIDL_vhFlushInternal* cmd )
             g_vhDevice->waitForIdle();
         }
         g_vhDevice->runGarbageCollection();
+        m_globalUniformBuffer.Step();
+        m_globalUniformBufferLastHash = 0;
     }
 
     // Notify caller that we're done.

src/vrhi_backend.h

@@ -69,6 +69,68 @@ struct vhBackendShader
     std::vector< vhSpecConstant > specConstants;
 };
 
+struct vhBackendTransientBuffer
+{
+    int64_t size = 0;
+    nvrhi::BufferHandle handle[VRHI_MAX_FRAMES_INFLIGHT];
+    uint32_t frameIdx = 0;
+    int64_t offset = 0;
+    uint8_t* ptr = nullptr; // For storing mapping.
+
+    inline void Reset() { size = 0; frameIdx = 0; offset = 0; ptr = 0; }
+
+    inline int64_t Alloc( int64_t bytes )
+    {
+        if ( offset + bytes > size )
+            return -1;
+        int64_t ret = offset;
+        offset += bytes;
+        return ret;
+    }
+
+    inline void Step()
+    {
+        frameIdx = ( frameIdx + 1 ) % VRHI_MAX_FRAMES_INFLIGHT;
+        offset = 0;
+    }
+
+    inline void Init_DeviceStateLocked( const nvrhi::BufferDesc& desc )
+    {
+        // WARNING: Lock g_nvRHIStateMutex before calling this.
+        Reset();
+        size = desc.byteSize;
+        for ( uint32_t i = 0; i < VRHI_MAX_FRAMES_INFLIGHT; i++ )
+        {
+            handle[i] = g_vhDevice->createBuffer( desc );
+            assert( handle[i] != nullptr );
+        }
+    }
+
+    inline uint8_t* Map_DeviceStateLocked()
+    {
+        // WARNING: Lock g_nvRHIStateMutex before calling this.
+        if ( ptr ) return ptr;
+        ptr = ( uint8_t* ) g_vhDevice->mapBuffer( handle[frameIdx], nvrhi::CpuAccessMode::Write );
+        return ptr;
+    }
+
+    inline void Unmap_DeviceStateLocked()
+    {
+        // WARNING: Lock g_nvRHIStateMutex before calling this.
+        if ( ptr ) g_vhDevice->unmapBuffer( handle[frameIdx] );
+        ptr = nullptr;
+    }
+
+    inline void Shutdown_DeviceStateLocked()
+    {
+        // WARNING: Lock g_nvRHIStateMutex before calling this.
+        for ( uint32_t i = 0; i < VRHI_MAX_FRAMES_INFLIGHT; i++ )
+        {
+            handle[i] = nullptr;
+        }
+    }
+};
+
 struct vhStateResolveCache
 {
     bool init = false;
@@ -120,25 +182,6 @@ struct vhStateResolveCache
 // Main Backend State
 // --------------------------------------------------------------------------
 
-struct vhGlobalUniform
-{
-    glm::vec4 u_viewRect;
-    glm::vec4 u_viewTexel;
-    glm::mat4 u_view;
-    glm::mat4 u_invView;
-    glm::mat4 u_proj;
-    glm::mat4 u_invProj;
-    glm::mat4 u_viewProj;
-    glm::mat4 u_invViewProj;
-    glm::mat4 u_model[4];
-    glm::mat4 u_modelView;
-    glm::mat4 u_modelViewProj;
-    glm::vec4 u_alphaRef4;
-    glm::vec4 u_global[32];
-};
-static_assert( sizeof( vhGlobalUniform ) < 16384, "vhGlobalUniform must be smaller than 16KB" );
-static_assert( sizeof( vhGlobalUniform ) == 1328, "vhGlobalUniform packing mismatch" );
-
 class vhCmdBackendState : public VIDLHandler
 {
     friend class vhCmdBackendStateTest;
@@ -150,9 +193,9 @@ class vhCmdBackendState : public VIDLHandler
     std::map< vhShader, std::unique_ptr< vhBackendShader > > backendShaders;
     std::map< vhStateId, vhState > backendStates;
     std::unordered_map< uint64_t, nvrhi::FramebufferHandle > backendFramebuffers;
-
-    // Internal Global Uniform Buffer
-    std::unique_ptr< vhBackendBuffer > m_globalUniformBuffer;
+    vhBackendTransientBuffer m_globalUniformBuffer;
+    uint64_t m_globalUniformBufferLastHash = 0;
+    
     // RAII for vhMem, takes ownership of the pointer and auto-destructs it.
     inline std::unique_ptr< vhMem > BE_MemRAII( const vhMem* mem )
     {
@@ -197,6 +240,8 @@ class vhCmdBackendState : public VIDLHandler
 
     bool BE_Util_ShaderStageMatches( uint64_t flags, bool useCompute, bool useGraphics );
 
+    int64_t BE_Util_WriteGlobalUniform( const vhState& state, vhBackendTransientBuffer& tbuf, uint64_t& lastHash );
+
     // --------------------------------------------------------------------------
     // Backend :: Complex BE Low Level NVRHI Device Functions
     // --------------------------------------------------------------------------

src/vrhi_buffer.cpp

@@ -403,4 +403,5 @@ uint64_t vhGetBufferInfo( vhBuffer buffer, uint32_t* outStride, uint64_t* outFla
 void* vhGetBufferNvrhiHandle( vhBuffer buffer )
 {
     return vhBackendQueryBufferHandle( buffer );
-}
+}
+

src/vrhi_device.cpp

@@ -452,10 +452,10 @@ void vhFlushInternal( std::atomic<bool>* fence, bool waitForGPU )
     vhCmdEnqueue( cmd );
 }
 
-void vhFlush()
+void vhFlush( bool wait )
 {
     std::atomic<bool> fence = false;
-    vhFlushInternal( &fence, false );
+    vhFlushInternal( wait ? &fence : nullptr, false );
 
     // Wait for fence to be signaled
     while ( !fence.load() )
@@ -994,4 +994,7 @@ uint64_t vhHashSamplerDesc( const nvrhi::SamplerDesc& desc )
     return h;
 }
 
-
+uint64_t vhHashGlobalUniform( const vhGlobalUniform& u )
+{
+    return komihash( &u, sizeof( u ), 0 );
+}

src/vrhi_internal.h

@@ -93,6 +93,10 @@
 // Internal Declarations
 // --------------------------------------------------------------------------
 
+// Triple buffering is the max we can support.
+#define VRHI_MAX_FRAMES_INFLIGHT 3
+#define VRHI_CBUF_ALIGN 256
+
 // Internal Vulkan State
 extern VkInstance g_vulkanInstance;
 extern VkPhysicalDevice g_vulkanPhysicalDevice;
@@ -198,6 +202,25 @@ nvrhi::GraphicsPipelineHandle vhPSOCacheGet( const nvrhi::GraphicsPipelineDesc&
 void vhBindingSetCacheClear();
 nvrhi::BindingSetHandle vhGetBindingSet( const nvrhi::BindingSetDesc& desc, nvrhi::BindingLayoutHandle layout );
 
+struct vhGlobalUniform
+{
+    glm::vec4 u_viewRect;
+    glm::vec4 u_viewTexel;
+    glm::mat4 u_view;
+    glm::mat4 u_invView;
+    glm::mat4 u_proj;
+    glm::mat4 u_invProj;
+    glm::mat4 u_viewProj;
+    glm::mat4 u_invViewProj;
+    glm::mat4 u_worldX[4]; // world[1+]. (world[0] is in pushbuffer).
+    glm::mat4 u_worldView;
+    glm::mat4 u_worldViewProj;
+    glm::vec4 u_alphaRef4;
+    glm::vec4 u_global[13];
+};
+static_assert( sizeof( vhGlobalUniform ) < 16384, "vhGlobalUniform must be smaller than 16KB" );
+static_assert( sizeof( vhGlobalUniform ) == 1024, "vhGlobalUniform packing mismatch" );
+
 bool vhReflectSpirv(
     const std::vector< uint32_t >& spirvBlob,
     nvrhi::BindingLayoutDesc& outDesc,
@@ -217,6 +240,8 @@ uint64_t vhHashShaderDebugName( nvrhi::ShaderHandle shader );
 uint64_t vhHashShaderSPIRV( const std::vector< uint32_t >& spirv );
 uint64_t vhHashInputLayout( nvrhi::InputLayoutHandle layout );
 uint64_t vhHashSamplerDesc( const nvrhi::SamplerDesc& desc );
+uint64_t vhHashGlobalUniform( const vhGlobalUniform& u );
+void vhWriteStateToGlobalUniform( const vhState& state, vhGlobalUniform& out );
 nvrhi::PrimitiveType vhTranslatePrimitiveType( uint64_t stateFlags );
 nvrhi::BlendState vhTranslateBlendState( uint64_t stateFlags );
 nvrhi::DepthStencilState vhTranslateDepthStencilState( uint64_t stateFlags, uint32_t frontStencil, uint32_t backStencil );