turnip: initial implementation of VK_KHR_push_descriptor

Add missing descriptor sets code for push descriptors, and a simple initial
implementation to enable the extension and pass dEQP tests.

Signed-off-by: Jonathan Marek <jonathan@marek.ca>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/6755>

.gitlab-ci/deqp-freedreno-a630-flakes.txt

@@ -38,6 +38,9 @@ dEQP-VK.memory_model.write_after_read.core11.u32.coherent.fence_fence.atomicwrit
 # https://gitlab.khronos.org/Tracker/vk-gl-cts/-/issues/2017
 dEQP-VK.renderpass.*separate_channels.*
 
+# These tests are broken (does not respect our minStorageBufferOffsetAlignment of 64)
+dEQP-VK.pipeline.push_descriptor.compute.*
+
 # Undiagnosed flakes appearing more than once in the last 2 months as
 # of 2020-08-19, in descending order of frequency.
 dEQP-GLES3.functional.fbo.msaa.2_samples.stencil_index8

src/freedreno/vulkan/tu_cmd_buffer.c

@@ -1771,6 +1771,76 @@ tu_CmdBindDescriptorSets(VkCommandBuffer commandBuffer,
    }
 }
 
+void tu_CmdPushDescriptorSetKHR(VkCommandBuffer commandBuffer,
+                                VkPipelineBindPoint pipelineBindPoint,
+                                VkPipelineLayout _layout,
+                                uint32_t _set,
+                                uint32_t descriptorWriteCount,
+                                const VkWriteDescriptorSet *pDescriptorWrites)
+{
+   TU_FROM_HANDLE(tu_cmd_buffer, cmd, commandBuffer);
+   TU_FROM_HANDLE(tu_pipeline_layout, pipe_layout, _layout);
+   struct tu_descriptor_set_layout *layout = pipe_layout->set[_set].layout;
+   struct tu_descriptor_set *set =
+      &tu_get_descriptors_state(cmd, pipelineBindPoint)->push_set;
+
+   struct tu_cs_memory set_mem;
+   VkResult result = tu_cs_alloc(&cmd->sub_cs,
+                                 DIV_ROUND_UP(layout->size, A6XX_TEX_CONST_DWORDS * 4),
+                                 A6XX_TEX_CONST_DWORDS, &set_mem);
+   assert(result == VK_SUCCESS);
+
+   /* preserve previous content if the layout is the same: */
+   if (set->layout == layout)
+      memcpy(set_mem.map, set->mapped_ptr, MIN2(set->size, layout->size));
+
+   set->layout = layout;
+   set->mapped_ptr = set_mem.map;
+   set->va = set_mem.iova;
+
+   tu_update_descriptor_sets(tu_descriptor_set_to_handle(set),
+                             descriptorWriteCount, pDescriptorWrites, 0, NULL);
+
+   tu_CmdBindDescriptorSets(commandBuffer, pipelineBindPoint, _layout, _set,
+                            1, (VkDescriptorSet[]) { tu_descriptor_set_to_handle(set) },
+                            0, NULL);
+}
+
+void tu_CmdPushDescriptorSetWithTemplateKHR(
+   VkCommandBuffer commandBuffer,
+   VkDescriptorUpdateTemplate descriptorUpdateTemplate,
+   VkPipelineLayout _layout,
+   uint32_t _set,
+   const void* pData)
+{
+   TU_FROM_HANDLE(tu_cmd_buffer, cmd, commandBuffer);
+   TU_FROM_HANDLE(tu_pipeline_layout, pipe_layout, _layout);
+   TU_FROM_HANDLE(tu_descriptor_update_template, templ, descriptorUpdateTemplate);
+   struct tu_descriptor_set_layout *layout = pipe_layout->set[_set].layout;
+   struct tu_descriptor_set *set =
+      &tu_get_descriptors_state(cmd, templ->bind_point)->push_set;
+
+   struct tu_cs_memory set_mem;
+   VkResult result = tu_cs_alloc(&cmd->sub_cs,
+                                 DIV_ROUND_UP(layout->size, A6XX_TEX_CONST_DWORDS * 4),
+                                 A6XX_TEX_CONST_DWORDS, &set_mem);
+   assert(result == VK_SUCCESS);
+
+   /* preserve previous content if the layout is the same: */
+   if (set->layout == layout)
+      memcpy(set_mem.map, set->mapped_ptr, MIN2(set->size, layout->size));
+
+   set->layout = layout;
+   set->mapped_ptr = set_mem.map;
+   set->va = set_mem.iova;
+
+   tu_update_descriptor_set_with_template(set, descriptorUpdateTemplate, pData);
+
+   tu_CmdBindDescriptorSets(commandBuffer, templ->bind_point, _layout, _set,
+                            1, (VkDescriptorSet[]) { tu_descriptor_set_to_handle(set) },
+                            0, NULL);
+}
+
 void tu_CmdBindTransformFeedbackBuffersEXT(VkCommandBuffer commandBuffer,
                                            uint32_t firstBinding,
                                            uint32_t bindingCount,

src/freedreno/vulkan/tu_descriptor_set.c

@@ -794,6 +794,13 @@ write_sampler_descriptor(uint32_t *dst, const VkDescriptorImageInfo *image_info)
    memcpy(dst, sampler->descriptor, sizeof(sampler->descriptor));
 }
 
