venus: implement vkCmdPushDescriptorSetWithTemplateKHR

Signed-off-by: Dawn Han <dawnhan@google.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18988>

src/virtio/vulkan/vn_command_buffer.c

@@ -13,6 +13,7 @@
 #include "venus-protocol/vn_protocol_driver_command_buffer.h"
 #include "venus-protocol/vn_protocol_driver_command_pool.h"
 
+#include "vn_descriptor_set.h"
 #include "vn_device.h"
 #include "vn_image.h"
 #include "vn_render_pass.h"
@@ -2031,3 +2032,38 @@ vn_CmdDrawMultiIndexedEXT(VkCommandBuffer commandBuffer,
    vn_cmd_count_draw_and_submit_on_batch_limit(
       vn_command_buffer_from_handle(commandBuffer));
 }
+
+void
+vn_CmdPushDescriptorSetKHR(VkCommandBuffer commandBuffer,
+                           VkPipelineBindPoint pipelineBindPoint,
+                           VkPipelineLayout layout,
+                           uint32_t set,
+                           uint32_t descriptorWriteCount,
+                           const VkWriteDescriptorSet *pDescriptorWrites)
+{
+   VN_CMD_ENQUEUE(vkCmdPushDescriptorSetKHR, commandBuffer, pipelineBindPoint,
+                  layout, set, descriptorWriteCount, pDescriptorWrites);
+}
+
+void
+vn_CmdPushDescriptorSetWithTemplateKHR(
+   VkCommandBuffer commandBuffer,
+   VkDescriptorUpdateTemplate descriptorUpdateTemplate,
+   VkPipelineLayout layout,
+   uint32_t set,
+   const void *pData)
+{
+   struct vn_descriptor_update_template *templ =
+      vn_descriptor_update_template_from_handle(descriptorUpdateTemplate);
+
+   mtx_lock(&templ->mutex);
+
+   struct vn_update_descriptor_sets *update =
+      vn_update_descriptor_set_with_template_locked(templ, VK_NULL_HANDLE,
+                                                    pData);
+   VN_CMD_ENQUEUE(vkCmdPushDescriptorSetKHR, commandBuffer,
+                  templ->pipeline_bind_point, layout, set,
+                  update->write_count, update->writes);
+
+   mtx_unlock(&templ->mutex);
+}

src/virtio/vulkan/vn_descriptor_set.c

@@ -16,6 +16,7 @@
 #include "venus-protocol/vn_protocol_driver_descriptor_update_template.h"
 
 #include "vn_device.h"
