Merge "Snapshot VkBuffer" into mainHEADmastermain

7 files changed, 558 insertions, 84 deletions

diff --git a/codegen/vulkan/vulkan-docs-next/scripts/cereal/decodersnapshot.py b/codegen/vulkan/vulkan-docs-next/scripts/cereal/decodersnapshot.py
index dc708502..fd2611b3 100644
--- a/codegen/vulkan/vulkan-docs-next/scripts/cereal/decodersnapshot.py
+++ b/codegen/vulkan/vulkan-docs-next/scripts/cereal/decodersnapshot.py
@@ -146,7 +146,13 @@ def extract_deps_vkCreateFramebuffer(param, access, lenExpr, api, cgen):
 
 def extract_deps_vkBindImageMemory(param, access, lenExpr, api, cgen):
     cgen.stmt("mReconstruction.addHandleDependency((const uint64_t*)%s, %s, (uint64_t)(uintptr_t)%s, VkReconstruction::BOUND_MEMORY)" % \
-              (access, lenExpr, "(uint64_t)(uintptr_t)unboxed_to_boxed_non_dispatchable_VkDeviceMemory(memory)"))
+              (access, lenExpr, "unboxed_to_boxed_non_dispatchable_VkDeviceMemory(memory)"))
+    cgen.stmt("mReconstruction.addHandleDependency((const uint64_t*)%s, %s, (uint64_t)(uintptr_t)((%s)[0]), VkReconstruction::BOUND_MEMORY)" % \
+              (access, lenExpr, access))
+
+def extract_deps_vkBindBufferMemory(param, access, lenExpr, api, cgen):
+    cgen.stmt("mReconstruction.addHandleDependency((const uint64_t*)%s, %s, (uint64_t)(uintptr_t)%s, VkReconstruction::BOUND_MEMORY)" % \
+              (access, lenExpr, "unboxed_to_boxed_non_dispatchable_VkDeviceMemory(memory)"))
     cgen.stmt("mReconstruction.addHandleDependency((const uint64_t*)%s, %s, (uint64_t)(uintptr_t)((%s)[0]), VkReconstruction::BOUND_MEMORY)" % \
               (access, lenExpr, access))
 