+/* note: this is used with immutable samplers in push descriptors */
+static void
+write_sampler_push(uint32_t *dst, const struct tu_sampler *sampler)
+{
+   memcpy(dst, sampler->descriptor, sizeof(sampler->descriptor));
+}
+
 void
 tu_update_descriptor_sets(VkDescriptorSet dstSetOverride,
                           uint32_t descriptorWriteCount,
@@ -804,11 +811,15 @@ tu_update_descriptor_sets(VkDescriptorSet dstSetOverride,
    uint32_t i, j;
    for (i = 0; i < descriptorWriteCount; i++) {
       const VkWriteDescriptorSet *writeset = &pDescriptorWrites[i];
-      TU_FROM_HANDLE(tu_descriptor_set, set,
+      TU_FROM_HANDLE(tu_descriptor_set, set, dstSetOverride ?: writeset->dstSet);
-                       dstSetOverride ? dstSetOverride : writeset->dstSet);
       const struct tu_descriptor_set_binding_layout *binding_layout =
          set->layout->binding + writeset->dstBinding;
       uint32_t *ptr = set->mapped_ptr;
+      /* for immutable samplers with push descriptors: */
+      const bool copy_immutable_samplers =
+         dstSetOverride && binding_layout->immutable_samplers_offset;
+      const struct tu_sampler *samplers =
+         tu_immutable_samplers(set->layout, binding_layout);
 
       ptr += binding_layout->offset / 4;
 
@@ -850,9 +861,15 @@ tu_update_descriptor_sets(VkDescriptorSet dstSetOverride,
                                                     writeset->descriptorType,
                                                     writeset->pImageInfo + j,
                                                     !binding_layout->immutable_samplers_offset);
+
+            if (copy_immutable_samplers)
+               write_sampler_push(ptr + A6XX_TEX_CONST_DWORDS, &samplers[writeset->dstArrayElement + j]);
             break;
          case VK_DESCRIPTOR_TYPE_SAMPLER:
-            write_sampler_descriptor(ptr, writeset->pImageInfo + j);
+            if (!binding_layout->immutable_samplers_offset)
+               write_sampler_descriptor(ptr, writeset->pImageInfo + j);
+            else if (copy_immutable_samplers)
+               write_sampler_push(ptr, &samplers[writeset->dstArrayElement + j]);
             break;
          case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
             /* nothing in descriptor set - framebuffer state is used instead */
@@ -952,6 +969,8 @@ tu_CreateDescriptorUpdateTemplate(
        */
       assert(pCreateInfo->set < MAX_SETS);
       set_layout = pipeline_layout->set[pCreateInfo->set].layout;
+
+      templ->bind_point = pCreateInfo->pipelineBindPoint;
    }
 
    for (uint32_t i = 0; i < entry_count; i++) {
@@ -960,6 +979,7 @@ tu_CreateDescriptorUpdateTemplate(
       const struct tu_descriptor_set_binding_layout *binding_layout =
          set_layout->binding + entry->dstBinding;
       uint32_t dst_offset, dst_stride;
+      const struct tu_sampler *immutable_samplers = NULL;
 
       /* dst_offset is an offset into dynamic_descriptors when the descriptor 
        * is dynamic, and an offset into mapped_ptr otherwise.
@@ -971,6 +991,14 @@ tu_CreateDescriptorUpdateTemplate(
             entry->dstArrayElement) * A6XX_TEX_CONST_DWORDS;
          dst_stride = A6XX_TEX_CONST_DWORDS;
          break;
+      case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
+      case VK_DESCRIPTOR_TYPE_SAMPLER:
+         if (pCreateInfo->templateType == VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_PUSH_DESCRIPTORS_KHR &&
+             binding_layout->immutable_samplers_offset) {
+            immutable_samplers =
+               tu_immutable_samplers(set_layout, binding_layout) + entry->dstArrayElement;
+         }
+         /* fallthrough */
       default:
          dst_offset = binding_layout->offset / 4;
          dst_offset += (binding_layout->size * entry->dstArrayElement) / 4;
@@ -985,6 +1013,7 @@ tu_CreateDescriptorUpdateTemplate(
          .dst_offset = dst_offset,
          .dst_stride = dst_stride,
          .has_sampler = !binding_layout->immutable_samplers_offset,
+         .immutable_samplers = immutable_samplers,
       };
    }
 
@@ -1022,6 +1051,7 @@ tu_update_descriptor_set_with_template(
    for (uint32_t i = 0; i < templ->entry_count; i++) {
       uint32_t *ptr = set->mapped_ptr;
       const void *src = ((const char *) pData) + templ->entry[i].src_offset;
+      const struct tu_sampler *samplers = templ->entry[i].immutable_samplers;
 
       ptr += templ->entry[i].dst_offset;
       unsigned dst_offset = templ->entry[i].dst_offset;
@@ -1057,9 +1087,14 @@ tu_update_descriptor_set_with_template(
                                                     templ->entry[i].descriptor_type,
                                                     src,
                                                     templ->entry[i].has_sampler);
+            if (samplers)
+               write_sampler_push(ptr + A6XX_TEX_CONST_DWORDS, &samplers[j]);
             break;
          case VK_DESCRIPTOR_TYPE_SAMPLER:
-            write_sampler_descriptor(ptr, src);
+            if (templ->entry[i].has_sampler)
+               write_sampler_descriptor(ptr, src);
+            else if (samplers)
+               write_sampler_push(ptr, &samplers[j]);
             break;
          case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
             /* nothing in descriptor set - framebuffer state is used instead */

src/freedreno/vulkan/tu_descriptor_set.h

@@ -108,7 +108,7 @@ struct tu_pipeline_layout
    uint32_t dynamic_offset_count;
 };
 
-static inline const uint32_t *
+static inline const struct tu_sampler *
 tu_immutable_samplers(const struct tu_descriptor_set_layout *set,
                       const struct tu_descriptor_set_binding_layout *binding)
 {

src/freedreno/vulkan/tu_extensions.py

@@ -96,6 +96,7 @@ EXTENSIONS = [
     Extension('VK_EXT_host_query_reset',                  1, True),
     Extension('VK_EXT_shader_viewport_index_layer',       1, True),
     Extension('VK_EXT_extended_dynamic_state',            1, True),
+    Extension('VK_KHR_push_descriptor',                   1, True),
 ]
 
 MAX_API_VERSION = VkVersion(MAX_API_VERSION)

src/freedreno/vulkan/tu_private.h

@@ -588,12 +588,6 @@ struct tu_descriptor_set
    uint32_t *dynamic_descriptors;
 };
 
-struct tu_push_descriptor_set
-{
-   struct tu_descriptor_set set;
-   uint32_t capacity;
-};
-
 struct tu_descriptor_pool_entry
 {
    uint32_t offset;
@@ -642,7 +636,7 @@ struct tu_descriptor_update_template_entry
    size_t src_stride;
 
    /* For push descriptors */
-   const uint32_t *immutable_samplers;
+   const struct tu_sampler *immutable_samplers;
 };
 
 struct tu_descriptor_update_template
@@ -650,6 +644,7 @@ struct tu_descriptor_update_template
    struct vk_object_base base;
 
    uint32_t entry_count;
+   VkPipelineBindPoint bind_point;
    struct tu_descriptor_update_template_entry entry[0];
 };
 
@@ -681,6 +676,7 @@ tu_get_perftest_option_name(int id);
 struct tu_descriptor_state
 {
    struct tu_descriptor_set *sets[MAX_SETS];
+   struct tu_descriptor_set push_set;
    uint32_t dynamic_descriptors[MAX_DYNAMIC_BUFFERS * A6XX_TEX_CONST_DWORDS];
 };
 
@@ -1514,12 +1510,6 @@ struct tu_semaphore
    struct tu_semaphore_part temporary;
 };
 
-void
-tu_set_descriptor_set(struct tu_cmd_buffer *cmd_buffer,
-                      VkPipelineBindPoint bind_point,
-                      struct tu_descriptor_set *set,
-                      unsigned idx);
-
 void
 tu_update_descriptor_sets(VkDescriptorSet overrideSet,
                           uint32_t descriptorWriteCount,