gfxstream: answer a deferred AHB image layout from the AHB itself

host/vulkan/vk_common_operations.cpp

@@ -1767,6 +1767,14 @@ void VkEmulation::initFeatures(Features features) {
 VkEmulation::~VkEmulation() {
     std::lock_guard<std::mutex> lock(mMutex);
 
+#ifdef __ANDROID__
+    // Nothing claimed these before teardown; release the extra reference each was stashed with.
+    for (auto& [key, ahb] : mPendingDeferredLayoutAhbs) {
+        AHardwareBuffer_release(ahb);
+    }
+    mPendingDeferredLayoutAhbs.clear();
+#endif
+
     mCompositorVk.reset();
     mDisplayVk.reset();
     mUdmabufCreator.reset();
@@ -2062,9 +2070,10 @@ MTLResource_id VkEmulation::getMtlResourceFromVkDeviceMemory(VulkanDispatch* vk,
 #endif
 
 #ifdef __ANDROID__
-// Allocate an AHardwareBuffer matching the given image's format, extent and usage.
-// Returns nullptr on failure (caller falls back to the non-AHB allocation path).
-static AHardwareBuffer* allocAhb(const VkImageCreateInfo* imageCreateInfo) {
+// Derives the AHB shape (format, usage, extent) that allocAhb() would allocate for this image,
+// without allocating anything. Shared with the deferred-layout pending-AHB pool so a probe AHB
+// and a ColorBuffer's own AHB request can be matched without ever calling allocAhb() twice.
+AhbShapeKey ComputeAhbShapeKey(const VkImageCreateInfo* imageCreateInfo) {
     // Map VkFormat to the corresponding AHB format — must match to avoid tiling mismatch.
     uint32_t ahbFormat;
     switch (imageCreateInfo->format) {
@@ -2099,12 +2108,25 @@ static AHardwareBuffer* allocAhb(const VkImageCreateInfo* imageCreateInfo) {
         ahbUsage = AHARDWAREBUFFER_USAGE_GPU_FRAMEBUFFER | AHARDWAREBUFFER_USAGE_GPU_SAMPLED_IMAGE;
     }
 
-    AHardwareBuffer_Desc desc = {
+    return AhbShapeKey{
         .width = imageCreateInfo->extent.width,
         .height = imageCreateInfo->extent.height,
+        .ahbFormat = ahbFormat,
+        .ahbUsage = ahbUsage,
+    };
+}
+
+// Allocate an AHardwareBuffer matching the given image's format, extent and usage.
+// Returns nullptr on failure (caller falls back to the non-AHB allocation path).
+AHardwareBuffer* allocAhb(const VkImageCreateInfo* imageCreateInfo) {
+    const AhbShapeKey shape = ComputeAhbShapeKey(imageCreateInfo);
+
+    AHardwareBuffer_Desc desc = {
+        .width = shape.width,
+        .height = shape.height,
         .layers = 1,
-        .format = ahbFormat,
-        .usage = ahbUsage,
+        .format = shape.ahbFormat,
+        .usage = shape.ahbUsage,
     };
 
     AHardwareBuffer* ahb = nullptr;
@@ -2116,6 +2138,37 @@ static AHardwareBuffer* allocAhb(const VkImageCreateInfo* imageCreateInfo) {
     }
     return ahb;
 }
+
+void VkEmulation::stashPendingDeferredLayoutAhb(const VkImageCreateInfo* imageCreateInfo,
+                                                AHardwareBuffer* ahb) {
+    AHardwareBuffer_acquire(ahb);
+    std::lock_guard<std::mutex> lock(mMutex);
+    mPendingDeferredLayoutAhbs.emplace(ComputeAhbShapeKey(imageCreateInfo), ahb);
+}
+
+void VkEmulation::unstashPendingDeferredLayoutAhb(const VkImageCreateInfo* imageCreateInfo,
+                                                  AHardwareBuffer* ahb) {
+    std::lock_guard<std::mutex> lock(mMutex);
+    auto range = mPendingDeferredLayoutAhbs.equal_range(ComputeAhbShapeKey(imageCreateInfo));
+    for (auto it = range.first; it != range.second; ++it) {
+        if (it->second == ahb) {
+            mPendingDeferredLayoutAhbs.erase(it);
+            AHardwareBuffer_release(ahb);
+            return;
+        }
+    }
+}
+
+AHardwareBuffer* VkEmulation::takePendingDeferredLayoutAhbLocked(
+    const VkImageCreateInfo* imageCreateInfo) {
+    auto it = mPendingDeferredLayoutAhbs.find(ComputeAhbShapeKey(imageCreateInfo));
+    if (it == mPendingDeferredLayoutAhbs.end()) {
+        return nullptr;
+    }
+    AHardwareBuffer* ahb = it->second;
+    mPendingDeferredLayoutAhbs.erase(it);
+    return ahb;
+}
 #endif
 
 // Precondition: sVkEmulation has valid device support info
@@ -2341,7 +2394,19 @@ bool VkEmulation::allocExternalMemory(VulkanDispatch* vk, VkEmulation::ExternalM
             if (colorBufferInfo) {
                 auto cbInfoPtr = *colorBufferInfo;
 
-                AHardwareBuffer* ahb = allocAhb(&cbInfoPtr->imageCreateInfoShallow);
+                // A same-shaped image that hit a deferred layout query earlier (see
+                // on_vkCreateImage) may have left its probe AHB here. Only probes not already
+                // imported as their own image's backing memory are still in the pool, so
+                // adopting one cannot alias -- see unstashPendingDeferredLayoutAhb().
+                AHardwareBuffer* ahb =
+                    takePendingDeferredLayoutAhbLocked(&cbInfoPtr->imageCreateInfoShallow);
+                if (ahb) {
+                    GFXSTREAM_INFO("DL-AHB adopted pending ahb=%p for ColorBuffer %u (%ux%u)",
+                                   (void*)ahb, cbInfoPtr->handle, cbInfoPtr->width,
+                                   cbInfoPtr->height);
+                } else {
+                    ahb = allocAhb(&cbInfoPtr->imageCreateInfoShallow);
+                }
                 if (!ahb) {
                     GFXSTREAM_WARNING(
                         "Falling back to non-exportable allocation for ColorBuffer %u (%ux%u).",

host/vulkan/vk_common_operations.h

@@ -13,6 +13,10 @@
 // limitations under the License.
 #pragma once
 
+#ifdef __ANDROID__
+#include <android/hardware_buffer.h>
+#endif
+
 #include <GLES2/gl2.h>
 #include <vulkan/vulkan.h>
 
@@ -85,6 +89,30 @@ enum class AstcEmulationMode {
     Gpu,       // Decompress ASTC textures on the GPU
 };
 
+#ifdef __ANDROID__
+// The AHardwareBuffer_Desc fields that fully determine an AHB's allocation/property-query
+// result, derived from a VkImageCreateInfo the same way allocAhb() derives its own AHB. Used to
+// match an image against an AHB allocated for a different image of the identical shape.
+struct AhbShapeKey {
+    uint32_t width = 0;
+    uint32_t height = 0;
+    uint32_t ahbFormat = 0;
+    uint64_t ahbUsage = 0;
+
+    bool operator==(const AhbShapeKey& other) const {
+        return width == other.width && height == other.height && ahbFormat == other.ahbFormat &&
+               ahbUsage == other.ahbUsage;
+    }
+};
+struct AhbShapeKeyHash {
+    size_t operator()(const AhbShapeKey& key) const {
+        return (static_cast<size_t>(key.width) << 48) ^ (static_cast<size_t>(key.height) << 32) ^
+               (static_cast<size_t>(key.ahbFormat) << 16) ^ static_cast<size_t>(key.ahbUsage);
+    }
+};
+AhbShapeKey ComputeAhbShapeKey(const VkImageCreateInfo* imageCreateInfo);
+#endif
+
 // Global state that holds a global Vulkan instance along with globally
 // exported memory allocations + images. This is in order to service things
 // like AndroidHardwareBuffer/FuchsiaImagePipeHandle. Each such allocation is
@@ -205,6 +233,7 @@ class VkEmulation {
     void onVkDeviceLost();
 
     VkExternalMemoryHandleTypeFlagBits getDefaultExternalMemoryHandleType();
+
     void appendExternalMemoryModeDeviceExtensions(std::vector<const char*>& outDeviceExtensions);
     ExternalMemory::Mode getExternalMemoryMode() const;
     bool supportsExternalMemory() {
@@ -466,6 +495,23 @@ class VkEmulation {
 
     uint32_t vulkanInstanceVersion() const;
 
+#ifdef __ANDROID__
+    // Stashes an AHB allocated to answer a deferred AHB-external image's zero-size layout query
+    // (see on_vkCreateImage), so that a same-shaped ColorBuffer's own backing allocation can
+    // adopt it instead of allocating a second one -- see allocExternalMemory's AndroidAHB case.
+    // Takes its own reference via AHardwareBuffer_acquire, independent of whatever reference the
+    // caller (the probing image's own DeferredLayoutInfo) already holds.
+    void stashPendingDeferredLayoutAhb(const VkImageCreateInfo* imageCreateInfo,
+                                       AHardwareBuffer* ahb);
+
+    // Withdraws a previously stashed AHB from the pending pool, releasing the pool's reference.
+    // Must be called once the probing image imports the AHB as its own backing memory: at that
+    // point it is in use, and letting a ColorBuffer adopt it would alias two unrelated images
+    // onto one allocation. No-op if it was already adopted or never stashed.
+    void unstashPendingDeferredLayoutAhb(const VkImageCreateInfo* imageCreateInfo,
+                                         AHardwareBuffer* ahb);
+#endif
+
    private:
     VkEmulation() = default;
 
@@ -707,6 +753,23 @@ class VkEmulation {
     // memory handles.
     std::unordered_map<uint32_t, ExternalMemoryInfo> mExternalMemories GUARDED_BY(mMutex);
 
+#ifdef __ANDROID__
+    // AHBs stashed by stashPendingDeferredLayoutAhb(), each holding its own
+    // AHardwareBuffer_acquire()'d reference, waiting for a same-shaped ColorBuffer to adopt
+    // instead of allocating a second AHB (see allocExternalMemory's AndroidAHB case). Any left
+    // unclaimed are released at VkEmulation teardown.
+    std::unordered_multimap<AhbShapeKey, AHardwareBuffer*, AhbShapeKeyHash>
+        mPendingDeferredLayoutAhbs GUARDED_BY(mMutex);
+
+    // Pops one AHB matching imageCreateInfo's shape from mPendingDeferredLayoutAhbs, or nullptr
+    // if none is pending. Caller already holds mMutex -- NO_THREAD_SAFETY_ANALYSIS rather than
+    // REQUIRES(mMutex) because its only caller, allocExternalMemory(), predates thread-safety
+    // annotations itself and isn't annotated, so REQUIRES here would just move the same
+    // unverifiable-by-the-analyzer assumption one frame up instead of removing it.
+    AHardwareBuffer* takePendingDeferredLayoutAhbLocked(const VkImageCreateInfo* imageCreateInfo)
+        NO_THREAD_SAFETY_ANALYSIS;
+#endif
+
     // The host keeps a set of occupied guest memory addresses to avoid a
     // host memory address mapped to guest twice.
     std::unordered_set<uint64_t> mOccupiedGpas GUARDED_BY(mMutex);
@@ -740,6 +803,11 @@ class VkEmulation {
     std::unique_ptr<UdmabufCreator> mUdmabufCreator;
 };
 
+#ifdef __ANDROID__
+// Allocates an AHardwareBuffer matching an image's create info.
+AHardwareBuffer* allocAhb(const VkImageCreateInfo* imageCreateInfo);
+#endif
+
 }  // namespace vk
 }  // namespace host
 }  // namespace gfxstream

host/vulkan/vk_decoder.cpp

@@ -3399,14 +3399,12 @@ size_t VkDecoder::Impl::decode(void* buf, size_t len, IOStream* ioStream,
                 VkImage image;
                 const VkImageSubresource* pSubresource;
                 VkSubresourceLayout* pLayout;
-                // Begin non wrapped dispatchable handle unboxing for device;
+                // Begin global wrapped dispatchable handle unboxing for device;
                 uint64_t cgen_var_0;
                 memcpy((uint64_t*)&cgen_var_0, *readStreamPtrPtr, 1 * 8);
                 *readStreamPtrPtr += 1 * 8;
                 *(VkDevice*)&device = (VkDevice)(VkDevice)((VkDevice)(*&cgen_var_0));
-                auto unboxed_device = unbox_VkDevice(device);
                 auto vk = dispatch_VkDevice(device);
-                // End manual dispatchable handle unboxing for device;
                 uint64_t cgen_var_1;
                 memcpy((uint64_t*)&cgen_var_1, *readStreamPtrPtr, 1 * 8);
                 *readStreamPtrPtr += 1 * 8;
@@ -3435,7 +3433,8 @@ size_t VkDecoder::Impl::decode(void* buf, size_t len, IOStream* ioStream,
                         (unsigned long long)pSubresource, (unsigned long long)pLayout);
                 }
                 if (CC_LIKELY(vk)) {
-                    vk->vkGetImageSubresourceLayout(unboxed_device, image, pSubresource, pLayout);
+                    m_state->on_vkGetImageSubresourceLayout(&m_pool, snapshotApiCallHandle, device,
+                                                            image, pSubresource, pLayout);
                 }
                 vkStream->unsetHandleMapping();
                 if (pLayout) {