@@ -157,6 +163,7 @@ specialCaseDependencyExtractors = {
     "vkCreateGraphicsPipelines" : extract_deps_vkCreateGraphicsPipelines,
     "vkCreateFramebuffer" : extract_deps_vkCreateFramebuffer,
     "vkBindImageMemory": extract_deps_vkBindImageMemory,
+    "vkBindBufferMemory": extract_deps_vkBindBufferMemory,
 }
 
 apiSequences = {
@@ -168,9 +175,10 @@ class VkObjectState:
     vk_object : str
     state : str = "VkReconstruction::CREATED"
 
-# TODO: add vkBindImageMemory2 into this list
+# TODO: add vkBindImageMemory2 and vkBindBufferMemory2 into this list
 apiChangeState = {
     "vkBindImageMemory": VkObjectState("image", "VkReconstruction::BOUND_MEMORY"),
+    "vkBindBufferMemory": VkObjectState("buffer", "VkReconstruction::BOUND_MEMORY"),
 }
 
 apiModifies = {
diff --git a/common/end2end/Android.bp b/common/end2end/Android.bp
index 568f1732..9d35c3ff 100644
--- a/common/end2end/Android.bp
+++ b/common/end2end/Android.bp
@@ -13,6 +13,7 @@ cc_test_host {
         "GfxstreamEnd2EndGlTests.cpp",
         "GfxstreamEnd2EndVkTests.cpp",
         "GfxstreamEnd2EndVkSnapshotBasicTests.cpp",
+        "GfxstreamEnd2EndVkSnapshotBufferTests.cpp",
         "GfxstreamEnd2EndVkSnapshotImageTests.cpp",
         "GfxstreamEnd2EndVkSnapshotPipelineTests.cpp",
     ],
diff --git a/common/end2end/GfxstreamEnd2EndVkSnapshotBufferTests.cpp b/common/end2end/GfxstreamEnd2EndVkSnapshotBufferTests.cpp
new file mode 100644
index 00000000..af0a75f2
--- /dev/null
+++ b/common/end2end/GfxstreamEnd2EndVkSnapshotBufferTests.cpp
@@ -0,0 +1,228 @@
+// Copyright (C) 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <string>
+
+#include "GfxstreamEnd2EndTests.h"
+#include "gfxstream/RutabagaLayerTestUtils.h"
+
+namespace gfxstream {
+namespace tests {
+namespace {
+
+using testing::Eq;
+using testing::Ge;
+using testing::IsEmpty;
+using testing::IsNull;
+using testing::Not;
+using testing::NotNull;
+
+class GfxstreamEnd2EndVkSnapshotBufferTest : public GfxstreamEnd2EndTest {};
+
+TEST_P(GfxstreamEnd2EndVkSnapshotBufferTest, BufferContent) {
+    static constexpr vkhpp::DeviceSize kSize = 256;
+
+    std::vector<uint8_t> srcBufferContent(kSize);
+    for (size_t i = 0; i < kSize; i++) {
+        srcBufferContent[i] = static_cast<uint8_t>(i & 0xff);
+    }
+    auto [instance, physicalDevice, device, queue, queueFamilyIndex] =
+        VK_ASSERT(SetUpTypicalVkTestEnvironment());
+
+    // Staging buffer
+    const vkhpp::BufferCreateInfo stagingBufferCreateInfo = {
+        .size = static_cast<VkDeviceSize>(kSize),
+        .usage = vkhpp::BufferUsageFlagBits::eTransferSrc,
+        .sharingMode = vkhpp::SharingMode::eExclusive,
+    };
+    auto stagingBuffer = device->createBufferUnique(stagingBufferCreateInfo).value;
+    ASSERT_THAT(stagingBuffer, IsValidHandle());
+
+    vkhpp::MemoryRequirements stagingBufferMemoryRequirements{};
+    device->getBufferMemoryRequirements(*stagingBuffer, &stagingBufferMemoryRequirements);
+
+    const auto stagingBufferMemoryType = GetMemoryType(
+        physicalDevice, stagingBufferMemoryRequirements,
+        vkhpp::MemoryPropertyFlagBits::eHostVisible | vkhpp::MemoryPropertyFlagBits::eHostCoherent);
+
+    // Staging memory
+    const vkhpp::MemoryAllocateInfo stagingBufferMemoryAllocateInfo = {
+        .allocationSize = stagingBufferMemoryRequirements.size,
+        .memoryTypeIndex = stagingBufferMemoryType,
+    };
+    auto stagingBufferMemory = device->allocateMemoryUnique(stagingBufferMemoryAllocateInfo).value;
+    ASSERT_THAT(stagingBufferMemory, IsValidHandle());
+    ASSERT_THAT(device->bindBufferMemory(*stagingBuffer, *stagingBufferMemory, 0), IsVkSuccess());
+
+    // Fill memory content
+    void* mapped = nullptr;
+    auto mapResult =
+        device->mapMemory(*stagingBufferMemory, 0, VK_WHOLE_SIZE, vkhpp::MemoryMapFlags{}, &mapped);
+    ASSERT_THAT(mapResult, IsVkSuccess());
+    ASSERT_THAT(mapped, NotNull());
+
+    auto* bytes = reinterpret_cast<uint8_t*>(mapped);
+    std::memcpy(bytes, srcBufferContent.data(), kSize);
+
+    const vkhpp::MappedMemoryRange range = {
+        .memory = *stagingBufferMemory,
+        .offset = 0,
+        .size = kSize,
+    };
+    device->unmapMemory(*stagingBufferMemory);
+
+    // Vertex buffer
+    const vkhpp::BufferCreateInfo vertexBufferCreateInfo = {
+        .size = static_cast<VkDeviceSize>(kSize),
+        .usage = vkhpp::BufferUsageFlagBits::eVertexBuffer,
+        .sharingMode = vkhpp::SharingMode::eExclusive,
+    };
+    auto vertexBuffer = device->createBufferUnique(vertexBufferCreateInfo).value;
+    ASSERT_THAT(vertexBuffer, IsValidHandle());
+
+    vkhpp::MemoryRequirements vertexBufferMemoryRequirements{};
+    device->getBufferMemoryRequirements(*vertexBuffer, &vertexBufferMemoryRequirements);
+
+    const auto vertexBufferMemoryType = GetMemoryType(
+        physicalDevice, vertexBufferMemoryRequirements,
+        vkhpp::MemoryPropertyFlagBits::eHostVisible | vkhpp::MemoryPropertyFlagBits::eHostCoherent);
+
+    // Vertex memory
+    const vkhpp::MemoryAllocateInfo vertexBufferMemoryAllocateInfo = {
+        .allocationSize = vertexBufferMemoryRequirements.size,
+        .memoryTypeIndex = vertexBufferMemoryType,
+    };
+    auto vertexBufferMemory = device->allocateMemoryUnique(vertexBufferMemoryAllocateInfo).value;
+    ASSERT_THAT(vertexBufferMemory, IsValidHandle());
+    ASSERT_THAT(device->bindBufferMemory(*vertexBuffer, *vertexBufferMemory, 0), IsVkSuccess());
+
+    // Command buffer
+    const vkhpp::CommandPoolCreateInfo commandPoolCreateInfo = {
+        .queueFamilyIndex = queueFamilyIndex,
+    };
+
+    auto commandPool = device->createCommandPoolUnique(commandPoolCreateInfo).value;
+    ASSERT_THAT(commandPool, IsValidHandle());
+
+    const vkhpp::CommandBufferAllocateInfo commandBufferAllocateInfo = {
+        .level = vkhpp::CommandBufferLevel::ePrimary,
+        .commandPool = *commandPool,
+        .commandBufferCount = 1,
+    };
+    auto commandBuffers = device->allocateCommandBuffersUnique(commandBufferAllocateInfo).value;
+    ASSERT_THAT(commandBuffers, Not(IsEmpty()));
+    auto commandBuffer = std::move(commandBuffers[0]);
+    ASSERT_THAT(commandBuffer, IsValidHandle());
+
+    const vkhpp::CommandBufferBeginInfo commandBufferBeginInfo = {
+        .flags = vkhpp::CommandBufferUsageFlagBits::eOneTimeSubmit,
+    };
+    commandBuffer->begin(commandBufferBeginInfo);
+    const vkhpp::BufferCopy bufferCopy = {
+        .size = kSize,
+    };
+    commandBuffer->copyBuffer(*stagingBuffer, *vertexBuffer, 1, &bufferCopy);
+    commandBuffer->end();
+
+    auto transferFence = device->createFenceUnique(vkhpp::FenceCreateInfo()).value;
+    ASSERT_THAT(transferFence, IsValidHandle());
+
+    // Execute the command to copy image
+    const vkhpp::SubmitInfo submitInfo = {
+        .commandBufferCount = 1,
+        .pCommandBuffers = &commandBuffer.get(),
+    };
+    queue.submit(submitInfo, *transferFence);
+
+    auto waitResult = device->waitForFences(*transferFence, VK_TRUE, 3000000000L);
+    ASSERT_THAT(waitResult, IsVkSuccess());
+
+    // Snapshot
+    SnapshotSaveAndLoad();
+
+    // Read-back buffer
+    const vkhpp::BufferCreateInfo readbackBufferCreateInfo = {
+        .size = static_cast<VkDeviceSize>(kSize),
+        .usage = vkhpp::BufferUsageFlagBits::eTransferDst,
+        .sharingMode = vkhpp::SharingMode::eExclusive,
+    };
+    auto readbackBuffer = device->createBufferUnique(readbackBufferCreateInfo).value;
+    ASSERT_THAT(readbackBuffer, IsValidHandle());
+
+    vkhpp::MemoryRequirements readbackBufferMemoryRequirements{};
+    device->getBufferMemoryRequirements(*readbackBuffer, &readbackBufferMemoryRequirements);
+
+    const auto readbackBufferMemoryType = GetMemoryType(
+        physicalDevice, readbackBufferMemoryRequirements,
+        vkhpp::MemoryPropertyFlagBits::eHostVisible | vkhpp::MemoryPropertyFlagBits::eHostCoherent);
+
+    // Read-back memory
+    const vkhpp::MemoryAllocateInfo readbackBufferMemoryAllocateInfo = {
+        .allocationSize = readbackBufferMemoryRequirements.size,
+        .memoryTypeIndex = readbackBufferMemoryType,
+    };
+    auto readbackBufferMemory =
+        device->allocateMemoryUnique(readbackBufferMemoryAllocateInfo).value;
+    ASSERT_THAT(readbackBufferMemory, IsValidHandle());
+    ASSERT_THAT(device->bindBufferMemory(*readbackBuffer, *readbackBufferMemory, 0), IsVkSuccess());
+
+    auto readbackCommandBuffers =
+        device->allocateCommandBuffersUnique(commandBufferAllocateInfo).value;
+    ASSERT_THAT(readbackCommandBuffers, Not(IsEmpty()));
+    auto readbackCommandBuffer = std::move(readbackCommandBuffers[0]);
+    ASSERT_THAT(readbackCommandBuffer, IsValidHandle());
+
+    readbackCommandBuffer->begin(commandBufferBeginInfo);
+    readbackCommandBuffer->copyBuffer(*vertexBuffer, *readbackBuffer, 1, &bufferCopy);
+    readbackCommandBuffer->end();
+
+    auto readbackFence = device->createFenceUnique(vkhpp::FenceCreateInfo()).value;
+    ASSERT_THAT(readbackCommandBuffer, IsValidHandle());
+
+    // Execute the command to copy image back to buffer
+    const vkhpp::SubmitInfo readbackSubmitInfo = {
+        .commandBufferCount = 1,
+        .pCommandBuffers = &readbackCommandBuffer.get(),
+    };
+    queue.submit(readbackSubmitInfo, *readbackFence);
+
+    auto readbackWaitResult = device->waitForFences(*readbackFence, VK_TRUE, 3000000000L);
+    ASSERT_THAT(readbackWaitResult, IsVkSuccess());
+
+    // Verify content
+    mapResult = device->mapMemory(*readbackBufferMemory, 0, VK_WHOLE_SIZE, vkhpp::MemoryMapFlags{},
+                                  &mapped);
+    ASSERT_THAT(mapResult, IsVkSuccess());
+    ASSERT_THAT(mapped, NotNull());
+    bytes = reinterpret_cast<uint8_t*>(mapped);
+
+    for (uint32_t i = 0; i < kSize; ++i) {
+        ASSERT_THAT(bytes[i], Eq(srcBufferContent[i]));
+    }
+    device->unmapMemory(*readbackBufferMemory);
+}
+
+INSTANTIATE_TEST_CASE_P(GfxstreamEnd2EndTests, GfxstreamEnd2EndVkSnapshotBufferTest,
+                        ::testing::ValuesIn({
+                            TestParams{
+                                .with_gl = false,
+                                .with_vk = true,
+                                .with_features = {"VulkanSnapshots"},
+                            },
+                        }),
+                        &GetTestName);
+
+}  // namespace
+}  // namespace tests
+}  // namespace gfxstream
\ No newline at end of file
diff --git a/host/vulkan/VkDecoderGlobalState.cpp b/host/vulkan/VkDecoderGlobalState.cpp
index 8d756f5e..45a90139 100644
--- a/host/vulkan/VkDecoderGlobalState.cpp
+++ b/host/vulkan/VkDecoderGlobalState.cpp
@@ -410,6 +410,56 @@ class VkDecoderGlobalState::Impl {
 
     const gfxstream::host::FeatureSet& getFeatures() const { return m_emu->features; }
 
+    StateBlock createSnapshotStateBlock(VkDevice unboxed_device) {
+            const auto& device = unboxed_device;
+            const auto& deviceInfo = android::base::find(mDeviceInfo, device);
+            const auto physicalDevice = deviceInfo->physicalDevice;
+            const auto& physicalDeviceInfo = android::base::find(mPhysdevInfo, physicalDevice);
+            const auto& instanceInfo = android::base::find(mInstanceInfo, physicalDeviceInfo->instance);
+
+            VulkanDispatch* ivk = dispatch_VkInstance(instanceInfo->boxed);
+            VulkanDispatch* dvk = dispatch_VkDevice(deviceInfo->boxed);
+
+            StateBlock stateBlock{
+                .physicalDevice = physicalDevice,
+                .physicalDeviceInfo = physicalDeviceInfo,
+                .device = device,
+                .deviceDispatch = dvk,
+                .queue = VK_NULL_HANDLE,
+                .commandPool = VK_NULL_HANDLE,
+            };
+
+            uint32_t queueFamilyCount = 0;
+            ivk->vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount,
+                                                          nullptr);
+            std::vector<VkQueueFamilyProperties> queueFamilyProps(queueFamilyCount);
+            ivk->vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount,
+                                                          queueFamilyProps.data());
+            uint32_t queueFamilyIndex = 0;
+            for (auto queue : deviceInfo->queues) {
+                int idx = queue.first;
+                if ((queueFamilyProps[idx].queueFlags & VK_QUEUE_GRAPHICS_BIT) == 0) {
+                    continue;
+                }
+                stateBlock.queue = queue.second[0];
+                queueFamilyIndex = idx;
+                break;
+            }
+
+            VkCommandPoolCreateInfo commandPoolCi = {
+                VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
+                0,
+                VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
+                queueFamilyIndex,
+            };
+            dvk->vkCreateCommandPool(device, &commandPoolCi, nullptr, &stateBlock.commandPool);
+            return stateBlock;
+    }
+
+    void releaseSnapshotStateBlock(const StateBlock* stateBlock) {
+        stateBlock->deviceDispatch->vkDestroyCommandPool(stateBlock->device, stateBlock->commandPool, nullptr);
+    }
+
     void save(android::base::Stream* stream) {
         mSnapshotState = SnapshotState::Saving;
         snapshot()->save(stream);
@@ -430,11 +480,10 @@ class VkDecoderGlobalState::Impl {
             stream->putBe64(it.second.size);
             stream->write(it.second.ptr, it.second.size);
         }
+
         // Set up VK structs to snapshot other Vulkan objects
         // TODO(b/323064243): group all images from the same device and reuse queue / command pool
 
-        VulkanDispatch* ivk = getGlobalVkEmulation()->ivk;
-        VulkanDispatch* dvk = getGlobalVkEmulation()->dvk;
         std::vector<VkImage> sortedBoxedImages;
         for (const auto& imageIte : mImageInfo) {
             sortedBoxedImages.push_back(unboxed_to_boxed_non_dispatchable_VkImage(imageIte.first));
@@ -450,46 +499,34 @@ class VkDecoderGlobalState::Impl {
             }
             // Vulkan command playback doesn't recover image layout. We need to do it here.
             stream->putBe32(imageInfo.layout);
-            const auto& device = imageInfo.device;
-            const auto& deviceInfo = android::base::find(mDeviceInfo, device);
-            const auto physicalDevice = deviceInfo->physicalDevice;
-            const auto& physicalDeviceInfo = android::base::find(mPhysdevInfo, physicalDevice);
-            StateBlock stateBlock{
-                .physicalDevice = physicalDevice,
-                .physicalDeviceInfo = physicalDeviceInfo,
-                .device = device,
-                .queue = VK_NULL_HANDLE,
-                .commandPool = VK_NULL_HANDLE,
-            };
 
-            uint32_t queueFamilyCount = 0;
-            ivk->vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount,
-                                                          nullptr);
-            std::vector<VkQueueFamilyProperties> queueFamilyProps(queueFamilyCount);
-            ivk->vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount,
-                                                          queueFamilyProps.data());
-            uint32_t queueFamilyIndex = 0;
-            for (auto queue : deviceInfo->queues) {
-                int idx = queue.first;
-                if ((queueFamilyProps[idx].queueFlags & VK_QUEUE_GRAPHICS_BIT) == 0) {
-                    continue;
-                }
-                stateBlock.queue = queue.second[0];
-                queueFamilyIndex = idx;
-                break;
+            StateBlock stateBlock = createSnapshotStateBlock(imageInfo.device);
+            // TODO(b/294277842): make sure the queue is empty before using.
+            saveImageContent(stream, &stateBlock, unboxedImage, &imageInfo);
+            releaseSnapshotStateBlock(&stateBlock);
+        }
+
+        // snapshot buffers
+        std::vector<VkBuffer> sortedBoxedBuffers;
+        for (const auto& bufferIte : mBufferInfo) {
+            sortedBoxedBuffers.push_back(
+                unboxed_to_boxed_non_dispatchable_VkBuffer(bufferIte.first));
+        }
+        sort(sortedBoxedBuffers.begin(), sortedBoxedBuffers.end());
+        for (const auto& boxedBuffer : sortedBoxedBuffers) {
+            auto unboxedBuffer = unbox_VkBuffer(boxedBuffer);
+            const BufferInfo& bufferInfo = mBufferInfo[unboxedBuffer];
+            if (bufferInfo.memory == VK_NULL_HANDLE) {
+                continue;
             }
+            // TODO: add a special case for host mapped memory
+            StateBlock stateBlock = createSnapshotStateBlock(bufferInfo.device);
 
-            VkCommandPoolCreateInfo commandPoolCi = {
-                VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
-                0,
-                VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
-                queueFamilyIndex,
-            };
-            dvk->vkCreateCommandPool(device, &commandPoolCi, nullptr, &stateBlock.commandPool);
             // TODO(b/294277842): make sure the queue is empty before using.
-            saveImageContent(stream, &stateBlock, unboxedImage, &imageInfo);
-            dvk->vkDestroyCommandPool(device, stateBlock.commandPool, nullptr);
+            saveBufferContent(stream, &stateBlock, unboxedBuffer, &bufferInfo);
+            releaseSnapshotStateBlock(&stateBlock);
         }
+
         mSnapshotState = SnapshotState::Normal;
     }
 
