From 832f67e187436194fe0815a6e5b5a322b5ff3142 Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Thu, 8 Aug 2024 13:30:03 -0500
Subject: [PATCH] nvk: Implement descriptor buffer binding

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30580>
---
 src/nouveau/vulkan/nvk_cmd_buffer.c | 70 ++++++++++++++++++++++++++++-
 src/nouveau/vulkan/nvk_cmd_buffer.h |  1 +
 src/nouveau/vulkan/nvk_cmd_meta.c   |  8 ++++
 3 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/nouveau/vulkan/nvk_cmd_buffer.c b/src/nouveau/vulkan/nvk_cmd_buffer.c
index 62de9a550f0..84a334372a3 100644
--- a/src/nouveau/vulkan/nvk_cmd_buffer.c
+++ b/src/nouveau/vulkan/nvk_cmd_buffer.c
@@ -775,6 +775,73 @@ nvk_CmdBindDescriptorSets2KHR(VkCommandBuffer commandBuffer,
    }
 }
 
+VKAPI_ATTR void VKAPI_CALL
+nvk_CmdBindDescriptorBuffersEXT(VkCommandBuffer commandBuffer,
+                                uint32_t bufferCount,
+                                const VkDescriptorBufferBindingInfoEXT *pBindingInfos)
+{
+   VK_FROM_HANDLE(nvk_cmd_buffer, cmd, commandBuffer);
+
+   for (uint32_t i = 0; i < bufferCount; i++)
+      cmd->state.descriptor_buffers[i] = pBindingInfos[i].address;
+}
+
+static void
+nvk_set_descriptor_buffer_offsets(struct nvk_cmd_buffer *cmd,
+                                  struct nvk_descriptor_state *desc,
+                                  const VkSetDescriptorBufferOffsetsInfoEXT *info)
+{
+   VK_FROM_HANDLE(vk_pipeline_layout, pipeline_layout, info->layout);
+
+   for (uint32_t i = 0; i < info->setCount; ++i) {
+      const uint32_t s = i + info->firstSet;
+
+      desc->sets[s].type = NVK_DESCRIPTOR_SET_TYPE_BUFFER;
+      desc->sets[s].set = NULL;
+
+      struct nvk_buffer_address set_addr;
+      if (pipeline_layout->set_layouts[s] != NULL) {
+         const struct nvk_descriptor_set_layout *set_layout =
+            vk_to_nvk_descriptor_set_layout(pipeline_layout->set_layouts[s]);
+         assert(set_layout->flags &
+                VK_DESCRIPTOR_SET_LAYOUT_CREATE_DESCRIPTOR_BUFFER_BIT_EXT);
+
+         const uint64_t buffer_base_addr =
+            cmd->state.descriptor_buffers[info->pBufferIndices[i]];
+
+         set_addr = (struct nvk_buffer_address) {
+            .base_addr = buffer_base_addr + info->pOffsets[i],
+            .size = set_layout->max_buffer_size,
+         };
+      } else {
+         set_addr = NVK_BUFFER_ADDRESS_NULL;
+      }
+      nvk_descriptor_state_set_root(cmd, desc, sets[s], set_addr);
+   }
+
+   nvk_cmd_dirty_cbufs_for_descriptors(cmd, info->stageFlags,
+                                       info->firstSet,
+                                       info->firstSet + info->setCount,
+                                       0, 0);
+}
+
+VKAPI_ATTR void VKAPI_CALL
+nvk_CmdSetDescriptorBufferOffsets2EXT(VkCommandBuffer commandBuffer,
+                                      const VkSetDescriptorBufferOffsetsInfoEXT *pInfo)
+{
+   VK_FROM_HANDLE(nvk_cmd_buffer, cmd, commandBuffer);
+
+   if (pInfo->stageFlags & NVK_VK_GRAPHICS_STAGE_BITS) {
+      nvk_set_descriptor_buffer_offsets(cmd, &cmd->state.gfx.descriptors,
+                                        pInfo);
+   }
+
+   if (pInfo->stageFlags & VK_SHADER_STAGE_COMPUTE_BIT) {
+      nvk_set_descriptor_buffer_offsets(cmd, &cmd->state.cs.descriptors,
+                                        pInfo);
+   }
+}
+
 static void
 nvk_push_constants(UNUSED struct nvk_cmd_buffer *cmd,
                    struct nvk_descriptor_state *desc,
@@ -964,7 +1031,8 @@ nvk_cmd_buffer_get_cbuf_descriptor_addr(struct nvk_cmd_buffer *cmd,
 {
    assert(cbuf->type == NVK_CBUF_TYPE_UBO_DESC);
    switch (desc->sets[cbuf->desc_set].type) {
-   case NVK_DESCRIPTOR_SET_TYPE_SET: {
+   case NVK_DESCRIPTOR_SET_TYPE_SET:
+   case NVK_DESCRIPTOR_SET_TYPE_BUFFER: {
       struct nvk_buffer_address set_addr;
       nvk_descriptor_state_get_root(desc, sets[cbuf->desc_set], &set_addr);
 
diff --git a/src/nouveau/vulkan/nvk_cmd_buffer.h b/src/nouveau/vulkan/nvk_cmd_buffer.h
index 9b6f5085f1d..14f72e2b8c9 100644
--- a/src/nouveau/vulkan/nvk_cmd_buffer.h
+++ b/src/nouveau/vulkan/nvk_cmd_buffer.h
@@ -73,6 +73,7 @@ enum ENUM_PACKED nvk_descriptor_set_type {
    NVK_DESCRIPTOR_SET_TYPE_NONE,
    NVK_DESCRIPTOR_SET_TYPE_SET,
    NVK_DESCRIPTOR_SET_TYPE_PUSH,
+   NVK_DESCRIPTOR_SET_TYPE_BUFFER,
 };
 
 struct nvk_descriptor_set_binding {
diff --git a/src/nouveau/vulkan/nvk_cmd_meta.c b/src/nouveau/vulkan/nvk_cmd_meta.c
index 1f018d3ba8a..3a9e8f310fb 100644
--- a/src/nouveau/vulkan/nvk_cmd_meta.c
+++ b/src/nouveau/vulkan/nvk_cmd_meta.c
@@ -64,6 +64,7 @@ struct nvk_meta_save {
    struct nvk_shader *shaders[MESA_SHADER_MESH + 1];
    struct nvk_addr_range vb0;
    struct nvk_descriptor_set_binding desc0;
+   struct nvk_buffer_address desc0_set_addr;
    struct nvk_push_descriptor_set push_desc0;
    uint8_t set_dynamic_buffer_start[NVK_MAX_SETS];
    uint8_t push[NVK_MAX_PUSH_SIZE];
@@ -85,6 +86,7 @@ nvk_meta_begin(struct nvk_cmd_buffer *cmd,
    save->vb0 = cmd->state.gfx.vb0;
 
    save->desc0 = desc->sets[0];
+   nvk_descriptor_state_get_root(desc, sets[0], &save->desc0_set_addr);
    if (desc->sets[0].push != NULL)
       save->push_desc0 = *desc->sets[0].push;
 
@@ -156,6 +158,12 @@ nvk_meta_end(struct nvk_cmd_buffer *cmd,
       desc->push_dirty |= BITFIELD_BIT(0);
       break;
 
+   case NVK_DESCRIPTOR_SET_TYPE_BUFFER:
+      desc->sets[0].type = NVK_DESCRIPTOR_SET_TYPE_BUFFER;
+      desc->sets[0].set = NULL;
+      nvk_descriptor_state_set_root(cmd, desc, sets[0], save->desc0_set_addr);
+      break;
+
    default:
       unreachable("Unknown descriptor set type");
    }