From 11f8dc0818f836b3e619ec7f07e7ba4728d04a95 Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Sun, 29 Sep 2024 15:21:19 +0200
Subject: [PATCH] [dxvk] Add functions to lock in and query GPU buffer
 addresses

---
 src/dxvk/dxvk_buffer.cpp  |  5 ++---
 src/dxvk/dxvk_buffer.h    | 22 ++++++++++++++++++++++
 src/dxvk/dxvk_context.cpp |  7 +++++++
 src/dxvk/dxvk_context.h   | 12 ++++++++++++
 4 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/src/dxvk/dxvk_buffer.cpp b/src/dxvk/dxvk_buffer.cpp
index 33d53e9c..0c3f291c 100644
--- a/src/dxvk/dxvk_buffer.cpp
+++ b/src/dxvk/dxvk_buffer.cpp
@@ -48,9 +48,8 @@ namespace dxvk {
 
 
   bool DxvkBuffer::canRelocate() const {
-    return !m_storage->flags().test(DxvkAllocationFlag::Imported)
-        && !m_bufferInfo.mapPtr
-        && !(m_info.usage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT)
+    return !m_bufferInfo.mapPtr && !m_stableAddress
+        && !m_storage->flags().test(DxvkAllocationFlag::Imported)
         && !(m_info.flags & VK_BUFFER_CREATE_SPARSE_BINDING_BIT);
   }
 
diff --git a/src/dxvk/dxvk_buffer.h b/src/dxvk/dxvk_buffer.h
index 46fa8390..7ca44b44 100644
--- a/src/dxvk/dxvk_buffer.h
+++ b/src/dxvk/dxvk_buffer.h
@@ -210,6 +210,16 @@ namespace dxvk {
         : nullptr;
     }
 
+    /**
+     * \brief GPU address
+     *
+     * May be 0 if the device address usage flag is not
+     * enabled for this buffer.
+     */
+    VkDeviceAddress gpuAddress() const {
+      return m_bufferInfo.gpuAddress;
+    }
+
     /**
      * \brief Queries shader stages that can access this buffer
      *
@@ -350,6 +360,16 @@ namespace dxvk {
      */
     bool canRelocate() const;
 
+    /**
+     * \brief Enables stable GPU address
+     *
+     * Subsequent calls to \c canRelocate will be \c false, preventing
+     * the buffer from being relocated or invalidated by the backend.
+     */
+    void enableStableAddress() {
+      m_stableAddress = true;
+    }
+
     /**
      * \brief Creates or retrieves a buffer view
      *
@@ -377,6 +397,8 @@ namespace dxvk {
     uint32_t                    m_xfbStride     = 0u;
     uint32_t                    m_version       = 0u;
 
+    bool                        m_stableAddress = false;
+
     DxvkResourceBufferInfo      m_bufferInfo    = { };
 
     Rc<DxvkResourceAllocation>  m_storage;
diff --git a/src/dxvk/dxvk_context.cpp b/src/dxvk/dxvk_context.cpp
index ac54e5b8..a1950c38 100644
--- a/src/dxvk/dxvk_context.cpp
+++ b/src/dxvk/dxvk_context.cpp
@@ -1894,6 +1894,13 @@ namespace dxvk {
   }
 
 
+  void DxvkContext::ensureBufferAddress(
+    const Rc<DxvkBuffer>&           buffer) {
+    // Really nothing else to do here but set the flag
+    buffer->enableStableAddress();
+  }
+
+
   void DxvkContext::invalidateImage(
     const Rc<DxvkImage>&            image,
           Rc<DxvkResourceAllocation>&& slice) {
diff --git a/src/dxvk/dxvk_context.h b/src/dxvk/dxvk_context.h
index 8475a85c..cbba27de 100644
--- a/src/dxvk/dxvk_context.h
+++ b/src/dxvk/dxvk_context.h
@@ -962,6 +962,18 @@ namespace dxvk {
       const Rc<DxvkBuffer>&           buffer,
             Rc<DxvkResourceAllocation>&& slice);
 
+    /**
+     * \brief Ensures that buffer will not be relocated
+     *
+     * This guarantees that the buffer's GPU address remains the same
+     * throughout its lifetime. Only prevents implicit invalidation or
+     * relocation by the backend, client APIs must take care to respect
+     * this too.
+     * \param [in] buffer Buffer to lock in place
+     */
+    void ensureBufferAddress(
+      const Rc<DxvkBuffer>&           buffer);
+
     /**
      * \brief Invalidates image content
      *