diff mbox series

[FFmpeg-devel,2/5] vulkan: move shader data execution state to execution pools

Message ID 20240929094256.396352-2-dev@lynne.ee
State New
Headers show
Series [FFmpeg-devel,1/5] vulkan: merge FFVkSPIRVShader and FFVkPipeline into FFVkShader | expand

Checks

Context Check Description
andriy/make_x86 success Make finished
andriy/make_fate_x86 success Make fate finished

Commit Message

Lynne Sept. 29, 2024, 9:42 a.m. UTC
This finally permits using fully compiled shaders across
multiple execution contexts.
---
 libavutil/vulkan.c | 157 +++++++++++++++++++++++++++++----------------
 libavutil/vulkan.h |  81 ++++++++++++++---------
 2 files changed, 153 insertions(+), 85 deletions(-)
diff mbox series

Patch

diff --git a/libavutil/vulkan.c b/libavutil/vulkan.c
index 9ea9e478df..71581b91a9 100644
--- a/libavutil/vulkan.c
+++ b/libavutil/vulkan.c
@@ -262,6 +262,28 @@  void ff_vk_exec_pool_free(FFVulkanContext *s, FFVkExecPool *pool)
         av_free(e->sem_wait);
     }
 
+    /* Free shader-specific data */
+    for (int i = 0; i < pool->nb_reg_shd; i++) {
+        FFVulkanShaderData *sd = &pool->reg_shd[i];
+
+        if (s->extensions & FF_VK_EXT_DESCRIPTOR_BUFFER) {
+            for (int j = 0; j < sd->nb_descriptor_sets; j++) {
+                FFVulkanDescriptorSetData *set_data = &sd->desc_set_buf[j];
+                if (set_data->buf.mem)
+                    ff_vk_unmap_buffer(s, &set_data->buf, 0);
+                ff_vk_free_buf(s, &set_data->buf);
+            }
+        }
+
+        if (sd->desc_pool)
+            vk->DestroyDescriptorPool(s->hwctx->act_dev, sd->desc_pool,
+                                      s->hwctx->alloc);
+
+        av_freep(&sd->desc_set_buf);
+        av_freep(&sd->desc_bind);
+        av_freep(&sd->desc_sets);
+    }
+
     if (pool->cmd_bufs)
         vk->FreeCommandBuffers(s->hwctx->act_dev, pool->cmd_buf_pool,
                                pool->pool_size, pool->cmd_bufs);
