Fix trim raytracing invalid device issue

framework/encode/vulkan_handle_wrapper_util.h

@@ -279,6 +279,17 @@ inline void CreateWrappedHandle<PhysicalDeviceWrapper, NoParentWrapper, DeviceWr
     CreateWrappedDispatchHandle<PhysicalDeviceWrapper, DeviceWrapper>(VK_NULL_HANDLE, handle, get_id);
 }
 
+template <>
+inline void CreateWrappedHandle<DeviceWrapper, NoParentWrapper, DeviceMemoryWrapper>(VkDevice device,
+                                                                                     NoParentWrapper::HandleType,
+                                                                                     VkDeviceMemory* handle,
+                                                                                     PFN_GetHandleId get_id)
+{
+    CreateWrappedNonDispatchHandle<DeviceMemoryWrapper>(handle, get_id);
+    auto memory_wrapper           = GetWrapper<DeviceMemoryWrapper>(*handle);
+    memory_wrapper->parent_device = GetWrapper<DeviceWrapper>(device);
+}
+
 template <>
 inline void CreateWrappedHandle<DeviceWrapper, NoParentWrapper, QueueWrapper>(
     VkDevice parent,

framework/encode/vulkan_handle_wrappers.h

@@ -162,9 +162,14 @@ struct EventWrapper : public HandleWrapper<VkEvent>
 
 struct DeviceMemoryWrapper : public HandleWrapper<VkDeviceMemory>
 {
-    uint32_t         memory_type_index{ std::numeric_limits<uint32_t>::max() };
-    VkDeviceSize     allocation_size{ 0 };
-    DeviceWrapper*   map_device{ nullptr };
+    uint32_t     memory_type_index{ std::numeric_limits<uint32_t>::max() };
+    VkDeviceSize allocation_size{ 0 };
+    // This is the device which was used to allocate the memory.
+    // Spec states if the memory can be mapped, the mapping device must be this device.
+    // The device wrapper will be initialized when allocating the memory. Some handling
+    // like VulkanStateTracker::TrackTlasToBlasDependencies may use it before mapping
+    // the memory.
+    DeviceWrapper*   parent_device{ nullptr };
     const void*      mapped_data{ nullptr };
     VkDeviceSize     mapped_offset{ 0 };
     VkDeviceSize     mapped_size{ 0 };

framework/encode/vulkan_state_tracker.cpp

@@ -433,7 +433,6 @@ void VulkanStateTracker::TrackMappedMemory(VkDevice         device,
     assert((device != VK_NULL_HANDLE) && (memory != VK_NULL_HANDLE));
 
     auto wrapper           = GetWrapper<DeviceMemoryWrapper>(memory);
-    wrapper->map_device    = GetWrapper<DeviceWrapper>(device);
     wrapper->mapped_data   = mapped_data;
     wrapper->mapped_offset = mapped_offset;
     wrapper->mapped_size   = mapped_size;
@@ -1561,7 +1560,7 @@ void VulkanStateTracker::TrackTlasToBlasDependencies(uint32_t               comm
             {
                 // If PageGuardManager is not used or if it couldn't find the memory id it means that
                 // we need to map the memory.
-                VkDevice           device        = dev_mem_wrapper->map_device->handle;
+                VkDevice           device        = dev_mem_wrapper->parent_device->handle;
                 const DeviceTable* device_table  = GetDeviceTable(device);
                 const VkDeviceSize map_size      = sizeof(VkAccelerationStructureInstanceKHR) * blas_count;
                 void*              mapped_memory = nullptr;
@@ -1596,7 +1595,7 @@ void VulkanStateTracker::TrackTlasToBlasDependencies(uint32_t               comm
                 // If we had to map the device memory unmap it now
                 if (needs_unmapping)
                 {
-                    VkDevice           device       = dev_mem_wrapper->map_device->handle;
+                    VkDevice           device       = dev_mem_wrapper->parent_device->handle;
                     const DeviceTable* device_table = GetDeviceTable(device);
                     device_table->UnmapMemory(device, dev_mem_wrapper->handle);
                 }

framework/encode/vulkan_state_writer.cpp

@@ -1816,7 +1816,7 @@ void VulkanStateWriter::WriteMappedMemoryState(const VulkanStateTable& state_tab
         if (wrapper->mapped_data != nullptr)
         {
             const VkResult result         = VK_SUCCESS;
-            const auto     device_wrapper = wrapper->map_device;
+            const auto     device_wrapper = wrapper->parent_device;
 
             // Map the replay memory.
             encoder_.EncodeHandleIdValue(device_wrapper->handle_id);