proresenc_kostya_vulkan: take the picture size from the gather

The trellis summed the estimated slice sizes into a separate host buffer
for the CPU to step over the payload with, and wrote each slice's start
offset, which nothing has read since slices went into fixed slots. The
gather writes the packed size of the slices after their sizes; use that,
and drop the buffer and the prefix sums.

The barrier after the trellis now covers the quantisers the encoder
reads, rather than the frame size buffer.

libavcodec/proresenc_kostya_vulkan.c

@@ -74,7 +74,6 @@ typedef struct SliceScore {
     int total_bits[MAX_STORED_Q];
     int total_error[MAX_STORED_Q];
     int overquant;
-    int buf_start;
     int quant;
 } SliceScore;
 
@@ -85,7 +84,6 @@ typedef struct VulkanEncodeProresFrameData {
     FFVkBuffer *gathered_ref;
     FFVkBuffer *slice_data_ref[2];
     FFVkBuffer *slice_score_ref[2];
-    FFVkBuffer *frame_size_ref[2];
 
     /* Copied from the source */
     enum AVColorTransferCharacteristic color_trc;
@@ -106,7 +104,6 @@ typedef struct ProresVulkanContext {
     AVRefStructPool *gathered_buf_pool;
     AVRefStructPool *slice_data_buf_pool;
     AVRefStructPool *slice_score_buf_pool;
-    AVRefStructPool *frame_size_buf_pool;
 
     size_t slice_slot_size;
     int estimate_slices_per_wg;
@@ -284,13 +281,9 @@ static int init_trellis_node_pipeline(ProresVulkanContext *pv, FFVulkanShader* s
     int err = 0;
     FFVulkanContext *vkctx = &pv->vkctx;
     FFVulkanDescriptorSetBinding *desc;
-    /* The driver picks the subgroup size; size for the most it may use */
-    int subgroup_size = vkctx->subgroup_props.minSubgroupSize;
-    int num_subgroups = FFALIGN(pv->ctx.mb_height, subgroup_size) / subgroup_size;
 
-    SPEC_LIST_CREATE(sl, 8, 8 * sizeof(uint32_t))
+    SPEC_LIST_CREATE(sl, 7, 7 * sizeof(uint32_t))
     SPEC_LIST_ADD(sl, 0, 32, pv->ctx.slices_width);
-    SPEC_LIST_ADD(sl, 1, 32, num_subgroups);
     SPEC_LIST_ADD(sl, 2, 32, pv->ctx.num_planes);
     SPEC_LIST_ADD(sl, 3, 32, pv->ctx.force_quant);
     SPEC_LIST_ADD(sl, 4, 32, pv->ctx.profile_info->min_quant);
@@ -302,18 +295,13 @@ static int init_trellis_node_pipeline(ProresVulkanContext *pv, FFVulkanShader* s
                       (uint32_t []) { pv->ctx.mb_height, 1, 1 }, 0);
 
     desc = (FFVulkanDescriptorSetBinding []) {
-        {
-            .name        = "FrameSize",
-            .type        = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
-            .stages      = VK_SHADER_STAGE_COMPUTE_BIT,
-        },
         {
             .name        = "SliceScores",
             .type        = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
             .stages      = VK_SHADER_STAGE_COMPUTE_BIT,
         },
     };
-    ff_vk_shader_add_descriptor_set(vkctx, shd, desc, 2, 0);
+    ff_vk_shader_add_descriptor_set(vkctx, shd, desc, 1, 0);
 
     RET(ff_vk_shader_link(vkctx, shd,
                           ff_prores_ks_trellis_node_comp_spv_data,
@@ -388,7 +376,7 @@ static int encode_picture(AVCodecContext *avctx, FFVkExecContext *exec,
     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(vkctx->frames->sw_format);
     VkImageView views[AV_NUM_DATA_POINTERS];
     VkImageMemoryBarrier2 img_bar[AV_NUM_DATA_POINTERS];
-    FFVkBuffer *pkt_vk_buf, *slice_data_buf, *slice_score_buf, *frame_size_buf;
+    FFVkBuffer *pkt_vk_buf, *slice_data_buf, *slice_score_buf;
     SliceDataInfo slice_data_info;
     EncodeSliceInfo encode_info;
     FFVulkanShader *shd;
@@ -442,17 +430,6 @@ static int encode_picture(AVCodecContext *avctx, FFVkExecContext *exec,
     slice_score_buf = pd->slice_score_ref[picture_idx];
     ff_vk_exec_add_dep_refstruct(vkctx, exec, pd->slice_score_ref[picture_idx]);
 
-    /* Allocate buffer for writing frame size */
-    RET(ff_vk_get_pooled_buffer(vkctx, &pv->frame_size_buf_pool, &pd->frame_size_ref[picture_idx],
-                                VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
-                                VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, NULL,
-                                sizeof(int),
-                                VK_MEMORY_PROPERTY_HOST_CACHED_BIT |
-                                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
-                                VK_MEMORY_PROPERTY_HOST_COHERENT_BIT));
-    frame_size_buf = pd->frame_size_ref[picture_idx];
-    ff_vk_exec_add_dep_refstruct(vkctx, exec, pd->frame_size_ref[picture_idx]);
-
     /* Generate barriers and image views for frame images. */
     RET(ff_vk_exec_add_dep_frame(vkctx, exec, frame,
                                  VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
@@ -557,9 +534,6 @@ static int encode_picture(AVCodecContext *avctx, FFVkExecContext *exec,
 
     /* Compute optimal quant value for each slice */
     ff_vk_shader_update_desc_buffer(vkctx, exec, &pv->trellis_node_shd, 0, 0, 0,
-                                    frame_size_buf, 0, frame_size_buf->size,
-                                    VK_FORMAT_UNDEFINED);
-    ff_vk_shader_update_desc_buffer(vkctx, exec, &pv->trellis_node_shd, 0, 1, 0,
                                     slice_score_buf, 0, slice_score_buf->size,
                                     VK_FORMAT_UNDEFINED);
     ff_vk_exec_bind_shader(vkctx, exec, &pv->trellis_node_shd);
@@ -577,9 +551,9 @@ static int encode_picture(AVCodecContext *avctx, FFVkExecContext *exec,
             .dstAccessMask = VK_ACCESS_2_SHADER_WRITE_BIT | VK_ACCESS_2_SHADER_READ_BIT,
             .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
             .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
-            .buffer = frame_size_buf->buf,
+            .buffer = slice_score_buf->buf,
             .offset = 0,
-            .size = frame_size_buf->size,
+            .size = slice_score_buf->size,
         },
         .bufferMemoryBarrierCount = 1,
     });
@@ -681,7 +655,6 @@ static int get_packet(AVCodecContext *avctx, FFVkExecContext *exec, AVPacket *pk
     uint8_t *picture_size_pos;
     int picture_idx;
     int frame_size, picture_size;
-    FFVkBuffer *frame_size_buf;
     VkMappedMemoryRange invalidate_data;
 
     /* Make sure encoding's done */
@@ -708,13 +681,7 @@ static int get_packet(AVCodecContext *avctx, FFVkExecContext *exec, AVPacket *pk
         FFVkBuffer *slice_sizes_buf = pd->slice_sizes_ref[picture_idx];
         const uint32_t *sizes;
 
-        frame_size_buf = pd->frame_size_ref[picture_idx];
-
         /* Invalidate slice sizes if needed */
-        if (!(frame_size_buf->flags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)) {
-            invalidate_data.memory = frame_size_buf->mem;
-            vk->InvalidateMappedMemoryRanges(vkctx->hwctx->act_dev, 1, &invalidate_data);
-        }
         if (!(slice_sizes_buf->flags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)) {
             invalidate_data.memory = slice_sizes_buf->mem;
             vk->InvalidateMappedMemoryRanges(vkctx->hwctx->act_dev, 1, &invalidate_data);
@@ -735,7 +702,7 @@ static int get_packet(AVCodecContext *avctx, FFVkExecContext *exec, AVPacket *pk
         av_assert1(picture_idx || buf - wrap_buf->mapped_mem == pv->payload_off);
 
         /* Calculate final size */
-        buf += *(int*)frame_size_buf->mapped_mem;
+        buf += sizes[ctx->slices_per_picture];
 
         /* Write picture size with header */
         picture_size = buf - (picture_size_pos - 1);
@@ -746,7 +713,6 @@ static int get_packet(AVCodecContext *avctx, FFVkExecContext *exec, AVPacket *pk
         av_refstruct_unref(&pd->slice_sizes_ref[picture_idx]);
         av_refstruct_unref(&pd->slice_data_ref[picture_idx]);
         av_refstruct_unref(&pd->slice_score_ref[picture_idx]);
-        av_refstruct_unref(&pd->frame_size_ref[picture_idx]);
     }
 
     /* Write frame size in header */
@@ -840,7 +806,6 @@ static av_cold int encode_close(AVCodecContext *avctx)
                 av_refstruct_unref(&pd->slice_sizes_ref[j]);
                 av_refstruct_unref(&pd->slice_data_ref[j]);
                 av_refstruct_unref(&pd->slice_score_ref[j]);
-                av_refstruct_unref(&pd->frame_size_ref[j]);
             }
             av_refstruct_unref(&pd->gathered_ref);
         }
@@ -853,7 +818,6 @@ static av_cold int encode_close(AVCodecContext *avctx)
     av_refstruct_pool_uninit(&pv->gathered_buf_pool);
     av_refstruct_pool_uninit(&pv->slice_data_buf_pool);
     av_refstruct_pool_uninit(&pv->slice_score_buf_pool);
-    av_refstruct_pool_uninit(&pv->frame_size_buf_pool);
 
     ff_vk_uninit(vkctx);
 

libavcodec/vulkan/prores_ks_encode_slice.comp.glsl

@@ -45,7 +45,6 @@ struct SliceScore {
     int total_bits[16];
     int total_score[16];
     int overquant;
-    int buf_start;
     int quant;
 };
 

libavcodec/vulkan/prores_ks_estimate_slice.comp.glsl

@@ -50,7 +50,6 @@ struct SliceScore {
     int total_bits[16];
     int total_score[16];
     int overquant;
-    int buf_start;
     int quant;
 };
 

libavcodec/vulkan/prores_ks_trellis_node.comp.glsl

@@ -22,13 +22,10 @@
 #extension GL_EXT_control_flow_attributes : require
 #extension GL_EXT_scalar_block_layout : require
 #extension GL_EXT_shader_explicit_arithmetic_types : require
-#extension GL_KHR_shader_subgroup_arithmetic : require
-#extension GL_EXT_maximal_reconvergence : require
 
 layout (local_size_x_id = 253, local_size_y_id = 254, local_size_z_id = 255) in;
 
 layout (constant_id = 0) const int slices_per_row = 1;
-layout (constant_id = 1) const int num_subgroups = 1;
 layout (constant_id = 2) const int num_planes = 0;
 layout (constant_id = 3) const int force_quant = 0;
 layout (constant_id = 4) const int min_quant = 0;
@@ -42,14 +39,10 @@ struct SliceScore {
     int total_bits[16];
     int total_score[16];
     int overquant;
-    int buf_start;
     int quant;
 };
 
-layout (set = 0, binding = 0, scalar) writeonly buffer FrameSize {
-    int frame_size;
-};
-layout (set = 0, binding = 1, scalar) buffer SliceScores {
+layout (set = 0, binding = 0, scalar) buffer SliceScores {
     SliceScore scores[];
 };
 
@@ -63,10 +56,6 @@ struct TrellisNode {
     int score;
 };
 
-shared int subgroup_sizes[num_subgroups];
-
-int slice_sizes[slices_per_row];
-
 TrellisNode nodes[(slices_per_row + 1) * TRELLIS_WIDTH];
 
 int find_slice_quant(int slice_x)
@@ -117,7 +106,7 @@ int find_slice_quant(int slice_x)
     return pq;
 }
 
-int find_slice_row_quants()
+void find_slice_row_quants()
 {
     for (int i = min_quant; i < max_quant + 2; i++) {
         nodes[i].prev_node = -1;
@@ -130,67 +119,21 @@ int find_slice_row_quants()
         q = find_slice_quant(slice_x);
     }
 
-    int slice_hdr_size = 2 * num_planes;
-    int slice_row_size = slice_hdr_size * slices_per_row;
     int y = int(gl_LocalInvocationID.x);
     for (int x = slices_per_row - 1; x >= 0; x--) {
         int slice = x + y * slices_per_row;
-        int quant = nodes[q].quant;
-        int q_idx = min(quant, max_quant + 1);
-        ivec4 bits = scores[slice].bits[q_idx];
-        slice_sizes[x] = (bits.x + bits.y + bits.z + bits.w) / 8;
-        slice_row_size += slice_sizes[x];
-        scores[slice].quant = quant;
+        scores[slice].quant = nodes[q].quant;
         q = nodes[q].prev_node;
     }
-
-    return slice_row_size;
 }
 
-int force_slice_row_quants()
+void main()
 {
-    int slice_hdr_size = 2 * num_planes;
-    int slice_row_size = slice_hdr_size * slices_per_row;
-    int y = int(gl_LocalInvocationID.x);
-    for (int x = slices_per_row - 1; x >= 0; x--) {
-        int slice = x + y * slices_per_row;
-        ivec4 bits = scores[slice].bits[0];
-        slice_sizes[x] = (bits.x + bits.y + bits.z + bits.w) / 8;
-        slice_row_size += slice_sizes[x];
-        scores[slice].quant = force_quant;
+    if (force_quant == 0) {
+        find_slice_row_quants();
+    } else {
+        int y = int(gl_LocalInvocationID.x);
+        for (int x = 0; x < slices_per_row; x++)
+            scores[x + y * slices_per_row].quant = force_quant;
     }
-
-    return slice_row_size;
-}
-
-void main() [[maximally_reconverges]]
-{
-    int slice_row_size;
-    if (force_quant == 0)
-        slice_row_size = find_slice_row_quants();
-    else
-        slice_row_size = force_slice_row_quants();
-
-    int subgroup_sum = subgroupAdd(slice_row_size);
-    if (subgroupElect())
-        subgroup_sizes[gl_SubgroupID] = subgroup_sum;
-    barrier();
-
-    int buf_start = subgroupExclusiveAdd(slice_row_size);
-    [[unroll]] for (int i = 0; i < num_subgroups; ++i) {
-        if (i >= gl_SubgroupID)
-            break;
-        buf_start += subgroup_sizes[i];
-    }
-
-    int slice_hdr_size = 2 * num_planes;
-    int y = int(gl_LocalInvocationID.x);
-    [[unroll]] for (int x = 0; x < slices_per_row; ++x) {
-        int slice = x + y * slices_per_row;
-        scores[slice].buf_start = buf_start;
-        buf_start += slice_hdr_size + slice_sizes[x];
-    }
-
-    if (y == gl_WorkGroupSize.x - 1)
-        frame_size = buf_start;
 }