Bugfix device memory report bind path hardening

layersvt/device_memory_report/device_memory_report.cpp

@@ -393,6 +393,16 @@ uint64_t DeviceMemoryReport::GetUsageCounterBytes(const std::string& track) {
     return it != usage_memory_bytes_.end() ? it->second : 0;
 }
 
+void DeviceMemoryReport::Reset() {
+    std::scoped_lock lock(map_mutex_, counter_mutex_);
+    vk_instance_map_.clear();
+    has_callback_map_.clear();
+    device_memory_properties_map_.clear();
+    resources_.clear();
+    resource_to_memory_map_.clear();
+    memory_allocations_.clear();
+    usage_memory_bytes_.clear();
+}
 void DeviceMemoryReport::OnCreateImage(uint64_t image_handle, VkImageUsageFlags usage) {
     std::lock_guard<std::mutex> lock(counter_mutex_);
     auto& res = resources_[image_handle];

layersvt/device_memory_report/device_memory_report.h

@@ -170,6 +170,11 @@ class DeviceMemoryReport {
      */
     VkDeviceSize GetRecordedResourceSize(uint64_t resource_handle);
 
+    /**
+     * @brief Resets all tracked internal state (for testing).
+     */
+    void Reset();
+
     /**
      * @brief Retrieves the current total memory bytes recorded on a usage track (for testing).
      * @param track The full usage track name (e.g. "vulkan.mem.app.usage.unbound_memory").

layersvt/device_memory_report/device_memory_report_handwritten_functions.h

@@ -258,72 +258,122 @@ EXPORT_FUNCTION VKAPI_ATTR VkResult VKAPI_CALL vkEnumerateDeviceLayerProperties(
 
 // Intercept memory binding to correlate buffer object handles with device memory allocations.
 VKAPI_ATTR VkResult VKAPI_CALL vkBindBufferMemory(VkDevice device, VkBuffer buffer, VkDeviceMemory memory, VkDeviceSize memoryOffset) {
-    VkResult result = device_dispatch_table(device)->BindBufferMemory(device, buffer, memory, memoryOffset);
-    if (result == VK_SUCCESS && buffer != VK_NULL_HANDLE && memory != VK_NULL_HANDLE) {
+    auto* table = device_dispatch_table(device);
+    assert(table->BindBufferMemory != nullptr);
+    assert(buffer != VK_NULL_HANDLE);
+    assert(memory != VK_NULL_HANDLE);
+    if (DeviceMemoryReport::Get().GetRecordedResourceSize(reinterpret_cast<uint64_t>(buffer)) == 0) {
+        if (table->GetBufferMemoryRequirements) {
+            VkMemoryRequirements mem_reqs;
+            table->GetBufferMemoryRequirements(device, buffer, &mem_reqs);
+            DeviceMemoryReport::Get().OnRecordResourceSize(reinterpret_cast<uint64_t>(buffer), mem_reqs.size);
+        }
+    }
+    VkResult result = table->BindBufferMemory(device, buffer, memory, memoryOffset);
+    if (result == VK_SUCCESS) {
         DeviceMemoryReport::Get().OnBindBufferMemory(reinterpret_cast<uint64_t>(buffer), reinterpret_cast<uint64_t>(memory), memoryOffset);
     }
     return result;
 }
 
 // Intercept memory binding to correlate image object handles with device memory allocations.
 VKAPI_ATTR VkResult VKAPI_CALL vkBindImageMemory(VkDevice device, VkImage image, VkDeviceMemory memory, VkDeviceSize memoryOffset) {
-    VkResult result = device_dispatch_table(device)->BindImageMemory(device, image, memory, memoryOffset);
-    if (result == VK_SUCCESS && image != VK_NULL_HANDLE && memory != VK_NULL_HANDLE) {
+    auto* table = device_dispatch_table(device);
+    assert(table->BindImageMemory != nullptr);
+    assert(image != VK_NULL_HANDLE);
+    assert(memory != VK_NULL_HANDLE);
+    if (DeviceMemoryReport::Get().GetRecordedResourceSize(reinterpret_cast<uint64_t>(image)) == 0) {
+        if (table->GetImageMemoryRequirements) {
+            VkMemoryRequirements mem_reqs;
+            table->GetImageMemoryRequirements(device, image, &mem_reqs);
+            DeviceMemoryReport::Get().OnRecordResourceSize(reinterpret_cast<uint64_t>(image), mem_reqs.size);
+        }
+    }
+    VkResult result = table->BindImageMemory(device, image, memory, memoryOffset);
+    if (result == VK_SUCCESS) {
         DeviceMemoryReport::Get().OnBindImageMemory(reinterpret_cast<uint64_t>(image), reinterpret_cast<uint64_t>(memory), memoryOffset);
     }
     return result;
 }
 
-static void RecordBufferBindings(uint32_t bindInfoCount, const VkBindBufferMemoryInfo* pBindInfos) {
+static void RecordBufferBindings(VkDevice device, uint32_t bindInfoCount, const VkBindBufferMemoryInfo* pBindInfos) {
+    auto* table = device_dispatch_table(device);
     for (uint32_t i = 0; i < bindInfoCount; ++i) {
-        if (pBindInfos[i].buffer != VK_NULL_HANDLE && pBindInfos[i].memory != VK_NULL_HANDLE) {
-            DeviceMemoryReport::Get().OnBindBufferMemory(reinterpret_cast<uint64_t>(pBindInfos[i].buffer), reinterpret_cast<uint64_t>(pBindInfos[i].memory), pBindInfos[i].memoryOffset);
+        assert(pBindInfos[i].buffer != VK_NULL_HANDLE);
+        assert(pBindInfos[i].memory != VK_NULL_HANDLE);
+        if (DeviceMemoryReport::Get().GetRecordedResourceSize(reinterpret_cast<uint64_t>(pBindInfos[i].buffer)) == 0) {
+            VkMemoryRequirements mem_reqs;
+            table->GetBufferMemoryRequirements(device, pBindInfos[i].buffer, &mem_reqs);
+            DeviceMemoryReport::Get().OnRecordResourceSize(reinterpret_cast<uint64_t>(pBindInfos[i].buffer), mem_reqs.size);
         }
+        DeviceMemoryReport::Get().OnBindBufferMemory(reinterpret_cast<uint64_t>(pBindInfos[i].buffer), reinterpret_cast<uint64_t>(pBindInfos[i].memory), pBindInfos[i].memoryOffset);
     }
 }
 
 // Intercept memory binding via vkBindBufferMemory2 to correlate buffer object handles with device memory allocations.
 VKAPI_ATTR VkResult VKAPI_CALL vkBindBufferMemory2(VkDevice device, uint32_t bindInfoCount, const VkBindBufferMemoryInfo* pBindInfos) {
-    VkResult result = device_dispatch_table(device)->BindBufferMemory2(device, bindInfoCount, pBindInfos);
+    auto* table = device_dispatch_table(device);
+    assert(table->BindBufferMemory2 != nullptr);
+    VkResult result = table->BindBufferMemory2(device, bindInfoCount, pBindInfos);
     if (result == VK_SUCCESS && pBindInfos != nullptr) {
-        RecordBufferBindings(bindInfoCount, pBindInfos);
+        RecordBufferBindings(device, bindInfoCount, pBindInfos);
     }
     return result;
 }
 
 // Intercept memory binding via vkBindBufferMemory2KHR to correlate buffer object handles with device memory allocations.
 VKAPI_ATTR VkResult VKAPI_CALL vkBindBufferMemory2KHR(VkDevice device, uint32_t bindInfoCount, const VkBindBufferMemoryInfo* pBindInfos) {
-    assert(device_dispatch_table(device)->BindBufferMemory2KHR != nullptr);
-    VkResult result = device_dispatch_table(device)->BindBufferMemory2KHR(device, bindInfoCount, pBindInfos);
+    auto* table = device_dispatch_table(device);
+    assert(table->BindBufferMemory2KHR != nullptr);
+    VkResult result = table->BindBufferMemory2KHR(device, bindInfoCount, pBindInfos);
     if (result == VK_SUCCESS && pBindInfos != nullptr) {
-        RecordBufferBindings(bindInfoCount, pBindInfos);
+        RecordBufferBindings(device, bindInfoCount, pBindInfos);
     }
     return result;
 }
 
-static void RecordImageBinds(uint32_t bindInfoCount, const VkBindImageMemoryInfo* pBindInfos) {
+static void RecordImageBinds(VkDevice device, uint32_t bindInfoCount, const VkBindImageMemoryInfo* pBindInfos) {
+    auto* table = device_dispatch_table(device);
     for (uint32_t i = 0; i < bindInfoCount; ++i) {
-        if (pBindInfos[i].image != VK_NULL_HANDLE && pBindInfos[i].memory != VK_NULL_HANDLE) {
-            DeviceMemoryReport::Get().OnBindImageMemory(reinterpret_cast<uint64_t>(pBindInfos[i].image), reinterpret_cast<uint64_t>(pBindInfos[i].memory), pBindInfos[i].memoryOffset);
+        assert(pBindInfos[i].image != VK_NULL_HANDLE);
+        assert(pBindInfos[i].memory != VK_NULL_HANDLE);
+        if (DeviceMemoryReport::Get().GetRecordedResourceSize(reinterpret_cast<uint64_t>(pBindInfos[i].image)) == 0) {
+            bool is_plane_bind = false;
+            for (const auto* header = reinterpret_cast<const VkBaseInStructure*>(pBindInfos[i].pNext);
+                 header != nullptr; header = header->pNext) {
+                if (header->sType == VK_STRUCTURE_TYPE_BIND_IMAGE_PLANE_MEMORY_INFO) {
+                    is_plane_bind = true;
+                    break;
+                }
+            }
+            if (!is_plane_bind) {
+                VkMemoryRequirements mem_reqs;
+                table->GetImageMemoryRequirements(device, pBindInfos[i].image, &mem_reqs);
+                DeviceMemoryReport::Get().OnRecordResourceSize(reinterpret_cast<uint64_t>(pBindInfos[i].image), mem_reqs.size);
+            }
         }
+        DeviceMemoryReport::Get().OnBindImageMemory(reinterpret_cast<uint64_t>(pBindInfos[i].image), reinterpret_cast<uint64_t>(pBindInfos[i].memory), pBindInfos[i].memoryOffset);
     }
 }
 
 // Intercept memory binding via vkBindImageMemory2 to correlate image object handles with device memory allocations.
 VKAPI_ATTR VkResult VKAPI_CALL vkBindImageMemory2(VkDevice device, uint32_t bindInfoCount, const VkBindImageMemoryInfo* pBindInfos) {
-    VkResult result = device_dispatch_table(device)->BindImageMemory2(device, bindInfoCount, pBindInfos);
+    auto* table = device_dispatch_table(device);
+    assert(table->BindImageMemory2 != nullptr);
+    VkResult result = table->BindImageMemory2(device, bindInfoCount, pBindInfos);
     if (result == VK_SUCCESS && pBindInfos != nullptr) {
-        RecordImageBinds(bindInfoCount, pBindInfos);
+        RecordImageBinds(device, bindInfoCount, pBindInfos);
     }
     return result;
 }
 
 // Intercept memory binding via vkBindImageMemory2KHR to correlate image object handles with device memory allocations.
 VKAPI_ATTR VkResult VKAPI_CALL vkBindImageMemory2KHR(VkDevice device, uint32_t bindInfoCount, const VkBindImageMemoryInfo* pBindInfos) {
-    assert(device_dispatch_table(device)->BindImageMemory2KHR != nullptr);
-    VkResult result = device_dispatch_table(device)->BindImageMemory2KHR(device, bindInfoCount, pBindInfos);
+    auto* table = device_dispatch_table(device);
+    assert(table->BindImageMemory2KHR != nullptr);
+    VkResult result = table->BindImageMemory2KHR(device, bindInfoCount, pBindInfos);
     if (result == VK_SUCCESS && pBindInfos != nullptr) {
-        RecordImageBinds(bindInfoCount, pBindInfos);
+        RecordImageBinds(device, bindInfoCount, pBindInfos);
     }
     return result;
 }

layersvt/test/CMakeLists.txt

@@ -61,3 +61,28 @@ foreach(test_item ${LAYER_TEST_FILES})
 
     LayerTest(${test_item})
 endforeach()
+
+# Unit tests for the DeviceMemoryReport layer's Vulkan entry points. The layer sources are linked
+# straight into the test binary and driven against a stub dispatch table, so no Vulkan ICD is
+# needed. The Vulkan loader is deliberately not linked here: the layer defines the vk* entry points
+# itself, and linking the loader as well would make it ambiguous which of the two is being called.
+if (TARGET VkLayer_DeviceMemoryReport)
+    add_executable(test_DeviceMemoryReport_dispatch
+                   test_devicememoryreport_dispatch.cpp
+                   ../device_memory_report/device_memory_report_handwritten_dispatch.cpp
+                   ../device_memory_report/device_memory_report.cpp
+                   ../device_memory_report/device_memory_report_perfetto.cpp
+                   ../perfetto/perfetto.cc
+                   ../vk_layer_table.cpp)
+    target_include_directories(test_DeviceMemoryReport_dispatch PRIVATE .. ../device_memory_report)
+    target_link_libraries(test_DeviceMemoryReport_dispatch
+                          Vulkan::Headers Vulkan::UtilityHeaders GTest::gtest GTest::gtest_main ${CMAKE_DL_LIBS})
+    target_compile_definitions(test_DeviceMemoryReport_dispatch PRIVATE VK_ENABLE_BETA_EXTENSIONS)
+    add_test(NAME test_DeviceMemoryReport_dispatch COMMAND test_DeviceMemoryReport_dispatch)
+    set_target_properties(test_DeviceMemoryReport_dispatch PROPERTIES FOLDER "layers/DeviceMemoryReport/Test")
+
+    if(WIN32 AND (QT_TARGET_TYPE STREQUAL STATIC_LIBRARY))
+      set_property(TARGET test_DeviceMemoryReport_dispatch PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded$<$<CONFIG:Debug>:Debug>")
+    endif()
+endif()
+