gfxstream: Implement private data extension

The logic is implemented on the guest side as we need special handling
for swapchain and device memory handle types.

Memory handles for coherent memory allocations on the guest can map into
a single handle on the host side, which makes it infeasible to pass
through the extension functions with handle remapping. Swapchain handles
are not passed to the host driver, so we need to keep a separate table
for them as well. Instead of separating the logic based on the handle
type, we manage all the private data set/get calls on the guest side
without encoding the commands to the host.

Test: dEQP-VK.api.object_management.private_data.*
Test: dEQP-VK.wsi.android.swapchain.private_data.*

Reviewed-by: Marcin Radomski <dextero@google.com>
Reviewed-by: Aaron Ruby <aruby@qnx.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/35323>

src/gfxstream/codegen/scripts/cereal/functable.py

@@ -95,6 +95,14 @@ RESOURCE_TRACKER_ENTRIES = [
     "vkSetBufferCollectionImageConstraintsFUCHSIA",
     "vkSetBufferCollectionBufferConstraintsFUCHSIA",
     "vkGetBufferCollectionPropertiesFUCHSIA",
+    "vkSetPrivateData",
+    "vkSetPrivateDataKHR",
+    "vkGetPrivateData",
+    "vkGetPrivateDataKHR",
+    "vkCreatePrivateDataSlot",
+    "vkCreatePrivateDataSlotEXT",
+    "vkDestroyPrivateDataSlot",
+    "vkDestroyPrivateDataSlotEXT",
 ]
 
 SUCCESS_VAL = {

src/gfxstream/guest/vulkan_enc/ResourceTracker.cpp

@@ -1059,6 +1059,13 @@ void ResourceTracker::unregister_VkSampler(VkSampler sampler) {
     info_VkSampler.erase(sampler);
 }
 
+void ResourceTracker::unregister_VkPrivateDataSlot(VkPrivateDataSlot privateSlot) {
+    if (!privateSlot) return;
+
+    std::lock_guard<std::recursive_mutex> lock(mLock);
+    info_VkPrivateDataSlot.erase(privateSlot);
+}
+
 void ResourceTracker::unregister_VkCommandBuffer(VkCommandBuffer commandBuffer) {
     resetCommandBufferStagingInfo(commandBuffer, true /* also reset primaries */,
                                   true /* also clear pending descriptor sets */);
@@ -1566,9 +1573,15 @@ void ResourceTracker::deviceMemoryTransform_tohost(VkDeviceMemory* memory, uint3
 
         for (uint32_t i = 0; i < memoryCount; ++i) {
             VkDeviceMemory mem = memory[i];
+            if (!mem) {
+                return;
+            }
 
             auto it = info_VkDeviceMemory.find(mem);
-            if (it == info_VkDeviceMemory.end()) return;
+            if (it == info_VkDeviceMemory.end()) {
+                mesa_logw("%s cannot find memory %p!", __func__, mem);
+                return;
+            }
 
             const auto& info = it->second;
 
@@ -1583,11 +1596,6 @@ void ResourceTracker::deviceMemoryTransform_tohost(VkDeviceMemory* memory, uint3
             if (size && size[i] == VK_WHOLE_SIZE) {
                 size[i] = info.allocationSize;
             }
-
-            // TODO
-            (void)memory;
-            (void)offset;
-            (void)size;
         }
     }
 }
@@ -4998,6 +5006,91 @@ VkResult ResourceTracker::on_vkWaitForFences(void* context, VkResult, VkDevice d
 #endif
 }
 
+VkResult ResourceTracker::on_vkSetPrivateData(void* context, VkResult input_result, VkDevice device,
+                                              VkObjectType objectType, uint64_t objectHandle,
+                                              VkPrivateDataSlot privateDataSlot, uint64_t data) {
+    if (input_result != VK_SUCCESS) return input_result;
+
+    VkPrivateDataSlot_Info::PrivateDataKey key = std::make_pair(objectHandle, objectType);
+
+    std::lock_guard<std::recursive_mutex> lock(mLock);
+    auto it = info_VkPrivateDataSlot.find(privateDataSlot);
+
+    // Do not forward calls with invalid handles to host.
+    if (it == info_VkPrivateDataSlot.end()) {
+        return VK_ERROR_OUT_OF_HOST_MEMORY;
+    }
+
+    auto& slotInfoTable = it->second.privateDataTable;
+    slotInfoTable[key] = data;
+    return VK_SUCCESS;
+}
+
+VkResult ResourceTracker::on_vkSetPrivateDataEXT(void* context, VkResult input_result,
+                                                 VkDevice device, VkObjectType objectType,
+                                                 uint64_t objectHandle,
+                                                 VkPrivateDataSlot privateDataSlot, uint64_t data) {
+    return on_vkSetPrivateData(context, input_result, device, objectType, objectHandle,
+                               privateDataSlot, data);
+}
+
+void ResourceTracker::on_vkGetPrivateData(void* context, VkDevice device, VkObjectType objectType,
+                                          uint64_t objectHandle, VkPrivateDataSlot privateDataSlot,
+                                          uint64_t* pData) {
+    VkPrivateDataSlot_Info::PrivateDataKey key = std::make_pair(objectHandle, objectType);
+
+    std::lock_guard<std::recursive_mutex> lock(mLock);
+    auto it = info_VkPrivateDataSlot.find(privateDataSlot);
+
+    // Do not forward calls with invalid handles to host.
+    if (it == info_VkPrivateDataSlot.end()) {
+        return;
+    }
+
+    auto& slotInfoTable = it->second.privateDataTable;
+    *pData = slotInfoTable[key];
+}
+
+void ResourceTracker::on_vkGetPrivateDataEXT(void* context, VkDevice device,
+                                             VkObjectType objectType, uint64_t objectHandle,
+                                             VkPrivateDataSlot privateDataSlot, uint64_t* pData) {
+    return on_vkGetPrivateData(context, device, objectType, objectHandle, privateDataSlot, pData);
+}
+
+VkResult ResourceTracker::on_vkCreatePrivateDataSlot(void* context, VkResult input_result,
+                                                     VkDevice device,
+                                                     const VkPrivateDataSlotCreateInfo* pCreateInfo,
+                                                     const VkAllocationCallbacks* pAllocator,
+                                                     VkPrivateDataSlot* pPrivateDataSlot) {
+    if (input_result != VK_SUCCESS) {
+        return input_result;
+    }
+    VkEncoder* enc = (VkEncoder*)context;
+    return enc->vkCreatePrivateDataSlot(device, pCreateInfo, pAllocator, pPrivateDataSlot,
+                                        true /* do lock */);
+}
+VkResult ResourceTracker::on_vkCreatePrivateDataSlotEXT(
+    void* context, VkResult input_result, VkDevice device,
+    const VkPrivateDataSlotCreateInfo* pCreateInfo, const VkAllocationCallbacks* pAllocator,
+    VkPrivateDataSlot* pPrivateDataSlot) {
+    return on_vkCreatePrivateDataSlot(context, input_result, device, pCreateInfo, pAllocator,
+                                      pPrivateDataSlot);
+}
+
+void ResourceTracker::on_vkDestroyPrivateDataSlot(void* context, VkDevice device,
+                                                  VkPrivateDataSlot privateDataSlot,
+                                                  const VkAllocationCallbacks* pAllocator) {
+    if (!privateDataSlot) return;
+
+    VkEncoder* enc = (VkEncoder*)context;
+    enc->vkDestroyPrivateDataSlot(device, privateDataSlot, pAllocator, true /* do lock */);
+}
+void ResourceTracker::on_vkDestroyPrivateDataSlotEXT(void* context, VkDevice device,
+                                                     VkPrivateDataSlot privateDataSlot,
+                                                     const VkAllocationCallbacks* pAllocator) {
+    return on_vkDestroyPrivateDataSlot(context, device, privateDataSlot, pAllocator);
+}
+
 VkResult ResourceTracker::on_vkCreateDescriptorPool(void* context, VkResult, VkDevice device,
                                                     const VkDescriptorPoolCreateInfo* pCreateInfo,
                                                     const VkAllocationCallbacks* pAllocator,

src/gfxstream/guest/vulkan_enc/ResourceTracker.h

@@ -395,6 +395,35 @@ class ResourceTracker {
                                 uint32_t fenceCount, const VkFence* pFences, VkBool32 waitAll,
                                 uint64_t timeout);
 
+    VkResult on_vkSetPrivateData(void* context, VkResult input_result, VkDevice device,
+                                 VkObjectType objectType, uint64_t objectHandle,
+                                 VkPrivateDataSlot privateDataSlot, uint64_t data);
+    VkResult on_vkSetPrivateDataEXT(void* context, VkResult input_result, VkDevice device,
+                                    VkObjectType objectType, uint64_t objectHandle,
+                                    VkPrivateDataSlot privateDataSlot, uint64_t data);
+
+    void on_vkGetPrivateData(void* context, VkDevice device, VkObjectType objectType,
+                             uint64_t objectHandle, VkPrivateDataSlot privateDataSlot,
+                             uint64_t* pData);
+    void on_vkGetPrivateDataEXT(void* context, VkDevice device, VkObjectType objectType,
+                                uint64_t objectHandle, VkPrivateDataSlot privateDataSlot,
+                                uint64_t* pData);
+
+    VkResult on_vkCreatePrivateDataSlot(void* context, VkResult input_result, VkDevice device,
+                                        const VkPrivateDataSlotCreateInfo* pCreateInfo,
+                                        const VkAllocationCallbacks* pAllocator,
+                                        VkPrivateDataSlot* pPrivateDataSlot);
+    VkResult on_vkCreatePrivateDataSlotEXT(void* context, VkResult input_result, VkDevice device,
+                                           const VkPrivateDataSlotCreateInfo* pCreateInfo,
+                                           const VkAllocationCallbacks* pAllocator,
+                                           VkPrivateDataSlot* pPrivateDataSlot);
+    void on_vkDestroyPrivateDataSlot(void* context, VkDevice device,
+                                     VkPrivateDataSlot privateDataSlot,
+                                     const VkAllocationCallbacks* pAllocator);
+    void on_vkDestroyPrivateDataSlotEXT(void* context, VkDevice device,
+                                        VkPrivateDataSlot privateDataSlot,
+                                        const VkAllocationCallbacks* pAllocator);
+
     VkResult on_vkCreateDescriptorPool(void* context, VkResult input_result, VkDevice device,
                                        const VkDescriptorPoolCreateInfo* pCreateInfo,
                                        const VkAllocationCallbacks* pAllocator,
@@ -887,6 +916,28 @@ class ResourceTracker {
         uint32_t unused;
     };
 
+    struct VkPrivateDataSlot_Info {
+        // We need special handling for device memory and swapchain object types for private data
+        // management. For memory, we can use a single handle on the host side, so setting a
+        // private data slot for guest handle can also set any other data set previously.
+        // For swapchains, we don't actually get create/destroy calls to keep track of object
+        // handles to be able to pass the call to the underlying host driver. Rather than handling
+        // the 2 cases separately, we handle all the private data management directly here with a
+        // single table, so vkSetPrivateData and vkGetPrivateData calls don't need to be encoded for
+        // the host.
+        typedef std::pair<uint64_t, VkObjectType> PrivateDataKey;
+        struct PrivateDataKeyHash {
+            template <class T1, class T2>
+            std::size_t operator()(const std::pair<T1, T2>& p) const {
+                std::size_t h1 = std::hash<T1>{}(p.first);
+                std::size_t h2 = std::hash<T2>{}(p.second);
+                return h1 ^ h2;
+            }
+        };
+
+        std::unordered_map<PrivateDataKey, uint64_t, PrivateDataKeyHash> privateDataTable;
+    };
+
     struct VkBufferCollectionFUCHSIA_Info {
 #ifdef VK_USE_PLATFORM_FUCHSIA
         std::optional<fuchsia_sysmem::wire::BufferCollectionConstraints> constraints;

src/gfxstream/guest/vulkan_enc/VulkanHandles.h

@@ -104,7 +104,6 @@ namespace vk {
     f(VkValidationCacheEXT)                                                           \
     f(VkDebugReportCallbackEXT)                                                       \
     f(VkDebugUtilsMessengerEXT)                                                       \
-    f(VkPrivateDataSlot)                                                              \
     f(VkMicromapEXT)                                                                  \
     __GOLDFISH_VK_LIST_NON_DISPATCHABLE_HANDLE_TYPES_NVX_BINARY_IMPORT(f)             \
     __GOLDFISH_VK_LIST_NON_DISPATCHABLE_HANDLE_TYPES_NVX_DEVICE_GENERATED_COMMANDS(f) \
@@ -124,6 +123,7 @@ namespace vk {
     f(VkDescriptorSetLayout)                                    \
     f(VkCommandPool)                                            \
     f(VkSampler)                                                \
+    f(VkPrivateDataSlot)                                        \
     __GOLDFISH_VK_LIST_NON_DISPATCHABLE_HANDLE_TYPES_FUCHSIA(f) \
     GOLDFISH_VK_LIST_TRIVIAL_NON_DISPATCHABLE_HANDLE_TYPES(f)
 
@@ -150,6 +150,7 @@ namespace vk {
     f(VkDescriptorUpdateTemplate)                                            \
     f(VkCommandPool)                                                         \
     f(VkSampler)                                                             \
+    f(VkPrivateDataSlot)                                                     \
     __GOLDFISH_VK_LIST_NON_DISPATCHABLE_HANDLE_TYPES_FUCHSIA(f)              \
     GOLDFISH_VK_LIST_TRIVIAL_NON_DISPATCHABLE_HANDLE_TYPES(f)