Fixes in assets' pages calculation

framework/encode/vulkan_capture_manager.cpp

@@ -2204,6 +2204,20 @@ void VulkanCaptureManager::PreProcess_vkUnmapMemory(VkDevice device, VkDeviceMem
 
     if (wrapper->mapped_data != nullptr)
     {
+        if (IsCaptureModeTrack())
+        {
+            assert(state_tracker_ != nullptr);
+            state_tracker_->TrackMappedMemory(device, memory, nullptr, 0, 0, 0);
+        }
+        else
+        {
+            // Perform subset of the state tracking performed by VulkanStateTracker::TrackMappedMemory, only storing
+            // values needed for non-tracking capture.
+            wrapper->mapped_data   = nullptr;
+            wrapper->mapped_offset = 0;
+            wrapper->mapped_size   = 0;
+        }
+
         if (GetMemoryTrackingMode() == CaptureSettings::MemoryTrackingMode::kPageGuard ||
             GetMemoryTrackingMode() == CaptureSettings::MemoryTrackingMode::kUserfaultfd)
         {
@@ -2235,20 +2249,6 @@ void VulkanCaptureManager::PreProcess_vkUnmapMemory(VkDevice device, VkDeviceMem
                 mapped_memory_.erase(wrapper);
             }
         }
-
-        if (IsCaptureModeTrack())
-        {
-            assert(state_tracker_ != nullptr);
-            state_tracker_->TrackMappedMemory(device, memory, nullptr, 0, 0, 0);
-        }
-        else
-        {
-            // Perform subset of the state tracking performed by VulkanStateTracker::TrackMappedMemory, only storing
-            // values needed for non-tracking capture.
-            wrapper->mapped_data   = nullptr;
-            wrapper->mapped_offset = 0;
-            wrapper->mapped_size   = 0;
-        }
     }
     else
     {

framework/encode/vulkan_state_tracker.cpp

@@ -31,6 +31,7 @@
 #include "encode/vulkan_track_struct.h"
 #include "graphics/vulkan_struct_get_pnext.h"
 #include "util/logging.h"
+#include "util/page_guard_manager.h"
 #include "util/platform.h"
 #include "Vulkan-Utility-Libraries/vk_format_utils.h"
 #include "util/to_string.h"
@@ -470,6 +471,12 @@ void VulkanStateTracker::TrackMappedMemory(VkDevice         device,
     wrapper->mapped_offset = mapped_offset;
     wrapper->mapped_size   = mapped_size;
     wrapper->mapped_flags  = mapped_flags;
+
+    // Scan assets on vkUnmapMemory
+    if (mapped_data == nullptr)
+    {
+        TrackMappedAssetsWrites(wrapper->handle_id);
+    }
 }
 
 void VulkanStateTracker::TrackBeginRenderPass(VkCommandBuffer command_buffer, const VkRenderPassBeginInfo* begin_info)
@@ -2578,16 +2585,16 @@ void VulkanStateTracker::TrackCmdDrawMeshTasksIndirectCountEXT(VkCommandBuffer c
     }
 }
 
-void VulkanStateTracker::TrackMappedAssetsWrites()
+void VulkanStateTracker::TrackMappedAssetsWrites(format::HandleId memory_id)
 {
     util::PageGuardManager* manager = util::PageGuardManager::Get();
     if (manager == nullptr)
     {
         return;
     }
 
-    std::unordered_map<uint64_t, const util::PageStatusTracker::PageStatus&> memories_page_status;
-    manager->GetDirtyMemoryRegions(memories_page_status);
+    std::unordered_map<uint64_t, const util::PageGuardManager::MemoryInfo&> memories_page_status;
+    manager->GetDirtyMemoryRegions(memory_id, memories_page_status);
     const size_t page_size = util::platform::GetSystemPageSize();
 
     for (const auto& entry : memories_page_status)
@@ -2598,19 +2605,23 @@ void VulkanStateTracker::TrackMappedAssetsWrites()
         dev_mem_wrapper->asset_map_lock.lock();
         for (auto& asset : dev_mem_wrapper->bound_assets)
         {
-            if (asset->dirty || !asset->size)
+            if (asset->dirty || !asset->size || asset->bind_offset < entry.second.mapped_offset)
             {
                 continue;
             }
 
-            const size_t first_page = asset->bind_offset / page_size;
-            const size_t last_page  = (asset->bind_offset + asset->size - 1) / page_size;
+            const size_t asset_offset_from_mapping = asset->bind_offset - entry.second.mapped_offset;
+            const size_t first_page                = asset_offset_from_mapping / page_size;
+            const size_t last_page                 = (asset_offset_from_mapping + asset->size - 1) / page_size;
             assert(first_page <= last_page);
-            assert(first_page <= entry.second.size());
 
-            for (size_t page = first_page; page < entry.second.size() && page <= last_page; ++page)
+            const util::PageStatusTracker::PageStatus& active_writes = entry.second.status_tracker.GetActiveWrites();
+            const size_t                               total_pages   = active_writes.size();
+            assert(first_page <= total_pages);
+
+            for (size_t page = first_page; page < total_pages && page <= last_page; ++page)
             {
-                if (entry.second[page])
+                if (active_writes[page])
                 {
                     asset->dirty = true;
                     break;

framework/encode/vulkan_state_tracker.h

@@ -767,7 +767,7 @@ class VulkanStateTracker
 
     void InsertBufferAssetInCommandBuffer(VkCommandBuffer command_buffer, VkBuffer buffer);
 
-    void TrackMappedAssetsWrites();
+    void TrackMappedAssetsWrites(format::HandleId memory_id = format::kNullHandleId);
 
     void MarkReferencedAssetsAsDirty(vulkan_wrappers::CommandBufferWrapper* cmd_buf_wrapper);
 

framework/util/page_guard_manager.cpp

@@ -1112,6 +1112,7 @@ void* PageGuardManager::AddTrackedMemory(uint64_t  memory_id,
                                      std::forward_as_tuple(memory_id),
                                      std::forward_as_tuple(mapped_memory,
                                                            mapped_range,
+                                                           mapped_offset,
                                                            shadow_memory,
                                                            shadow_size,
                                                            aligned_address,
@@ -1370,16 +1371,27 @@ const void* PageGuardManager::GetMappedMemory(uint64_t memory_id) const
     return nullptr;
 }
 
-void PageGuardManager::GetDirtyMemoryRegions(
-    std::unordered_map<uint64_t, const PageStatusTracker::PageStatus&>& memories_page_status)
+void PageGuardManager::GetDirtyMemoryRegions(uint64_t                                         memory_id,
+                                             std::unordered_map<uint64_t, const MemoryInfo&>& memories_page_status)
 {
     std::lock_guard<std::mutex> lock(tracked_memory_lock_);
 
-    for (auto& entry : memory_info_)
+    if (memory_id)
+    {
+        for (auto& entry : memory_info_)
+        {
+            if (entry.second.is_modified)
+            {
+                memories_page_status.emplace(entry.first, entry.second);
+            }
+        }
+    }
+    else
     {
-        if (entry.second.is_modified)
+        const auto entry = memory_info_.find(memory_id);
+        if (entry != memory_info_.end())
         {
-            auto new_entry = memories_page_status.emplace(entry.first, entry.second.status_tracker.GetActiveWrites());
+            memories_page_status.emplace(entry->first, entry->second);
         }
     }
 }

framework/util/page_guard_manager.h

@@ -145,9 +145,6 @@ class PageGuardManager
 
     void UffdUnblockRtSignal();
 
-    void
-    GetDirtyMemoryRegions(std::unordered_map<uint64_t, const PageStatusTracker::PageStatus&>& memories_page_status);
-
   protected:
     PageGuardManager();
 
@@ -161,11 +158,12 @@ class PageGuardManager
 
     ~PageGuardManager();
 
-  private:
+  public:
     struct MemoryInfo
     {
         MemoryInfo(void*       mm,
                    size_t      mr,
+                   size_t      mo,
                    void*       sm,
                    size_t      sr,
                    void*       aa,
@@ -193,6 +191,7 @@ class PageGuardManager
 
         void*  mapped_memory;  // Pointer to mapped memory to be tracked.
         size_t mapped_range;   // Size of the mapped memory range.
+        size_t mapped_offset;  // The offset of the mapped memory from the VkDeviceMemory's base address.
         void*  shadow_memory;  // Shadow memory for mapped memory types that cannot be tracked by guard pages.
         size_t shadow_range;   // Size of the shadow memory allocation, which is the mapped memory size adjusted to be a
                                // multiple of system page size.
@@ -213,6 +212,12 @@ class PageGuardManager
 #endif
     };
 
+    void GetDirtyMemoryRegions(uint64_t                                         memory_id,
+                               std::unordered_map<uint64_t, const MemoryInfo&>& memories_page_status);
+
+    typedef std::unordered_map<uint64_t, MemoryInfo> MemoryInfoMap;
+
+  private:
     struct ShadowMemoryInfo
     {
         ShadowMemoryInfo(void* sm, size_t ss, size_t tp, size_t lss) :
@@ -226,8 +231,6 @@ class PageGuardManager
         std::vector<bool> page_loaded;       // Tracks which pages have been loaded.
     };
 
-    typedef std::unordered_map<uint64_t, MemoryInfo> MemoryInfoMap;
-
   private:
     size_t GetSystemPagePotShift() const;
     void   InitializeSystemExceptionContext();

framework/util/page_status_tracker.h

@@ -53,7 +53,7 @@ class PageStatusTracker
 
     void SetAllBlocksActiveWrite() { std::fill(active_writes_.begin(), active_writes_.end(), 1); }
 
-    const PageStatus& GetActiveWrites() { return active_writes_; }
+    const PageStatus& GetActiveWrites() const { return active_writes_; }
 
   private:
     PageStatus active_writes_; //< Track blocks that have been written.