diff --git a/OptiScaler/hooks/VulkanwDx12/CommandBuffer_StateTracker.h b/OptiScaler/hooks/VulkanwDx12/CommandBuffer_StateTracker.h index 2f6a670c..0ea4cb9c 100644 --- a/OptiScaler/hooks/VulkanwDx12/CommandBuffer_StateTracker.h +++ b/OptiScaler/hooks/VulkanwDx12/CommandBuffer_StateTracker.h @@ -15,6 +15,8 @@ #include #include #include +#include +#include #define LOW_PRECISION_TRACKING @@ -234,29 +236,49 @@ class CommandBufferStateTracker CommandBufferStateTracker() { _state = &State::Instance(); } // Call this when command buffers are allocated from a pool - void OnAllocateCommandBuffers(VkCommandPool pool, uint32_t count, const VkCommandBuffer* pCommandBuffers) + void OnAllocateCommandBuffers(VkCommandPool pool, uint32_t count, const VkCommandBuffer* pCommandBuffers, + uint32_t queueFamilyIndex) { - if (_state->currentFeature == nullptr || !_state->currentFeature->IsWithDx12()) - return; - - std::scoped_lock lock(_mtx); - - // Initialize pool epoch ONLY if this is a truly new pool (not seen before) - // Don't overwrite existing epochs - OnResetPool may have already incremented it - if (_poolEpochs.find(pool) == _poolEpochs.end()) + // Fast path: check if pool metadata already exists (read-only lock) { - _poolEpochs[pool] = _globalEpochCounter; - // LOG_DEBUG("Pool {:X} initialized with epoch {}", (size_t) pool, _globalEpochCounter); + std::shared_lock poolLock(_poolMetadataMutex); + if (_poolToQueueFamily.find(pool) != _poolToQueueFamily.end()) + { + // Pool already initialized - only need to map command buffers + std::unique_lock lock(_mtx); + for (uint32_t i = 0; i < count; ++i) + { + _cmdBufferToPool[pCommandBuffers[i]] = pool; + } + return; + } } - // Map each command buffer to its pool - for (uint32_t i = 0; i < count; ++i) + // Slow path: initialize new pool (exclusive lock only for this rare case) { - _cmdBufferToPool[pCommandBuffers[i]] = pool; - } + std::unique_lock poolLock(_poolMetadataMutex); - // LOG_DEBUG("Allocated {} command buffers from pool {:X} (current pool epoch: {})", count, (size_t) pool, - // _poolEpochs[pool]); + // Double-check after acquiring exclusive lock + if (_poolToQueueFamily.find(pool) == _poolToQueueFamily.end()) + { + _poolToQueueFamily[pool] = queueFamilyIndex; + + // Initialize epoch atomically using shared_ptr + auto epochIt = _poolEpochs.find(pool); + if (epochIt == _poolEpochs.end()) + { + _poolEpochs[pool] = + std::make_shared>(_globalEpochCounter.load(std::memory_order_acquire)); + } + } + + // Map command buffers + std::unique_lock lock(_mtx); + for (uint32_t i = 0; i < count; ++i) + { + _cmdBufferToPool[pCommandBuffers[i]] = pool; + } + } } void OnBegin(VkCommandBuffer cmd, const VkCommandBufferBeginInfo* pBeginInfo) @@ -267,45 +289,45 @@ class CommandBufferStateTracker const uint32_t flags = (pBeginInfo) ? pBeginInfo->flags : 0; std::scoped_lock lock(_mtx); - // Check if this command buffer is tracked in a pool auto poolIt = _cmdBufferToPool.find(cmd); if (poolIt == _cmdBufferToPool.end()) { - // Command buffer not mapped to any pool - allocation hook may have been missed - // We'll allow recording but log a warning since this affects epoch validation accuracy - LOG_WARN("Command buffer {:p} not tracked in any pool (allocation hook missed?). " - "Recording will proceed but epoch validation will be disabled for safety.", - (void*) cmd); - - // Create/reset state but mark with epoch 0 to signal "untrusted" status auto& statePtr = _states[cmd]; if (!statePtr) statePtr = std::make_shared(); - - statePtr->ResetForNewRecording(flags, 0); // epoch 0 = untrusted/unvalidated + statePtr->ResetForNewRecording(flags, 0); return; } - // Get the current epoch for this command buffer's pool VkCommandPool pool = poolIt->second; + + // Lock-free atomic read + std::shared_lock poolLock(_poolMetadataMutex); auto epochIt = _poolEpochs.find(pool); uint64_t currentEpoch; if (epochIt == _poolEpochs.end()) { - // Pool not in epoch map - should have been initialized during allocation - // Initialize it now with current global epoch - currentEpoch = _globalEpochCounter; - _poolEpochs[pool] = currentEpoch; - LOG_WARN("Pool {:X} had no epoch - initializing to {} during vkBeginCommandBuffer", (size_t) pool, - currentEpoch); + poolLock.unlock(); + std::unique_lock exclusiveLock(_poolMetadataMutex); + + // Double-check after exclusive lock + epochIt = _poolEpochs.find(pool); + if (epochIt == _poolEpochs.end()) + { + currentEpoch = _globalEpochCounter.load(std::memory_order_acquire); + _poolEpochs[pool] = std::make_shared>(currentEpoch); + } + else + { + currentEpoch = epochIt->second->load(std::memory_order_acquire); + } } else { - currentEpoch = epochIt->second; + currentEpoch = epochIt->second->load(std::memory_order_acquire); } - // Create new state or reset existing auto& statePtr = _states[cmd]; if (!statePtr) statePtr = std::make_shared(); @@ -340,15 +362,34 @@ class CommandBufferStateTracker if (_state->currentFeature == nullptr || !_state->currentFeature->IsWithDx12()) return; - // Invalidate only command buffers from this specific pool by incrementing its epoch - std::scoped_lock lock(_mtx); + // Atomic increment - lock-free for the epoch counter + uint64_t newEpoch = _globalEpochCounter.fetch_add(1, std::memory_order_acq_rel) + 1; - _globalEpochCounter++; - _poolEpochs[pool] = _globalEpochCounter; + // Need lock only to access/modify the map itself + std::unique_lock poolLock(_poolMetadataMutex); - // LOG_DEBUG("Pool {:X} reset - pool epoch set to {} - command buffers from THIS POOL invalidated until next " - // "vkBeginCommandBuffer", - // (size_t) pool, _globalEpochCounter); + auto epochIt = _poolEpochs.find(pool); + if (epochIt != _poolEpochs.end()) + { + epochIt->second->store(newEpoch, std::memory_order_release); + } + else + { + _poolEpochs[pool] = std::make_shared>(newEpoch); + } + + // if (_state->currentFeature == nullptr || !_state->currentFeature->IsWithDx12()) + // return; + + //// Invalidate only command buffers from this specific pool by incrementing its epoch + // std::scoped_lock lock(_mtx); + + //_globalEpochCounter++; + //_poolEpochs[pool] = _globalEpochCounter; + + //// LOG_DEBUG("Pool {:X} reset - pool epoch set to {} - command buffers from THIS POOL invalidated until next " + //// "vkBeginCommandBuffer", + //// (size_t) pool, _globalEpochCounter); } void OnBindPipeline(VkCommandBuffer cmd, VkPipelineBindPoint bindPoint, VkPipeline pipeline) @@ -950,6 +991,27 @@ class CommandBufferStateTracker _hasCachedFns = true; } + std::optional GetCommandBufferQueueFamily(VkCommandBuffer cmd) const + { + std::scoped_lock lock(_mtx); + + auto poolIt = _cmdBufferToPool.find(cmd); + if (poolIt == _cmdBufferToPool.end()) + { + LOG_WARN("Command buffer {:X} not tracked in any pool", (size_t) cmd); + return std::nullopt; + } + + auto familyIt = _poolToQueueFamily.find(poolIt->second); + if (familyIt == _poolToQueueFamily.end()) + { + LOG_WARN("Pool {:X} has no queue family info", (size_t) poolIt->second); + return std::nullopt; + } + + return familyIt->second; + } + private: inline static State* _state; @@ -1278,42 +1340,34 @@ class CommandBufferStateTracker if (it == _states.end() || !it->second) return false; - // Check if this command buffer is tracked in a pool auto poolIt = _cmdBufferToPool.find(cmd); if (poolIt == _cmdBufferToPool.end()) { - // Command buffer not mapped to any pool - allocation hook may have been missed - // This is potentially unsafe as we can't validate against pool-specific epochs - LOG_WARN("Command buffer {:p} not tracked in any pool (allocation hook missed?). " - "Cannot validate epoch - refusing replay for safety.", - (void*) cmd); + LOG_WARN("Command buffer {:p} not tracked in any pool", (void*) cmd); return false; } - // Get the current epoch for this command buffer's pool VkCommandPool pool = poolIt->second; + + // Lock-free atomic read + std::shared_lock poolLock(_poolMetadataMutex); auto epochIt = _poolEpochs.find(pool); if (epochIt == _poolEpochs.end()) { - // Pool exists in mapping but has no epoch - should not happen if properly initialized - LOG_ERROR("Pool {:X} for command buffer {:p} has no epoch entry. Internal state corruption?", (size_t) pool, - (void*) cmd); + LOG_ERROR("Pool {:X} has no epoch entry", (size_t) pool); return false; } - uint64_t currentPoolEpoch = epochIt->second; + uint64_t currentPoolEpoch = epochIt->second->load(std::memory_order_acquire); + poolLock.unlock(); - // Check if this command buffer's state is stale (invalidated by its pool's reset) if (it->second->BeginEpoch < currentPoolEpoch) { - LOG_WARN("Command buffer {:p} has stale state (epoch {} < pool {:X} epoch {}), refusing replay. " - "This command buffer was invalidated by vkResetCommandPool and must not be used until " - "vkBeginCommandBuffer is called.", - (void*) cmd, it->second->BeginEpoch, (size_t) pool, currentPoolEpoch); + LOG_WARN("Command buffer {:p} has stale state (epoch {} < pool {:X} epoch {})", (void*) cmd, + it->second->BeginEpoch, (size_t) pool, currentPoolEpoch); return false; } - // Deep copy the state under lock - this is now a true immutable snapshot out = *it->second; return true; } @@ -1347,26 +1401,26 @@ class CommandBufferStateTracker auto poolIt = _cmdBufferToPool.find(srcCmd); if (poolIt == _cmdBufferToPool.end()) { - // Command buffer not mapped to any pool - allocation hook may have been missed - // This is potentially unsafe as we can't validate against pool-specific epochs LOG_WARN("Command buffer {:p} not tracked in any pool (allocation hook missed?). " "Cannot validate epoch - refusing replay for safety.", (void*) srcCmd); return false; } - // Get the current epoch for this command buffer's pool + // Get the current epoch for this command buffer's pool - LOCK-FREE ATOMIC READ VkCommandPool pool = poolIt->second; + + std::shared_lock poolLock(_poolMetadataMutex); auto epochIt = _poolEpochs.find(pool); if (epochIt == _poolEpochs.end()) { - // Pool exists in mapping but has no epoch - should not happen if properly initialized LOG_ERROR("Pool {:X} for command buffer {:p} has no epoch entry. Internal state corruption?", (size_t) pool, (void*) srcCmd); return false; } - uint64_t currentPoolEpoch = epochIt->second; + uint64_t currentPoolEpoch = epochIt->second->load(std::memory_order_acquire); + poolLock.unlock(); // Check if this command buffer's state is stale (invalidated by pool reset) if (it->second->BeginEpoch < currentPoolEpoch) @@ -1542,16 +1596,9 @@ class CommandBufferStateTracker return true; } - CommandBufferState* GetState(VkCommandBuffer cmd) - { - std::scoped_lock lock(_mtx); - auto& statePtr = _states[cmd]; - if (!statePtr) - statePtr = std::make_shared(); - return statePtr.get(); // Return raw pointer - } + mutable std::shared_mutex _mtx; + mutable std::shared_mutex _poolMetadataMutex; // Separate lock for pool metadata - mutable std::mutex _mtx; std::unordered_map> _states; VulkanCmdFns _cachedFns {}; @@ -1559,7 +1606,8 @@ class CommandBufferStateTracker // Per-pool epoch tracking for accurate invalidation std::unordered_map _cmdBufferToPool; - std::unordered_map _poolEpochs; - uint64_t _globalEpochCounter = 1; + std::unordered_map _poolToQueueFamily; + std::unordered_map>> _poolEpochs; + std::atomic _globalEpochCounter { 1 }; }; } // namespace vk_state