@@ -523,8 +560,6 @@ class VkDecoderGlobalState::Impl {
         // Set up VK structs to snapshot other Vulkan objects
         // TODO(b/323064243): group all images from the same device and reuse queue / command pool
 
-        VulkanDispatch* ivk = getGlobalVkEmulation()->ivk;
-        VulkanDispatch* dvk = getGlobalVkEmulation()->dvk;
         std::vector<VkImage> sortedBoxedImages;
         for (const auto& imageIte : mImageInfo) {
             sortedBoxedImages.push_back(unboxed_to_boxed_non_dispatchable_VkImage(imageIte.first));
@@ -547,46 +582,32 @@ class VkDecoderGlobalState::Impl {
             //
             // TODO(b/323059453): fix corner cases when image contents cannot be properly loaded.
             imageInfo.layout = static_cast<VkImageLayout>(stream->getBe32());
-            const auto& device = imageInfo.device;
-            const auto& deviceInfo = android::base::find(mDeviceInfo, device);
-            const auto physicalDevice = deviceInfo->physicalDevice;
-            const auto& physicalDeviceInfo = android::base::find(mPhysdevInfo, physicalDevice);
-            StateBlock stateBlock{
-                .physicalDevice = physicalDevice,
-                .physicalDeviceInfo = physicalDeviceInfo,
-                .device = device,
-                .queue = VK_NULL_HANDLE,
-                .commandPool = VK_NULL_HANDLE,
-            };
+            StateBlock stateBlock = createSnapshotStateBlock(imageInfo.device);
+            // TODO(b/294277842): make sure the queue is empty before using.
+            loadImageContent(stream, &stateBlock, unboxedImage, &imageInfo);
+            releaseSnapshotStateBlock(&stateBlock);
+        }
 
-            uint32_t queueFamilyCount = 0;
-            ivk->vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount,
-                                                          nullptr);
-            std::vector<VkQueueFamilyProperties> queueFamilyProps(queueFamilyCount);
-            ivk->vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount,
-                                                          queueFamilyProps.data());
-            uint32_t queueFamilyIndex = 0;
-            for (auto queue : deviceInfo->queues) {
-                int idx = queue.first;
-                if ((queueFamilyProps[idx].queueFlags & VK_QUEUE_GRAPHICS_BIT) == 0) {
-                    continue;
-                }
-                stateBlock.queue = queue.second[0];
-                queueFamilyIndex = idx;
-                break;
+        // snapshot buffers
+        std::vector<VkBuffer> sortedBoxedBuffers;
+        for (const auto& bufferIte : mBufferInfo) {
+            sortedBoxedBuffers.push_back(
+                unboxed_to_boxed_non_dispatchable_VkBuffer(bufferIte.first));
+        }
+        sort(sortedBoxedBuffers.begin(), sortedBoxedBuffers.end());
+        for (const auto& boxedBuffer : sortedBoxedBuffers) {
+            auto unboxedBuffer = unbox_VkBuffer(boxedBuffer);
+            const BufferInfo& bufferInfo = mBufferInfo[unboxedBuffer];
+            if (bufferInfo.memory == VK_NULL_HANDLE) {
+                continue;
             }
-
-            VkCommandPoolCreateInfo commandPoolCi = {
-                VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
-                0,
-                VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
-                queueFamilyIndex,
-            };
-            dvk->vkCreateCommandPool(device, &commandPoolCi, nullptr, &stateBlock.commandPool);
+            // TODO: add a special case for host mapped memory
+            StateBlock stateBlock = createSnapshotStateBlock(bufferInfo.device);
             // TODO(b/294277842): make sure the queue is empty before using.
-            loadImageContent(stream, &stateBlock, unboxedImage, &imageInfo);
-            dvk->vkDestroyCommandPool(device, stateBlock.commandPool, nullptr);
+            loadBufferContent(stream, &stateBlock, unboxedBuffer, &bufferInfo);
+            releaseSnapshotStateBlock(&stateBlock);
         }
+
         mSnapshotState = SnapshotState::Normal;
     }
 