@@ -1658,6 +1680,16 @@  int ff_vk_shader_link(FFVulkanContext *s, FFVulkanShader *shd,
     if (err < 0)
         return err;
 
+    if (s->extensions & FF_VK_EXT_DESCRIPTOR_BUFFER) {
+        shd->bound_buffer_indices = av_calloc(shd->nb_descriptor_sets,
+                                              sizeof(*shd->bound_buffer_indices));
+        if (!shd->bound_buffer_indices)
+            return AVERROR(ENOMEM);
+
+        for (int i = 0; i < shd->nb_descriptor_sets; i++)
+            shd->bound_buffer_indices[i] = i;
+    }
+
     {
         VkShaderModule mod;
         err = create_shader_module(s, shd, &mod, spirv, spirv_len);
@@ -1820,25 +1852,40 @@  int ff_vk_shader_register_exec(FFVulkanContext *s, FFVkExecPool *pool,
                                FFVulkanShader *shd)
 {
     int err;
+    FFVulkanShaderData *sd;
 
     if (!shd->nb_descriptor_sets)
         return 0;
 
+    sd = av_realloc_array(pool->reg_shd,
+                          sizeof(*pool->reg_shd),
+                          pool->nb_reg_shd + 1);
+    if (!sd)
+        return AVERROR(ENOMEM);
+
+    pool->reg_shd = sd;
+    sd = &sd[pool->nb_reg_shd++];
+    memset(sd, 0, sizeof(*sd));
+
+    sd->shd = shd;
+    sd->nb_descriptor_sets = shd->nb_descriptor_sets;
+
     if (s->extensions & FF_VK_EXT_DESCRIPTOR_BUFFER) {
-        shd->desc_bind = av_calloc(shd->nb_descriptor_sets, sizeof(*shd->desc_bind));
-        if (!shd->desc_bind)
+        sd->desc_bind = av_malloc_array(sd->nb_descriptor_sets, sizeof(*sd->desc_bind));
+        if (!sd->desc_bind)
             return AVERROR(ENOMEM);
 
-        shd->bound_buffer_indices = av_calloc(shd->nb_descriptor_sets,
-                                             sizeof(*shd->bound_buffer_indices));
-        if (!shd->bound_buffer_indices)
+        sd->desc_set_buf = av_calloc(sd->nb_descriptor_sets, sizeof(*sd->desc_set_buf));
+        if (!sd->desc_set_buf)
             return AVERROR(ENOMEM);
 
-        for (int i = 0; i < shd->nb_descriptor_sets; i++) {
+        for (int i = 0; i < sd->nb_descriptor_sets; i++) {
             FFVulkanDescriptorSet *set = &shd->desc_set[i];
+            FFVulkanDescriptorSetData *sdb = &sd->desc_set_buf[i];
             int nb = set->singular ? 1 : pool->pool_size;
 
-            err = ff_vk_create_buf(s, &set->buf, set->aligned_size*nb,
+            err = ff_vk_create_buf(s, &sdb->buf,
+                                   set->aligned_size*nb,
                                    NULL, NULL, set->usage,
                                    VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT |
                                    VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
@@ -1846,17 +1893,15 @@  int ff_vk_shader_register_exec(FFVulkanContext *s, FFVkExecPool *pool,
             if (err < 0)
                 return err;
 
-            err = ff_vk_map_buffer(s, &set->buf, &set->desc_mem, 0);
+            err = ff_vk_map_buffer(s, &sdb->buf, &sdb->desc_mem, 0);
             if (err < 0)
                 return err;
 
-            shd->desc_bind[i] = (VkDescriptorBufferBindingInfoEXT) {
+            sd->desc_bind[i] = (VkDescriptorBufferBindingInfoEXT) {
                 .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_BUFFER_BINDING_INFO_EXT,
                 .usage = set->usage,
-                .address = set->buf.address,
+                .address = sdb->buf.address,
             };
-
-            shd->bound_buffer_indices[i] = i;
         }
     } else if (!shd->use_push) {
         VkResult ret;
@@ -1873,11 +1918,11 @@  int ff_vk_shader_register_exec(FFVulkanContext *s, FFVkExecPool *pool,
             .flags = 0,
             .pPoolSizes = shd->desc_pool_size,
             .poolSizeCount = shd->nb_desc_pool_size,
-            .maxSets = shd->nb_descriptor_sets*pool->pool_size,
+            .maxSets = sd->nb_descriptor_sets*pool->pool_size,
         };
 
         ret = vk->CreateDescriptorPool(s->hwctx->act_dev, &pool_create_info,
-                                       s->hwctx->alloc, &shd->desc_pool);
+                                       s->hwctx->alloc, &sd->desc_pool);
         if (ret != VK_SUCCESS) {
             av_log(s, AV_LOG_ERROR, "Unable to create descriptor pool: %s\n",
                    ff_vk_ret2str(ret));
@@ -1890,63 +1935,75 @@  int ff_vk_shader_register_exec(FFVulkanContext *s, FFVkExecPool *pool,
 
         /* Colate each execution context's descriptor set layouts */
         for (int i = 0; i < pool->pool_size; i++)
-            for (int j = 0; j < shd->nb_descriptor_sets; j++)
-                tmp_layouts[i*shd->nb_descriptor_sets + j] = shd->desc_layout[j];
+            for (int j = 0; j < sd->nb_descriptor_sets; j++)
+                tmp_layouts[i*sd->nb_descriptor_sets + j] = shd->desc_layout[j];
 
         set_alloc_info = (VkDescriptorSetAllocateInfo) {
             .sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
-            .descriptorPool     = shd->desc_pool,
+            .descriptorPool     = sd->desc_pool,
             .pSetLayouts        = tmp_layouts,
             .descriptorSetCount = pool_create_info.maxSets,
         };
 
-        shd->desc_sets = av_malloc_array(pool_create_info.maxSets,
+        sd->desc_sets = av_malloc_array(pool_create_info.maxSets,
                                         sizeof(*tmp_layouts));
-        if (!shd->desc_sets) {
+        if (!sd->desc_sets) {
             av_free(tmp_layouts);
             return AVERROR(ENOMEM);
         }
         ret = vk->AllocateDescriptorSets(s->hwctx->act_dev, &set_alloc_info,
-                                         shd->desc_sets);
+                                         sd->desc_sets);
         av_free(tmp_layouts);
         if (ret != VK_SUCCESS) {
             av_log(s, AV_LOG_ERROR, "Unable to allocate descriptor set: %s\n",
                    ff_vk_ret2str(ret));
-            av_freep(&shd->desc_sets);
+            av_freep(&sd->desc_sets);
             return AVERROR_EXTERNAL;
         }
-
-        shd->assoc_pool = pool;
     }
 
     return 0;
 }
 
+static inline FFVulkanShaderData *get_shd_data(FFVkExecContext *e,
+                                               FFVulkanShader *shd)
+{
+    for (int i = 0; i < e->parent->nb_reg_shd; i++)
+        if (e->parent->reg_shd[i].shd == shd)
+            return &e->parent->reg_shd[i];
+    av_assert0(0);
+}
+
 static inline void update_set_descriptor(FFVulkanContext *s, FFVkExecContext *e,
-                                         FFVulkanDescriptorSet *set,
+                                         FFVulkanShader *shd, int set,
                                          int bind_idx, int array_idx,
                                          VkDescriptorGetInfoEXT *desc_get_info,
                                          size_t desc_size)
 {
     FFVulkanFunctions *vk = &s->vkfn;
-    const size_t exec_offset = set->singular ? 0 : set->aligned_size*e->idx;
-    void *desc = set->desc_mem +                 /* Base */
-                 exec_offset +                   /* Execution context */
-                 set->binding_offset[bind_idx] + /* Descriptor binding */
-                 array_idx*desc_size;            /* Array position */
+    FFVulkanDescriptorSet *desc_set = &shd->desc_set[set];
+    FFVulkanShaderData *sd = get_shd_data(e, shd);
+    const size_t exec_offset = desc_set->singular ? 0 : desc_set->aligned_size*e->idx;
+
+    void *desc = sd->desc_set_buf[set].desc_mem +     /* Base */
+                 exec_offset +                        /* Execution context */
+                 desc_set->binding_offset[bind_idx] + /* Descriptor binding */
+                 array_idx*desc_size;                 /* Array position */
 
     vk->GetDescriptorEXT(s->hwctx->act_dev, desc_get_info, desc_size, desc);
 }
 
-static inline void update_set_pool_write(FFVulkanContext *s, FFVulkanShader *shd,
-                                         FFVkExecContext *e,
-                                         FFVulkanDescriptorSet *desc_set, int set,
+static inline void update_set_pool_write(FFVulkanContext *s, FFVkExecContext *e,
+                                         FFVulkanShader *shd, int set,
                                          VkWriteDescriptorSet *write_info)
 {
     FFVulkanFunctions *vk = &s->vkfn;
+    FFVulkanDescriptorSet *desc_set = &shd->desc_set[set];
+    FFVulkanShaderData *sd = get_shd_data(e, shd);
+
     if (desc_set->singular) {
-        for (int i = 0; i < shd->assoc_pool->pool_size; i++) {
-            write_info->dstSet = shd->desc_sets[i*shd->nb_descriptor_sets + set];
+        for (int i = 0; i < e->parent->pool_size; i++) {
+            write_info->dstSet = sd->desc_sets[i*sd->nb_descriptor_sets + set];
             vk->UpdateDescriptorSets(s->hwctx->act_dev, 1, write_info, 0, NULL);
         }
     } else {
@@ -1957,7 +2014,7 @@  static inline void update_set_pool_write(FFVulkanContext *s, FFVulkanShader *shd
                                         set, 1,
                                         write_info);
         } else {
-            write_info->dstSet = shd->desc_sets[e->idx*shd->nb_descriptor_sets + set];
+            write_info->dstSet = sd->desc_sets[e->idx*sd->nb_descriptor_sets + set];
             vk->UpdateDescriptorSets(s->hwctx->act_dev, 1, write_info, 0, NULL);
         }
     }
@@ -2006,7 +2063,7 @@  static int vk_set_descriptor_image(FFVulkanContext *s, FFVulkanShader *shd,
             break;
         };
 
-        update_set_descriptor(s, e, desc_set, bind, offs,
+        update_set_descriptor(s, e, shd, set, bind, offs,
                               &desc_get_info, desc_size);
     } else {
         VkDescriptorImageInfo desc_pool_write_info_img = {
@@ -2022,7 +2079,7 @@  static int vk_set_descriptor_image(FFVulkanContext *s, FFVulkanShader *shd,
             .descriptorType = desc_set->binding[bind].descriptorType,
             .pImageInfo = &desc_pool_write_info_img,
         };
-        update_set_pool_write(s, shd, e, desc_set, set, &desc_pool_write_info);
+        update_set_pool_write(s, e, shd, set, &desc_pool_write_info);
     }
 
     return 0;
@@ -2073,7 +2130,7 @@  int ff_vk_shader_update_desc_buffer(FFVulkanContext *s, FFVkExecContext *e,
             break;
         };
 
-        update_set_descriptor(s, e, desc_set, bind, elem, &desc_get_info, desc_size);
+        update_set_descriptor(s, e, shd, set, bind, elem, &desc_get_info, desc_size);
     } else {
         VkDescriptorBufferInfo desc_pool_write_info_buf = {
             .buffer = buf->buf,
@@ -2088,7 +2145,7 @@  int ff_vk_shader_update_desc_buffer(FFVulkanContext *s, FFVkExecContext *e,
             .descriptorType = desc_set->binding[bind].descriptorType,
             .pBufferInfo = &desc_pool_write_info_buf,
         };
-        update_set_pool_write(s, shd, e, desc_set, set, &desc_pool_write_info);
+        update_set_pool_write(s, e, shd, set, &desc_pool_write_info);
     }
 
     return 0;
@@ -2122,25 +2179,26 @@  void ff_vk_exec_bind_shader(FFVulkanContext *s, FFVkExecContext *e,
 {
     FFVulkanFunctions *vk = &s->vkfn;
     VkDeviceSize offsets[1024];
+    FFVulkanShaderData *sd = get_shd_data(e, shd);
 
     /* Bind pipeline */
     vk->CmdBindPipeline(e->buf, shd->bind_point, shd->pipeline);
 
-    if (shd->nb_descriptor_sets) {
+    if (sd->nb_descriptor_sets) {
         if (s->extensions & FF_VK_EXT_DESCRIPTOR_BUFFER) {
-            for (int i = 0; i < shd->nb_descriptor_sets; i++)
+            for (int i = 0; i < sd->nb_descriptor_sets; i++)
                 offsets[i] = shd->desc_set[i].singular ? 0 : shd->desc_set[i].aligned_size*e->idx;
 
             /* Bind descriptor buffers */
-            vk->CmdBindDescriptorBuffersEXT(e->buf, shd->nb_descriptor_sets, shd->desc_bind);
+            vk->CmdBindDescriptorBuffersEXT(e->buf, sd->nb_descriptor_sets, sd->desc_bind);
             /* Binding offsets */
             vk->CmdSetDescriptorBufferOffsetsEXT(e->buf, shd->bind_point, shd->pipeline_layout,
-                                                 0, shd->nb_descriptor_sets,
+                                                 0, sd->nb_descriptor_sets,
                                                  shd->bound_buffer_indices, offsets);
         } else if (!shd->use_push) {
             vk->CmdBindDescriptorSets(e->buf, shd->bind_point, shd->pipeline_layout,
-                                      0, shd->nb_descriptor_sets,
-                                      &shd->desc_sets[e->idx*shd->nb_descriptor_sets],
+                                      0, sd->nb_descriptor_sets,
+                                      &sd->desc_sets[e->idx*sd->nb_descriptor_sets],
                                       0, NULL);
         }
     }
@@ -2166,9 +2224,6 @@  void ff_vk_shader_free(FFVulkanContext *s, FFVulkanShader *shd)
 
     for (int i = 0; i < shd->nb_descriptor_sets; i++) {
         FFVulkanDescriptorSet *set = &shd->desc_set[i];
-        if (set->buf.mem)
-            ff_vk_unmap_buffer(s, &set->buf, 0);
-        ff_vk_free_buf(s, &set->buf);
         av_free(set->binding);
         av_free(set->binding_offset);
     }
@@ -2178,15 +2233,9 @@  void ff_vk_shader_free(FFVulkanContext *s, FFVulkanShader *shd)
             vk->DestroyDescriptorSetLayout(s->hwctx->act_dev, shd->desc_layout[i],
                                            s->hwctx->alloc);
 
-    if (shd->desc_pool)
-        vk->DestroyDescriptorPool(s->hwctx->act_dev, shd->desc_pool,
-                                  s->hwctx->alloc);
-
     av_freep(&shd->desc_pool_size);
     av_freep(&shd->desc_layout);
-    av_freep(&shd->desc_sets);
     av_freep(&shd->desc_set);
-    av_freep(&shd->desc_bind);
     av_freep(&shd->bound_buffer_indices);
     av_freep(&shd->push_consts);
     shd->push_consts_num = 0;
diff --git a/libavutil/vulkan.h b/libavutil/vulkan.h
index 9c87dd29cf..5391afa4e7 100644
--- a/libavutil/vulkan.h
+++ b/libavutil/vulkan.h
@@ -162,31 +162,11 @@  typedef struct FFVkExecContext {
     unsigned int frame_update_alloc_size;
 } FFVkExecContext;
 
-typedef struct FFVkExecPool {
-    FFVkExecContext *contexts;
-    atomic_int_least64_t idx;
-
-    VkCommandPool cmd_buf_pool;
-    VkCommandBuffer *cmd_bufs;
-    int pool_size;
-
-    VkQueryPool query_pool;
-    void *query_data;
-    int query_results;
-    int query_statuses;
-    int query_64bit;
-    int query_status_stride;
-    int nb_queries;
-    size_t qd_size;
-} FFVkExecPool;
-
 typedef struct FFVulkanDescriptorSet {
-    FFVkBuffer             buf;
-    uint8_t               *desc_mem;
-    VkDeviceSize           layout_size;
-    VkDeviceSize           aligned_size; /* descriptorBufferOffsetAlignment */
-    VkDeviceSize           total_size; /* Once registered to an exec context */
-    VkBufferUsageFlags     usage;
+    /* Descriptor buffer */
+    VkDeviceSize layout_size;
+    VkDeviceSize aligned_size; /* descriptorBufferOffsetAlignment */
+    VkBufferUsageFlags usage;
 
     VkDescriptorSetLayoutBinding *binding;
     VkDeviceSize *binding_offset;
@@ -225,23 +205,62 @@  typedef struct FFVulkanShader {
     VkPushConstantRange *push_consts;
     int push_consts_num;
 
+    /* Descriptor sets */
+    FFVulkanDescriptorSet *desc_set;
+    int nb_descriptor_sets;
+
     /* Descriptor buffer */
     VkDescriptorSetLayout *desc_layout;
-    FFVulkanDescriptorSet *desc_set;
-    VkDescriptorBufferBindingInfoEXT *desc_bind;
     uint32_t *bound_buffer_indices;
-    int nb_descriptor_sets;
 
     /* Descriptor pool */
     int use_push;
-    VkDescriptorSet *desc_sets;
-    VkDescriptorPool desc_pool;
     VkDescriptorPoolSize *desc_pool_size;
     int nb_desc_pool_size;
-    int total_desc_sets;
-    FFVkExecPool *assoc_pool;
 } FFVulkanShader;
 
+typedef struct FFVulkanDescriptorSetData {
+    /* Descriptor buffer */
+    FFVkBuffer buf;
+    uint8_t *desc_mem;
+} FFVulkanDescriptorSetData;
+
+typedef struct FFVulkanShaderData {
+    /* Shader to which this data belongs to */
+    FFVulkanShader *shd;
+    int nb_descriptor_sets;
+
+    /* Descriptor buffer */
+    FFVulkanDescriptorSetData *desc_set_buf;
+    VkDescriptorBufferBindingInfoEXT *desc_bind;
+
+    /* Descriptor pools */
+    VkDescriptorSet *desc_sets;
+    VkDescriptorPool desc_pool;
+} FFVulkanShaderData;
+
+typedef struct FFVkExecPool {
+    FFVkExecContext *contexts;
+    atomic_int_least64_t idx;
+
+    VkCommandPool cmd_buf_pool;
+    VkCommandBuffer *cmd_bufs;
+    int pool_size;
+
+    VkQueryPool query_pool;
+    void *query_data;
+    int query_results;
+    int query_statuses;
+    int query_64bit;
+    int query_status_stride;
+    int nb_queries;
+    size_t qd_size;
+
+    /* Registered shaders' data */
+    FFVulkanShaderData *reg_shd;
+    int nb_reg_shd;
+} FFVkExecPool;
+
 typedef struct FFVulkanContext {
     const AVClass *class;
     void *log_parent;