7 files changed, 245 insertions, 73 deletions

diff --git a/src/video_core/renderer_base.cpp b/src/video_core/renderer_base.cpp
index 3ea72fda9..a99c33c37 100644
--- a/src/video_core/renderer_base.cpp
+++ b/src/video_core/renderer_base.cpp
@@ -27,7 +27,7 @@ void RendererBase::UpdateCurrentFramebufferLayout() {
     render_window.UpdateCurrentFramebufferLayout(layout.width, layout.height);
 }
 
-void RendererBase::RequestScreenshot(void* data, std::function<void()> callback,
+void RendererBase::RequestScreenshot(void* data, std::function<void(bool)> callback,
                                      const Layout::FramebufferLayout& layout) {
     if (renderer_settings.screenshot_requested) {
         LOG_ERROR(Render, "A screenshot is already requested or in progress, ignoring the request");
diff --git a/src/video_core/renderer_base.h b/src/video_core/renderer_base.h
index 22b80c328..bb204454e 100644
--- a/src/video_core/renderer_base.h
+++ b/src/video_core/renderer_base.h
@@ -24,7 +24,7 @@ struct RendererSettings {
     // Screenshot
     std::atomic<bool> screenshot_requested{false};
     void* screenshot_bits{};
-    std::function<void()> screenshot_complete_callback;
+    std::function<void(bool)> screenshot_complete_callback;
     Layout::FramebufferLayout screenshot_framebuffer_layout;
 };
 
@@ -80,7 +80,7 @@ public:
     void RefreshBaseSettings();
 
     /// Request a screenshot of the next frame
-    void RequestScreenshot(void* data, std::function<void()> callback,
+    void RequestScreenshot(void* data, std::function<void(bool)> callback,
                            const Layout::FramebufferLayout& layout);
 
 protected:
diff --git a/src/video_core/renderer_opengl/renderer_opengl.cpp b/src/video_core/renderer_opengl/renderer_opengl.cpp
index f1b00c24c..7d7cba69c 100644
--- a/src/video_core/renderer_opengl/renderer_opengl.cpp
+++ b/src/video_core/renderer_opengl/renderer_opengl.cpp
@@ -486,7 +486,7 @@ void RendererOpenGL::RenderScreenshot() {
     glBindFramebuffer(GL_READ_FRAMEBUFFER, old_read_fb);
     glBindFramebuffer(GL_DRAW_FRAMEBUFFER, old_draw_fb);
 
-    renderer_settings.screenshot_complete_callback();
+    renderer_settings.screenshot_complete_callback(true);
     renderer_settings.screenshot_requested = false;
 }
 
diff --git a/src/video_core/renderer_vulkan/renderer_vulkan.cpp b/src/video_core/renderer_vulkan/renderer_vulkan.cpp
index a8d04dc61..7c9b0d6db 100644
--- a/src/video_core/renderer_vulkan/renderer_vulkan.cpp
+++ b/src/video_core/renderer_vulkan/renderer_vulkan.cpp
@@ -138,6 +138,7 @@ void RendererVulkan::SwapBuffers(const Tegra::FramebufferConfig* framebuffer) {
     const bool use_accelerated =
         rasterizer.AccelerateDisplay(*framebuffer, framebuffer_addr, framebuffer->stride);
     const bool is_srgb = use_accelerated && screen_info.is_srgb;
+    RenderScreenshot(*framebuffer, use_accelerated);
 
     bool has_been_recreated = false;
     const auto recreate_swapchain = [&] {
@@ -162,7 +163,7 @@ void RendererVulkan::SwapBuffers(const Tegra::FramebufferConfig* framebuffer) {
     if (has_been_recreated) {
         blit_screen.Recreate();
     }
-    const VkSemaphore render_semaphore = blit_screen.Draw(*framebuffer, use_accelerated);
+    const VkSemaphore render_semaphore = blit_screen.DrawToSwapchain(*framebuffer, use_accelerated);
     scheduler.Flush(render_semaphore);
     scheduler.WaitWorker();
     swapchain.Present(render_semaphore);
@@ -193,4 +194,153 @@ void RendererVulkan::Report() const {
     telemetry_session.AddField(field, "GPU_Vulkan_Extensions", extensions);
 }
 
+void Vulkan::RendererVulkan::RenderScreenshot(const Tegra::FramebufferConfig& framebuffer,
+                                              bool use_accelerated) {
+    if (!renderer_settings.screenshot_requested) {
+        return;
+    }
+    const Layout::FramebufferLayout layout{renderer_settings.screenshot_framebuffer_layout};
+    vk::Image staging_image = device.GetLogical().CreateImage(VkImageCreateInfo{
+        .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
+        .pNext = nullptr,
+        .flags = VK_IMAGE_CREATE_MUTABLE_FORMAT_BIT,
+        .imageType = VK_IMAGE_TYPE_2D,
+        .format = VK_FORMAT_B8G8R8A8_UNORM,
+        .extent =
+            {
+                .width = layout.width,
+                .height = layout.height,
+                .depth = 1,
+            },
+        .mipLevels = 1,
+        .arrayLayers = 1,
+        .samples = VK_SAMPLE_COUNT_1_BIT,
+        .tiling = VK_IMAGE_TILING_OPTIMAL,
+        .usage = VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT |
+                 VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
+        .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
+        .queueFamilyIndexCount = 0,
+        .pQueueFamilyIndices = nullptr,
+        .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
+    });
+    const auto image_commit = memory_allocator.Commit(staging_image, MemoryUsage::DeviceLocal);
+
+    const vk::ImageView dst_view = device.GetLogical().CreateImageView(VkImageViewCreateInfo{
+        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
+        .pNext = nullptr,
+        .flags = 0,
+        .image = *staging_image,
+        .viewType = VK_IMAGE_VIEW_TYPE_2D,
+        .format = screen_info.is_srgb ? VK_FORMAT_B8G8R8A8_SRGB : VK_FORMAT_B8G8R8A8_UNORM,
+        .components{
+            .r = VK_COMPONENT_SWIZZLE_IDENTITY,
+            .g = VK_COMPONENT_SWIZZLE_IDENTITY,
+            .b = VK_COMPONENT_SWIZZLE_IDENTITY,
+            .a = VK_COMPONENT_SWIZZLE_IDENTITY,
+        },
+        .subresourceRange{
+            .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+            .baseMipLevel = 0,
+            .levelCount = 1,
+            .baseArrayLayer = 0,
+            .layerCount = VK_REMAINING_ARRAY_LAYERS,
+        },
+    });
+    const VkExtent2D render_area{.width = layout.width, .height = layout.height};
+    const vk::Framebuffer screenshot_fb = blit_screen.CreateFramebuffer(*dst_view, render_area);
+    // Since we're not rendering to the screen, ignore the render semaphore.
+    void(blit_screen.Draw(framebuffer, *screenshot_fb, layout, render_area, use_accelerated));
+
+    const auto buffer_size = static_cast<VkDeviceSize>(layout.width * layout.height * 4);
+    const VkBufferCreateInfo dst_buffer_info{
+        .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
+        .pNext = nullptr,
+        .flags = 0,
+        .size = buffer_size,
+        .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT,
+        .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
+        .queueFamilyIndexCount = 0,
+        .pQueueFamilyIndices = nullptr,
+    };
+    const vk::Buffer dst_buffer = device.GetLogical().CreateBuffer(dst_buffer_info);
+    MemoryCommit dst_buffer_memory = memory_allocator.Commit(dst_buffer, MemoryUsage::Download);
+
+    scheduler.RequestOutsideRenderPassOperationContext();
+    scheduler.Record([&](vk::CommandBuffer cmdbuf) {
+        const VkImageMemoryBarrier read_barrier{
+            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
+            .pNext = nullptr,
+            .srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT,
+            .dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT,
+            .oldLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
+            .newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
+            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+            .image = *staging_image,
+            .subresourceRange{
+                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+                .baseMipLevel = 0,
+                .levelCount = VK_REMAINING_MIP_LEVELS,
+                .baseArrayLayer = 0,
+                .layerCount = VK_REMAINING_ARRAY_LAYERS,
+            },
+        };
+        const VkImageMemoryBarrier image_write_barrier{
+            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
+            .pNext = nullptr,
+            .srcAccessMask = 0,
+            .dstAccessMask = VK_ACCESS_MEMORY_WRITE_BIT,
+            .oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
+            .newLayout = VK_IMAGE_LAYOUT_GENERAL,
+            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+            .image = *staging_image,
+            .subresourceRange{
+                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+                .baseMipLevel = 0,
+                .levelCount = VK_REMAINING_MIP_LEVELS,
+                .baseArrayLayer = 0,
+                .layerCount = VK_REMAINING_ARRAY_LAYERS,
+            },
+        };
+        static constexpr VkMemoryBarrier memory_write_barrier{
+            .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
+            .pNext = nullptr,
+            .srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT,
+            .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT,
+        };
+        const VkBufferImageCopy copy{
+            .bufferOffset = 0,
+            .bufferRowLength = 0,
+            .bufferImageHeight = 0,
+            .imageSubresource{
+                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+                .mipLevel = 0,
+                .baseArrayLayer = 0,
+                .layerCount = 1,
+            },
+            .imageOffset{.x = 0, .y = 0, .z = 0},
+            .imageExtent{
+                .width = layout.width,
+                .height = layout.height,
+                .depth = 1,
+            },
+        };
+        cmdbuf.PipelineBarrier(VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT,
+                               0, read_barrier);
+        cmdbuf.CopyImageToBuffer(*staging_image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, *dst_buffer,
+                                 copy);
+        cmdbuf.PipelineBarrier(VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
+                               0, memory_write_barrier, nullptr, image_write_barrier);
+    });
+    // Ensure the copy is fully completed before saving the screenshot
+    scheduler.Finish();
+
+    // Copy backing image data to the QImage screenshot buffer
+    const auto dst_memory_map = dst_buffer_memory.Map();
+    std::memcpy(renderer_settings.screenshot_bits, dst_memory_map.data(), dst_memory_map.size());
+    renderer_settings.screenshot_complete_callback(false);
+    renderer_settings.screenshot_requested = false;
+}
+
 } // namespace Vulkan
diff --git a/src/video_core/renderer_vulkan/renderer_vulkan.h b/src/video_core/renderer_vulkan/renderer_vulkan.h
index d7d17e110..6dc985109 100644
--- a/src/video_core/renderer_vulkan/renderer_vulkan.h
+++ b/src/video_core/renderer_vulkan/renderer_vulkan.h
@@ -54,6 +54,8 @@ public:
 private:
     void Report() const;
 
+    void RenderScreenshot(const Tegra::FramebufferConfig& framebuffer, bool use_accelerated);
+
     Core::TelemetrySession& telemetry_session;
     Core::Memory::Memory& cpu_memory;
     Tegra::GPU& gpu;
diff --git a/src/video_core/renderer_vulkan/vk_blit_screen.cpp b/src/video_core/renderer_vulkan/vk_blit_screen.cpp
index 516f428e7..5c43b8acf 100644
--- a/src/video_core/renderer_vulkan/vk_blit_screen.cpp
+++ b/src/video_core/renderer_vulkan/vk_blit_screen.cpp
@@ -130,7 +130,10 @@ void VKBlitScreen::Recreate() {
     CreateDynamicResources();
 }
 
-VkSemaphore VKBlitScreen::Draw(const Tegra::FramebufferConfig& framebuffer, bool use_accelerated) {
+VkSemaphore VKBlitScreen::Draw(const Tegra::FramebufferConfig& framebuffer,
+                               const VkFramebuffer& host_framebuffer,
+                               const Layout::FramebufferLayout layout, VkExtent2D render_area,
+                               bool use_accelerated) {
     RefreshResources(framebuffer);
 
     // Finish any pending renderpass
@@ -145,8 +148,8 @@ VkSemaphore VKBlitScreen::Draw(const Tegra::FramebufferConfig& framebuffer, bool
                         use_accelerated ? screen_info.image_view : *raw_image_views[image_index]);
 
     BufferData data;
-    SetUniformData(data, framebuffer);
-    SetVertexData(data, framebuffer);
+    SetUniformData(data, layout);
+    SetVertexData(data, framebuffer, layout);
 
     const std::span<u8> mapped_span = buffer_commit.Map();
     std::memcpy(mapped_span.data(), &data, sizeof(data));
@@ -220,52 +223,75 @@ VkSemaphore VKBlitScreen::Draw(const Tegra::FramebufferConfig& framebuffer, bool
                                    VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, write_barrier);
         });
     }
-    scheduler.Record([this, image_index, size = swapchain.GetSize()](vk::CommandBuffer cmdbuf) {
-        const f32 bg_red = Settings::values.bg_red.GetValue() / 255.0f;
-        const f32 bg_green = Settings::values.bg_green.GetValue() / 255.0f;
-        const f32 bg_blue = Settings::values.bg_blue.GetValue() / 255.0f;
-        const VkClearValue clear_color{
-            .color = {.float32 = {bg_red, bg_green, bg_blue, 1.0f}},
-        };
-        const VkRenderPassBeginInfo renderpass_bi{
-            .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
-            .pNext = nullptr,
-            .renderPass = *renderpass,
-            .framebuffer = *framebuffers[image_index],
-            .renderArea =
-                {
-                    .offset = {0, 0},
-                    .extent = size,
-                },
-            .clearValueCount = 1,
-            .pClearValues = &clear_color,
-        };
-        const VkViewport viewport{
-            .x = 0.0f,
-            .y = 0.0f,
-            .width = static_cast<float>(size.width),
-            .height = static_cast<float>(size.height),
-            .minDepth = 0.0f,
-            .maxDepth = 1.0f,
-        };
-        const VkRect2D scissor{
-            .offset = {0, 0},
-            .extent = size,
-        };
-        cmdbuf.BeginRenderPass(renderpass_bi, VK_SUBPASS_CONTENTS_INLINE);
-        cmdbuf.BindPipeline(VK_PIPELINE_BIND_POINT_GRAPHICS, *pipeline);
-        cmdbuf.SetViewport(0, viewport);
-        cmdbuf.SetScissor(0, scissor);
-
-        cmdbuf.BindVertexBuffer(0, *buffer, offsetof(BufferData, vertices));
-        cmdbuf.BindDescriptorSets(VK_PIPELINE_BIND_POINT_GRAPHICS, *pipeline_layout, 0,
-                                  descriptor_sets[image_index], {});
-        cmdbuf.Draw(4, 1, 0, 0);
-        cmdbuf.EndRenderPass();
-    });
+    scheduler.Record(
+        [this, host_framebuffer, image_index, size = render_area](vk::CommandBuffer cmdbuf) {
+            const f32 bg_red = Settings::values.bg_red.GetValue() / 255.0f;
+            const f32 bg_green = Settings::values.bg_green.GetValue() / 255.0f;
+            const f32 bg_blue = Settings::values.bg_blue.GetValue() / 255.0f;
+            const VkClearValue clear_color{
+                .color = {.float32 = {bg_red, bg_green, bg_blue, 1.0f}},
+            };
+            const VkRenderPassBeginInfo renderpass_bi{
+                .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
+                .pNext = nullptr,
+                .renderPass = *renderpass,
+                .framebuffer = host_framebuffer,
+                .renderArea =
+                    {
+                        .offset = {0, 0},
+                        .extent = size,
+                    },
+                .clearValueCount = 1,
+                .pClearValues = &clear_color,
+            };
+            const VkViewport viewport{
+                .x = 0.0f,
+                .y = 0.0f,
+                .width = static_cast<float>(size.width),
+                .height = static_cast<float>(size.height),
+                .minDepth = 0.0f,
+                .maxDepth = 1.0f,
+            };
+            const VkRect2D scissor{
+                .offset = {0, 0},
+                .extent = size,
+            };
+            cmdbuf.BeginRenderPass(renderpass_bi, VK_SUBPASS_CONTENTS_INLINE);
+            cmdbuf.BindPipeline(VK_PIPELINE_BIND_POINT_GRAPHICS, *pipeline);
+            cmdbuf.SetViewport(0, viewport);
+            cmdbuf.SetScissor(0, scissor);
+
+            cmdbuf.BindVertexBuffer(0, *buffer, offsetof(BufferData, vertices));
+            cmdbuf.BindDescriptorSets(VK_PIPELINE_BIND_POINT_GRAPHICS, *pipeline_layout, 0,
+                                      descriptor_sets[image_index], {});
+            cmdbuf.Draw(4, 1, 0, 0);
+            cmdbuf.EndRenderPass();
+        });
     return *semaphores[image_index];
 }
 
+VkSemaphore VKBlitScreen::DrawToSwapchain(const Tegra::FramebufferConfig& framebuffer,
+                                          bool use_accelerated) {
+    const std::size_t image_index = swapchain.GetImageIndex();
+    const VkExtent2D render_area = swapchain.GetSize();
+    const Layout::FramebufferLayout layout = render_window.GetFramebufferLayout();
+    return Draw(framebuffer, *framebuffers[image_index], layout, render_area, use_accelerated);
+}
+
+vk::Framebuffer VKBlitScreen::CreateFramebuffer(const VkImageView& image_view, VkExtent2D extent) {
+    return device.GetLogical().CreateFramebuffer(VkFramebufferCreateInfo{
+        .sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
+        .pNext = nullptr,
+        .flags = 0,
+        .renderPass = *renderpass,
+        .attachmentCount = 1,
+        .pAttachments = &image_view,
+        .width = extent.width,
+        .height = extent.height,
+        .layers = 1,
+    });
+}
+
 void VKBlitScreen::CreateStaticResources() {
     CreateShaders();
     CreateSemaphores();
@@ -609,22 +635,9 @@ void VKBlitScreen::CreateFramebuffers() {
     const VkExtent2D size{swapchain.GetSize()};
     framebuffers.resize(image_count);
 
-    VkFramebufferCreateInfo ci{
-        .sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
-        .pNext = nullptr,
-        .flags = 0,
-        .renderPass = *renderpass,
-        .attachmentCount = 1,
-        .pAttachments = nullptr,
-        .width = size.width,
-        .height = size.height,
-        .layers = 1,
-    };
-
     for (std::size_t i = 0; i < image_count; ++i) {
         const VkImageView image_view{swapchain.GetImageViewIndex(i)};
-        ci.pAttachments = &image_view;
-        framebuffers[i] = device.GetLogical().CreateFramebuffer(ci);
+        framebuffers[i] = CreateFramebuffer(image_view, size);
     }
 }
 
@@ -752,15 +765,13 @@ void VKBlitScreen::UpdateDescriptorSet(std::size_t image_index, VkImageView imag
     device.GetLogical().UpdateDescriptorSets(std::array{ubo_write, sampler_write}, {});
 }
 
-void VKBlitScreen::SetUniformData(BufferData& data,
-                                  const Tegra::FramebufferConfig& framebuffer) const {
-    const auto& layout = render_window.GetFramebufferLayout();
+void VKBlitScreen::SetUniformData(BufferData& data, const Layout::FramebufferLayout layout) const {
     data.uniform.modelview_matrix =
         MakeOrthographicMatrix(static_cast<f32>(layout.width), static_cast<f32>(layout.height));
 }
 
-void VKBlitScreen::SetVertexData(BufferData& data,
-                                 const Tegra::FramebufferConfig& framebuffer) const {
+void VKBlitScreen::SetVertexData(BufferData& data, const Tegra::FramebufferConfig& framebuffer,
+                                 const Layout::FramebufferLayout layout) const {
     const auto& framebuffer_transform_flags = framebuffer.transform_flags;
     const auto& framebuffer_crop_rect = framebuffer.crop_rect;
 
@@ -798,7 +809,7 @@ void VKBlitScreen::SetVertexData(BufferData& data,
                   static_cast<f32>(screen_info.height);
     }
 
-    const auto& screen = render_window.GetFramebufferLayout().screen;
+    const auto& screen = layout.screen;
     const auto x = static_cast<f32>(screen.left);
     const auto y = static_cast<f32>(screen.top);
     const auto w = static_cast<f32>(screen.GetWidth());
diff --git a/src/video_core/renderer_vulkan/vk_blit_screen.h b/src/video_core/renderer_vulkan/vk_blit_screen.h
index 5e3177685..430bcfbca 100644
--- a/src/video_core/renderer_vulkan/vk_blit_screen.h
+++ b/src/video_core/renderer_vulkan/vk_blit_screen.h
@@ -56,8 +56,16 @@ public:
     void Recreate();
 
     [[nodiscard]] VkSemaphore Draw(const Tegra::FramebufferConfig& framebuffer,
+                                   const VkFramebuffer& host_framebuffer,
+                                   const Layout::FramebufferLayout layout, VkExtent2D render_area,
                                    bool use_accelerated);
 
+    [[nodiscard]] VkSemaphore DrawToSwapchain(const Tegra::FramebufferConfig& framebuffer,
+                                              bool use_accelerated);
+
+    [[nodiscard]] vk::Framebuffer CreateFramebuffer(const VkImageView& image_view,
+                                                    VkExtent2D extent);
+
 private:
     struct BufferData;
 
@@ -81,8 +89,9 @@ private:
     void CreateRawImages(const Tegra::FramebufferConfig& framebuffer);
 
     void UpdateDescriptorSet(std::size_t image_index, VkImageView image_view) const;
-    void SetUniformData(BufferData& data, const Tegra::FramebufferConfig& framebuffer) const;
-    void SetVertexData(BufferData& data, const Tegra::FramebufferConfig& framebuffer) const;
+    void SetUniformData(BufferData& data, const Layout::FramebufferLayout layout) const;
+    void SetVertexData(BufferData& data, const Tegra::FramebufferConfig& framebuffer,
+                       const Layout::FramebufferLayout layout) const;
 
     u64 CalculateBufferSize(const Tegra::FramebufferConfig& framebuffer) const;
     u64 GetRawImageOffset(const Tegra::FramebufferConfig& framebuffer,