+#include "vn_pipeline.h"
 
 void
 vn_descriptor_set_layout_destroy(struct vn_device *dev,
@@ -126,6 +127,10 @@ vn_descriptor_set_layout_init(
    if (binding_flags && !binding_flags->bindingCount)
       binding_flags = NULL;
 
+   layout->is_push_descriptor =
+      create_info->flags &
+      VK_DESCRIPTOR_SET_LAYOUT_CREATE_PUSH_DESCRIPTOR_BIT_KHR;
+
    layout->refcount = VN_REFCOUNT_INIT(1);
    layout->last_binding = last_binding;
 
@@ -1066,6 +1071,15 @@ vn_CreateDescriptorUpdateTemplate(
       return vn_error(dev->instance, VK_ERROR_OUT_OF_HOST_MEMORY);
    }
 
+   templ->is_push_descriptor =
+      pCreateInfo->templateType ==
+      VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_PUSH_DESCRIPTORS_KHR;
+   if (templ->is_push_descriptor) {
+      templ->pipeline_bind_point = pCreateInfo->pipelineBindPoint;
+      templ->pipeline_layout =
+         vn_pipeline_layout_from_handle(pCreateInfo->pipelineLayout);
+   }
+
    mtx_init(&templ->mutex, mtx_plain);
 
    /* no host object */
@@ -1100,32 +1114,30 @@ vn_DestroyDescriptorUpdateTemplate(
    vk_free(alloc, templ);
 }
 
-void
-vn_UpdateDescriptorSetWithTemplate(
-   VkDevice device,
-   VkDescriptorSet descriptorSet,
-   VkDescriptorUpdateTemplate descriptorUpdateTemplate,
-   const void *pData)
+struct vn_update_descriptor_sets *
+vn_update_descriptor_set_with_template_locked(
+   struct vn_descriptor_update_template *templ,
+   struct vn_descriptor_set *set,
+   const void *data)
 {
-   VN_TRACE_FUNC();
-   struct vn_device *dev = vn_device_from_handle(device);
-   struct vn_descriptor_set *set =
-      vn_descriptor_set_from_handle(descriptorSet);
-   struct vn_descriptor_update_template *templ =
-      vn_descriptor_update_template_from_handle(descriptorUpdateTemplate);
    struct vn_update_descriptor_sets *update = templ->update;
 
-   /* duplicate update instead to avoid locking? */
-   mtx_lock(&templ->mutex);
-
    for (uint32_t i = 0; i < update->write_count; i++) {
       const struct vn_descriptor_update_template_entry *entry =
          &templ->entries[i];
+
+      const struct vn_descriptor_set_layout *set_layout =
+         templ->is_push_descriptor
+            ? templ->pipeline_layout->push_descriptor_set_layout
+            : set->layout;
       const struct vn_descriptor_set_layout_binding *binding =
-         &set->layout->bindings[update->writes[i].dstBinding];
+         &set_layout->bindings[update->writes[i].dstBinding];
+
       VkWriteDescriptorSet *write = &update->writes[i];
 
-      write->dstSet = vn_descriptor_set_to_handle(set);
+      write->dstSet = templ->is_push_descriptor
+                         ? VK_NULL_HANDLE
+                         : vn_descriptor_set_to_handle(set);
 
       switch (write->descriptorType) {
       case VK_DESCRIPTOR_TYPE_SAMPLER:
@@ -1142,7 +1154,7 @@ vn_UpdateDescriptorSetWithTemplate(
             const bool need_view =
                write->descriptorType != VK_DESCRIPTOR_TYPE_SAMPLER;
             const VkDescriptorImageInfo *src =
-               pData + entry->offset + entry->stride * j;
+               data + entry->offset + entry->stride * j;
             VkDescriptorImageInfo *dst =
                (VkDescriptorImageInfo *)&write->pImageInfo[j];
 
@@ -1155,7 +1167,7 @@ vn_UpdateDescriptorSetWithTemplate(
       case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER:
          for (uint32_t j = 0; j < write->descriptorCount; j++) {
             const VkBufferView *src =
-               pData + entry->offset + entry->stride * j;
+               data + entry->offset + entry->stride * j;
             VkBufferView *dst = (VkBufferView *)&write->pTexelBufferView[j];
             *dst = *src;
          }
@@ -1166,7 +1178,7 @@ vn_UpdateDescriptorSetWithTemplate(
       case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC:
          for (uint32_t j = 0; j < write->descriptorCount; j++) {
             const VkDescriptorBufferInfo *src =
-               pData + entry->offset + entry->stride * j;
+               data + entry->offset + entry->stride * j;
             VkDescriptorBufferInfo *dst =
                (VkDescriptorBufferInfo *)&write->pBufferInfo[j];
             *dst = *src;
@@ -1176,7 +1188,7 @@ vn_UpdateDescriptorSetWithTemplate(
          VkWriteDescriptorSetInlineUniformBlock *iub_data =
             (VkWriteDescriptorSetInlineUniformBlock *)vk_find_struct_const(
                write->pNext, WRITE_DESCRIPTOR_SET_INLINE_UNIFORM_BLOCK);
-         iub_data->pData = pData + entry->offset;
+         iub_data->pData = data + entry->offset;
          break;
       case VK_DESCRIPTOR_TYPE_MUTABLE_EXT:
          break;
@@ -1185,6 +1197,26 @@ vn_UpdateDescriptorSetWithTemplate(
          break;
       }
    }
+   return update;
+}
+
+void
+vn_UpdateDescriptorSetWithTemplate(
+   VkDevice device,
+   VkDescriptorSet descriptorSet,
+   VkDescriptorUpdateTemplate descriptorUpdateTemplate,
+   const void *pData)
+{
+   VN_TRACE_FUNC();
+   struct vn_device *dev = vn_device_from_handle(device);
+   struct vn_descriptor_update_template *templ =
+      vn_descriptor_update_template_from_handle(descriptorUpdateTemplate);
+   struct vn_descriptor_set *set =
+      vn_descriptor_set_from_handle(descriptorSet);
+   mtx_lock(&templ->mutex);
+
+   struct vn_update_descriptor_sets *update =
+      vn_update_descriptor_set_with_template_locked(templ, set, pData);
 
    vn_async_vkUpdateDescriptorSets(dev->instance, device, update->write_count,
                                    update->writes, 0, NULL);

src/virtio/vulkan/vn_descriptor_set.h

@@ -51,6 +51,7 @@ struct vn_descriptor_set_layout {
 
    uint32_t last_binding;
    bool has_variable_descriptor_count;
+   bool is_push_descriptor;
 
    /* bindings must be the last field in the layout */
    struct vn_descriptor_set_layout_binding bindings[];
@@ -120,6 +121,10 @@ struct vn_descriptor_update_template_entry {
 struct vn_descriptor_update_template {
    struct vn_object_base base;
 
+   bool is_push_descriptor;
+   VkPipelineBindPoint pipeline_bind_point;
+   struct vn_pipeline_layout *pipeline_layout;
+
    mtx_t mutex;
    struct vn_update_descriptor_sets *update;
 
@@ -130,7 +135,13 @@ VK_DEFINE_NONDISP_HANDLE_CASTS(vn_descriptor_update_template,
                                VkDescriptorUpdateTemplate,
                                VK_OBJECT_TYPE_DESCRIPTOR_UPDATE_TEMPLATE)
 
- void
+struct vn_update_descriptor_sets *
+vn_update_descriptor_set_with_template_locked(
+   struct vn_descriptor_update_template *templ,
+   struct vn_descriptor_set *set,
+   const void *data);
+
+void
 vn_descriptor_set_layout_destroy(struct vn_device *dev,
                                  struct vn_descriptor_set_layout *layout);