diff --git a/host/vulkan/VkDecoderSnapshot.cpp b/host/vulkan/VkDecoderSnapshot.cpp
index 0260df9f..fcc54493 100644
--- a/host/vulkan/VkDecoderSnapshot.cpp
+++ b/host/vulkan/VkDecoderSnapshot.cpp
@@ -270,7 +270,24 @@ class VkDecoderSnapshot::Impl {
     }
     void vkBindBufferMemory(const uint8_t* snapshotTraceBegin, size_t snapshotTraceBytes,
                             android::base::BumpPool* pool, VkResult input_result, VkDevice device,
-                            VkBuffer buffer, VkDeviceMemory memory, VkDeviceSize memoryOffset) {}
+                            VkBuffer buffer, VkDeviceMemory memory, VkDeviceSize memoryOffset) {
+        VkBuffer boxed_VkBuffer = unboxed_to_boxed_non_dispatchable_VkBuffer((&buffer)[0]);
+        android::base::AutoLock lock(mLock);
+        // buffer create
+        mReconstruction.addHandleDependency(
+            (const uint64_t*)&boxed_VkBuffer, 1,
+            (uint64_t)(uintptr_t)unboxed_to_boxed_non_dispatchable_VkDeviceMemory(memory),
+            VkReconstruction::BOUND_MEMORY);
+        mReconstruction.addHandleDependency((const uint64_t*)&boxed_VkBuffer, 1,
+                                            (uint64_t)(uintptr_t)((&boxed_VkBuffer)[0]),
+                                            VkReconstruction::BOUND_MEMORY);
+        auto apiHandle = mReconstruction.createApiInfo();
+        auto apiInfo = mReconstruction.getApiInfo(apiHandle);
+        mReconstruction.setApiTrace(apiInfo, OP_vkBindBufferMemory, snapshotTraceBegin,
+                                    snapshotTraceBytes);
+        mReconstruction.forEachHandleAddApi((const uint64_t*)&boxed_VkBuffer, 1, apiHandle,
+                                            VkReconstruction::BOUND_MEMORY);
+    }
     void vkBindImageMemory(const uint8_t* snapshotTraceBegin, size_t snapshotTraceBytes,
                            android::base::BumpPool* pool, VkResult input_result, VkDevice device,
                            VkImage image, VkDeviceMemory memory, VkDeviceSize memoryOffset) {
@@ -279,8 +296,7 @@ class VkDecoderSnapshot::Impl {
         // image create
         mReconstruction.addHandleDependency(
             (const uint64_t*)&boxed_VkImage, 1,
-            (uint64_t)(uintptr_t)(uint64_t)(uintptr_t)
-                unboxed_to_boxed_non_dispatchable_VkDeviceMemory(memory),
+            (uint64_t)(uintptr_t)unboxed_to_boxed_non_dispatchable_VkDeviceMemory(memory),
             VkReconstruction::BOUND_MEMORY);
         mReconstruction.addHandleDependency((const uint64_t*)&boxed_VkImage, 1,
                                             (uint64_t)(uintptr_t)((&boxed_VkImage)[0]),
diff --git a/host/vulkan/VkDecoderSnapshotUtils.cpp b/host/vulkan/VkDecoderSnapshotUtils.cpp
index e7fe8c43..921d0753 100644
--- a/host/vulkan/VkDecoderSnapshotUtils.cpp
+++ b/host/vulkan/VkDecoderSnapshotUtils.cpp
@@ -131,8 +131,7 @@ void saveImageContent(android::base::Stream* stream, StateBlock* stateBlock, VkI
     if (imageInfo->imageCreateInfoShallow.samples != VK_SAMPLE_COUNT_1_BIT) {
         return;
     }
-    VkEmulation* vkEmulation = getGlobalVkEmulation();
-    VulkanDispatch* dispatch = vkEmulation->dvk;
+    VulkanDispatch* dispatch = stateBlock->deviceDispatch;
     const VkImageCreateInfo& imageCreateInfo = imageInfo->imageCreateInfoShallow;
     VkCommandBufferAllocateInfo allocInfo{
         .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
@@ -280,8 +279,7 @@ void loadImageContent(android::base::Stream* stream, StateBlock* stateBlock, VkI
     if (imageInfo->layout == VK_IMAGE_LAYOUT_UNDEFINED) {
         return;
     }
-    VkEmulation* vkEmulation = getGlobalVkEmulation();
-    VulkanDispatch* dispatch = vkEmulation->dvk;
+    VulkanDispatch* dispatch = stateBlock->deviceDispatch;
     const VkImageCreateInfo& imageCreateInfo = imageInfo->imageCreateInfoShallow;
     VkCommandBufferAllocateInfo allocInfo{
         .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
@@ -468,5 +466,202 @@ void loadImageContent(android::base::Stream* stream, StateBlock* stateBlock, VkI
     dispatch->vkFreeCommandBuffers(stateBlock->device, stateBlock->commandPool, 1, &commandBuffer);
 }
 
+void saveBufferContent(android::base::Stream* stream, StateBlock* stateBlock, VkBuffer buffer,
+                       const BufferInfo* bufferInfo) {
+    VulkanDispatch* dispatch = stateBlock->deviceDispatch;
+    VkCommandBufferAllocateInfo allocInfo{
+        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
+        .commandPool = stateBlock->commandPool,
+        .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
+        .commandBufferCount = 1,
+    };
+    VkCommandBuffer commandBuffer;
+    _RUN_AND_CHECK(dispatch->vkAllocateCommandBuffers(stateBlock->device, &allocInfo,
+                                                      &commandBuffer) != VK_SUCCESS);
+    VkFenceCreateInfo fenceCreateInfo{
+        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
+    };
+    VkFence fence;
+    _RUN_AND_CHECK(dispatch->vkCreateFence(stateBlock->device, &fenceCreateInfo, nullptr, &fence));
+    VkBufferCreateInfo bufferCreateInfo = {
+        .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
+        .size = static_cast<VkDeviceSize>(bufferInfo->size),
+        .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT,
+        .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
+    };
+    VkBuffer readbackBuffer;
+    _RUN_AND_CHECK(
+        dispatch->vkCreateBuffer(stateBlock->device, &bufferCreateInfo, nullptr, &readbackBuffer));
+
+    VkMemoryRequirements readbackBufferMemoryRequirements{};
+    dispatch->vkGetBufferMemoryRequirements(stateBlock->device, readbackBuffer,
+                                            &readbackBufferMemoryRequirements);
+
+    const auto readbackBufferMemoryType =
+        GetMemoryType(*stateBlock->physicalDeviceInfo, readbackBufferMemoryRequirements,
+                      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
+    // Staging memory
+    // TODO(b/323064243): reuse staging memory
+    VkMemoryAllocateInfo readbackBufferMemoryAllocateInfo = {
+        .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
+        .allocationSize = readbackBufferMemoryRequirements.size,
+        .memoryTypeIndex = readbackBufferMemoryType,
+    };
+    VkDeviceMemory readbackMemory;
+    _RUN_AND_CHECK(dispatch->vkAllocateMemory(stateBlock->device, &readbackBufferMemoryAllocateInfo,
+                                              nullptr, &readbackMemory));
+    _RUN_AND_CHECK(
+        dispatch->vkBindBufferMemory(stateBlock->device, readbackBuffer, readbackMemory, 0));
+
+    void* mapped = nullptr;
+    _RUN_AND_CHECK(dispatch->vkMapMemory(stateBlock->device, readbackMemory, 0, VK_WHOLE_SIZE,
+                                         VkMemoryMapFlags{}, &mapped));
+
+    VkBufferCopy bufferCopy = {
+        .srcOffset = 0,
+        .dstOffset = 0,
+        .size = bufferInfo->size,
+    };
+
+    VkCommandBufferBeginInfo beginInfo{
+        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
+    };
+    if (dispatch->vkBeginCommandBuffer(commandBuffer, &beginInfo) != VK_SUCCESS) {
+        GFXSTREAM_ABORT(emugl::FatalError(emugl::ABORT_REASON_OTHER))
+            << "Failed to start command buffer on snapshot save";
+    }
+    dispatch->vkCmdCopyBuffer(commandBuffer, buffer, readbackBuffer, 1, &bufferCopy);
+    VkBufferMemoryBarrier barrier{.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
+                                  .pNext = nullptr,
+                                  .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
+                                  .dstAccessMask = VK_ACCESS_HOST_READ_BIT,
+                                  .srcQueueFamilyIndex = 0xFFFFFFFF,
+                                  .dstQueueFamilyIndex = 0xFFFFFFFF,
+                                  .buffer = readbackBuffer,
+                                  .offset = 0,
+                                  .size = bufferInfo->size};
+    dispatch->vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT,
+                                   VK_PIPELINE_STAGE_HOST_BIT, 0, 0, nullptr, 1, &barrier, 0,
+                                   nullptr);
+
+    // Execute the command to copy buffer
+    VkSubmitInfo submitInfo = {
+        .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
+        .commandBufferCount = 1,
+        .pCommandBuffers = &commandBuffer,
+    };
+    _RUN_AND_CHECK(dispatch->vkEndCommandBuffer(commandBuffer));
+    _RUN_AND_CHECK(dispatch->vkQueueSubmit(stateBlock->queue, 1, &submitInfo, fence));
+    _RUN_AND_CHECK(dispatch->vkWaitForFences(stateBlock->device, 1, &fence, VK_TRUE, 3000000000L));
+    _RUN_AND_CHECK(dispatch->vkResetFences(stateBlock->device, 1, &fence));
+    stream->putBe64(bufferInfo->size);
+    stream->write(mapped, bufferInfo->size);
+
+    dispatch->vkDestroyFence(stateBlock->device, fence, nullptr);
+    dispatch->vkUnmapMemory(stateBlock->device, readbackMemory);
+    dispatch->vkDestroyBuffer(stateBlock->device, readbackBuffer, nullptr);
+    dispatch->vkFreeMemory(stateBlock->device, readbackMemory, nullptr);
+    dispatch->vkFreeCommandBuffers(stateBlock->device, stateBlock->commandPool, 1, &commandBuffer);
+}
+
+void loadBufferContent(android::base::Stream* stream, StateBlock* stateBlock, VkBuffer buffer,
+                       const BufferInfo* bufferInfo) {
+    VulkanDispatch* dispatch = stateBlock->deviceDispatch;
+    VkCommandBufferAllocateInfo allocInfo{
+        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
+        .commandPool = stateBlock->commandPool,
+        .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
+        .commandBufferCount = 1,
+    };
+    VkCommandBuffer commandBuffer;
+    _RUN_AND_CHECK(dispatch->vkAllocateCommandBuffers(stateBlock->device, &allocInfo,
+                                                      &commandBuffer) != VK_SUCCESS);
+    VkFenceCreateInfo fenceCreateInfo{
+        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
+    };
+    VkFence fence;
+    _RUN_AND_CHECK(dispatch->vkCreateFence(stateBlock->device, &fenceCreateInfo, nullptr, &fence));
+    VkBufferCreateInfo bufferCreateInfo = {
+        .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
+        .size = static_cast<VkDeviceSize>(bufferInfo->size),
+        .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT,
+        .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
+    };
+    VkBuffer stagingBuffer;
+    _RUN_AND_CHECK(
+        dispatch->vkCreateBuffer(stateBlock->device, &bufferCreateInfo, nullptr, &stagingBuffer));
+
+    VkMemoryRequirements stagingBufferMemoryRequirements{};
+    dispatch->vkGetBufferMemoryRequirements(stateBlock->device, stagingBuffer,
+                                            &stagingBufferMemoryRequirements);
+
+    const auto stagingBufferMemoryType =
+        GetMemoryType(*stateBlock->physicalDeviceInfo, stagingBufferMemoryRequirements,
+                      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
+    // Staging memory
+    // TODO(b/323064243): reuse staging memory
+    VkMemoryAllocateInfo stagingBufferMemoryAllocateInfo = {
+        .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
+        .allocationSize = stagingBufferMemoryRequirements.size,
+        .memoryTypeIndex = stagingBufferMemoryType,
+    };
+    VkDeviceMemory stagingMemory;
+    _RUN_AND_CHECK(dispatch->vkAllocateMemory(stateBlock->device, &stagingBufferMemoryAllocateInfo,
+                                              nullptr, &stagingMemory));
+    _RUN_AND_CHECK(
+        dispatch->vkBindBufferMemory(stateBlock->device, stagingBuffer, stagingMemory, 0));
+
+    void* mapped = nullptr;
+    _RUN_AND_CHECK(dispatch->vkMapMemory(stateBlock->device, stagingMemory, 0, VK_WHOLE_SIZE,
+                                         VkMemoryMapFlags{}, &mapped));
+    size_t bufferSize = stream->getBe64();
+    assert(bufferSize == bufferInfo->size);
+    stream->read(mapped, bufferInfo->size);
+
+    VkBufferCopy bufferCopy = {
+        .srcOffset = 0,
+        .dstOffset = 0,
+        .size = bufferInfo->size,
+    };
+
+    VkCommandBufferBeginInfo beginInfo{
+        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
+    };
+    if (dispatch->vkBeginCommandBuffer(commandBuffer, &beginInfo) != VK_SUCCESS) {
+        GFXSTREAM_ABORT(emugl::FatalError(emugl::ABORT_REASON_OTHER))
+            << "Failed to start command buffer on snapshot save";
+    }
+    dispatch->vkCmdCopyBuffer(commandBuffer, stagingBuffer, buffer, 1, &bufferCopy);
+    VkBufferMemoryBarrier barrier{.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
+                                  .pNext = nullptr,
+                                  .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
+                                  .dstAccessMask = static_cast<VkAccessFlags>(~VK_ACCESS_NONE_KHR),
+                                  .srcQueueFamilyIndex = 0xFFFFFFFF,
+                                  .dstQueueFamilyIndex = 0xFFFFFFFF,
+                                  .buffer = buffer,
+                                  .offset = 0,
+                                  .size = bufferInfo->size};
+    dispatch->vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT,
+                                   VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, nullptr, 1, &barrier,
+                                   0, nullptr);
+
+    // Execute the command to copy buffer
+    VkSubmitInfo submitInfo = {
+        .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
+        .commandBufferCount = 1,
+        .pCommandBuffers = &commandBuffer,
+    };
+    _RUN_AND_CHECK(dispatch->vkEndCommandBuffer(commandBuffer));
+    _RUN_AND_CHECK(dispatch->vkQueueSubmit(stateBlock->queue, 1, &submitInfo, fence));
+    _RUN_AND_CHECK(dispatch->vkWaitForFences(stateBlock->device, 1, &fence, VK_TRUE, 3000000000L));
+    _RUN_AND_CHECK(dispatch->vkResetFences(stateBlock->device, 1, &fence));
+
+    dispatch->vkDestroyFence(stateBlock->device, fence, nullptr);
+    dispatch->vkUnmapMemory(stateBlock->device, stagingMemory);
+    dispatch->vkDestroyBuffer(stateBlock->device, stagingBuffer, nullptr);
+    dispatch->vkFreeMemory(stateBlock->device, stagingMemory, nullptr);
+    dispatch->vkFreeCommandBuffers(stateBlock->device, stateBlock->commandPool, 1, &commandBuffer);
+}
+
 }  // namespace vk
 }  // namespace gfxstream
\ No newline at end of file
diff --git a/host/vulkan/VkDecoderSnapshotUtils.h b/host/vulkan/VkDecoderSnapshotUtils.h
index 7a626ecd..c07701e1 100644
--- a/host/vulkan/VkDecoderSnapshotUtils.h
+++ b/host/vulkan/VkDecoderSnapshotUtils.h
@@ -22,6 +22,7 @@ struct StateBlock {
     VkPhysicalDevice physicalDevice;
     const PhysicalDeviceInfo* physicalDeviceInfo;
     VkDevice device;
+    VulkanDispatch* deviceDispatch;
     VkQueue queue;
     VkCommandPool commandPool;
 };
@@ -29,5 +30,9 @@ void saveImageContent(android::base::Stream* stream, StateBlock* stateBlock, VkI
                       const ImageInfo* imageInfo);
 void loadImageContent(android::base::Stream* stream, StateBlock* stateBlock, VkImage image,
                       const ImageInfo* imageInfo);
+void saveBufferContent(android::base::Stream* stream, StateBlock* stateBlock, VkBuffer buffer,
+                       const BufferInfo* bufferInfo);
+void loadBufferContent(android::base::Stream* stream, StateBlock* stateBlock, VkBuffer buffer,
+                       const BufferInfo* bufferInfo);
 }  // namespace vk
 }  // namespace gfxstream
\ No newline at end of file