From 50878f2846cf4ea94eb2ab91e319e3fcb56fa60f Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Thu, 26 Sep 2024 03:21:46 +0200
Subject: [PATCH] [dxvk] Add function to invalidate images

---
 src/dxvk/dxvk_context.cpp | 26 +++++++++++++++++++++++++-
 src/dxvk/dxvk_context.h   | 13 +++++++++++++
 src/dxvk/dxvk_image.cpp   | 16 ++++++++++++----
 src/dxvk/dxvk_image.h     | 19 +++++++++++++++++++
 4 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/src/dxvk/dxvk_context.cpp b/src/dxvk/dxvk_context.cpp
index fc518606b..681d18b52 100644
--- a/src/dxvk/dxvk_context.cpp
+++ b/src/dxvk/dxvk_context.cpp
@@ -1815,7 +1815,6 @@ namespace dxvk {
   void DxvkContext::invalidateBuffer(
     const Rc<DxvkBuffer>&           buffer,
           Rc<DxvkResourceAllocation>&& slice) {
-    // Allocate new backing resource
     Rc<DxvkResourceAllocation> prevAllocation = buffer->assignSlice(std::move(slice));
     m_cmd->trackResource<DxvkAccess::None>(std::move(prevAllocation));
 
@@ -1851,6 +1850,31 @@ namespace dxvk {
   }
 
 
+  void DxvkContext::invalidateImage(
+    const Rc<DxvkImage>&            image,
+          Rc<DxvkResourceAllocation>&& slice) {
+    Rc<DxvkResourceAllocation> prevAllocation = image->assignResource(std::move(slice));
+    m_cmd->trackResource<DxvkAccess::None>(std::move(prevAllocation));
+
+    VkImageUsageFlags usage = image->info().usage;
+
+    if (usage & (VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT))
+      m_descriptorState.dirtyViews(image->getShaderStages());
+
+    if (usage & (VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT)) {
+      // Interrupt the current render pass if the image is bound for rendering
+      for (uint32_t i = 0; i < m_state.om.framebufferInfo.numAttachments(); i++) {
+        if (m_state.om.framebufferInfo.getAttachment(i).view->image() == image) {
+          this->spillRenderPass(false);
+
+          m_flags.set(DxvkContextFlag::GpDirtyFramebuffer);
+          break;
+        }
+      }
+    }
+  }
+
+
   void DxvkContext::resolveImage(
     const Rc<DxvkImage>&            dstImage,
     const Rc<DxvkImage>&            srcImage,
diff --git a/src/dxvk/dxvk_context.h b/src/dxvk/dxvk_context.h
index 2d8513514..55788fcbb 100644
--- a/src/dxvk/dxvk_context.h
+++ b/src/dxvk/dxvk_context.h
@@ -964,6 +964,19 @@ namespace dxvk {
     void invalidateBuffer(
       const Rc<DxvkBuffer>&           buffer,
             Rc<DxvkResourceAllocation>&& slice);
+
+    /**
+     * \brief Invalidates image content
+     *
+     * Replaces the backing storage of an image.
+     * \warning If the image is used by another context,
+     * invalidating it will result in undefined behaviour.
+     * \param [in] buffer The buffer to invalidate
+     * \param [in] slice New buffer slice
+     */
+    void invalidateImage(
+      const Rc<DxvkImage>&            image,
+            Rc<DxvkResourceAllocation>&& slice);
     
     /**
      * \brief Updates push constants
diff --git a/src/dxvk/dxvk_image.cpp b/src/dxvk/dxvk_image.cpp
index 770cdfdb8..de8e47ad6 100644
--- a/src/dxvk/dxvk_image.cpp
+++ b/src/dxvk/dxvk_image.cpp
@@ -9,7 +9,11 @@ namespace dxvk {
     const DxvkImageCreateInfo&  createInfo,
           DxvkMemoryAllocator&  memAlloc,
           VkMemoryPropertyFlags memFlags)
-  : m_vkd(device->vkd()), m_allocator(&memAlloc), m_properties(memFlags), m_info(createInfo) {
+  : m_vkd           (device->vkd()),
+    m_allocator     (&memAlloc),
+    m_properties    (memFlags),
+    m_shaderStages  (util::shaderStages(createInfo.stages)),
+    m_info          (createInfo) {
     copyFormatList(createInfo.viewFormatCount, createInfo.viewFormats);
 
     // Determine whether the image is shareable before creating the resource
@@ -22,12 +26,16 @@ namespace dxvk {
 
   DxvkImage::DxvkImage(
           DxvkDevice*           device,
-    const DxvkImageCreateInfo&  info,
+    const DxvkImageCreateInfo&  createInfo,
           VkImage               imageHandle,
           DxvkMemoryAllocator&  memAlloc,
           VkMemoryPropertyFlags memFlags)
-  : m_vkd(device->vkd()), m_allocator(&memAlloc), m_properties(memFlags), m_info(info) {
-    copyFormatList(info.viewFormatCount, info.viewFormats);
+  : m_vkd           (device->vkd()),
+    m_allocator     (&memAlloc),
+    m_properties    (memFlags),
+    m_shaderStages  (util::shaderStages(createInfo.stages)),
+    m_info          (createInfo) {
+    copyFormatList(createInfo.viewFormatCount, createInfo.viewFormats);
 
     // Create backing storage for existing image resource
     VkImageCreateInfo imageInfo = getImageCreateInfo();
diff --git a/src/dxvk/dxvk_image.h b/src/dxvk/dxvk_image.h
index 5f90d20fc..5000fbe94 100644
--- a/src/dxvk/dxvk_image.h
+++ b/src/dxvk/dxvk_image.h
@@ -373,6 +373,16 @@ namespace dxvk {
       return m_properties;
     }
     
+    /**
+     * \brief Queries shader stages that can access this image
+     *
+     * Derived from the pipeline stage mask passed in during creation.
+     * \returns Shader stages that may access this image
+     */
+    VkShaderStageFlags getShaderStages() const {
+      return m_shaderStages;
+    }
+
     /**
      * \brief Map pointer
      * 
@@ -543,6 +553,14 @@ namespace dxvk {
       return old;
     }
 
+    /**
+     * \brief Retrieves current backing storage
+     * \returns Backing storage for this image
+     */
+    Rc<DxvkResourceAllocation> getAllocation() const {
+      return m_storage;
+    }
+
     /**
      * \brief Creates or retrieves an image view
      *
@@ -557,6 +575,7 @@ namespace dxvk {
     Rc<vk::DeviceFn>            m_vkd;
     DxvkMemoryAllocator*        m_allocator   = nullptr;
     VkMemoryPropertyFlags       m_properties  = 0u;
+    VkShaderStageFlags          m_shaderStages = 0u;
 
     DxvkImageCreateInfo         m_info        = { };