History

ffv1enc_vulkan: implement RCT search for level >= 4

Lynne committed 1 year ago

Commit 75764101 parent 01566804 files · +346 −5Browse files

Showing the changes to libavcodec (4 of 4 files). Show all

Unified · Split

4 files
  1. libavcodec/ffv1enc_vulkan.c +202 −2
  2. libavcodec/vulkan/Makefile +1 −1
  3. libavcodec/vulkan/ffv1_enc_setup.comp +4 −2
  4. libavcodec/vulkan/ffv1_rct_search.comp +139 −0
libavcodec/ffv1enc_vulkan.c+202−2
@@ -74,6 +74,7 @@ typedef struct VulkanEncodeFFv1Context { · 73 unmodified lines
7474 size_t max_heap_size;
7575
7676 FFVulkanShader setup;
77+ FFVulkanShader rct_search;
7778 FFVulkanShader reset;
7879 FFVulkanShader enc;
7980
@@ -101,6 +102,7 @@ typedef struct VulkanEncodeFFv1Context { · 21 unmodified lines
101102 int num_h_slices;
102103 int num_v_slices;
103104 int force_pcm;
105+ int optimize_rct;
104106
105107 int is_rgb;
106108 int ppi;
@@ -112,6 +114,7 @@ extern const char *ff_source_rangecoder_comp; · 5 unmodified lines
112114 extern const char *ff_source_ffv1_vlc_comp;
113115 extern const char *ff_source_ffv1_common_comp;
114116 extern const char *ff_source_ffv1_reset_comp;
117+extern const char *ff_source_ffv1_rct_search_comp;
115118 extern const char *ff_source_ffv1_enc_setup_comp;
116119 extern const char *ff_source_ffv1_enc_comp;
117120
@@ -147,7 +150,8 @@ typedef struct FFv1VkParameters { · 29 unmodified lines
147150 uint8_t ec;
148151 uint8_t ppi;
149152 uint8_t chunks;
150− uint8_t padding[4];
153+ uint8_t rct_search;
154+ uint8_t padding[3];
151155 } FFv1VkParameters;
152156
153157 static void add_push_data(FFVulkanShader *shd)
@@ -184,12 +188,76 @@ static void add_push_data(FFVulkanShader *shd) · 30 unmodified lines
184188 GLSLC(1, uint8_t ec; );
185189 GLSLC(1, uint8_t ppi; );
186190 GLSLC(1, uint8_t chunks; );
187− GLSLC(1, uint8_t padding[4]; );
191+ GLSLC(1, uint8_t rct_search; );
192+ GLSLC(1, uint8_t padding[3]; );
188193 GLSLC(0, }; );
189194 ff_vk_shader_add_push_const(shd, 0, sizeof(FFv1VkParameters),
190195 VK_SHADER_STAGE_COMPUTE_BIT);
191196 }
192197
198+typedef struct FFv1VkRCTSearchParameters {
199+ int fmt_lut[4];
200+ int rct_offset;
201+ uint8_t planar_rgb;
202+ uint8_t transparency;
203+ uint8_t key_frame;
204+ uint8_t force_pcm;
205+ uint8_t version;
206+ uint8_t micro_version;
207+ uint8_t padding[2];
208+} FFv1VkRCTSearchParameters;
209+
210+static int run_rct_search(AVCodecContext *avctx, FFVkExecContext *exec,
211+ AVFrame *enc_in, VkImageView *enc_in_views,
212+ FFVkBuffer *slice_data_buf, uint32_t slice_data_size)
213+{
214+ VulkanEncodeFFv1Context *fv = avctx->priv_data;
215+ FFV1Context *f = &fv->ctx;
216+ FFVulkanFunctions *vk = &fv->s.vkfn;
217+ AVHWFramesContext *src_hwfc = (AVHWFramesContext *)enc_in->hw_frames_ctx->data;
218+ FFv1VkRCTSearchParameters pd;
219+
220+ /* Update descriptors */
221+ ff_vk_shader_update_desc_buffer(&fv->s, exec, &fv->rct_search,
222+ 0, 0, 0,
223+ slice_data_buf,
224+ 0, slice_data_size*f->slice_count,
225+ VK_FORMAT_UNDEFINED);
226+ ff_vk_shader_update_img_array(&fv->s, exec, &fv->rct_search,
227+ enc_in, enc_in_views,
228+ 0, 1,
229+ VK_IMAGE_LAYOUT_GENERAL,
230+ VK_NULL_HANDLE);
231+
232+ ff_vk_exec_bind_shader(&fv->s, exec, &fv->rct_search);
233+
234+ pd = (FFv1VkRCTSearchParameters) {
235+ .rct_offset = 1 << f->bits_per_raw_sample,
236+ .planar_rgb = ff_vk_mt_is_np_rgb(src_hwfc->sw_format) &&
237+ (ff_vk_count_images((AVVkFrame *)enc_in->data[0]) > 1),
238+ .transparency = f->transparency,
239+ .key_frame = f->key_frame,
240+ .force_pcm = fv->force_pcm,
241+ .version = f->version,
242+ .micro_version = f->micro_version,
243+ };
244+
245+ if (avctx->sw_pix_fmt == AV_PIX_FMT_GBRP10 ||
246+ avctx->sw_pix_fmt == AV_PIX_FMT_GBRP12 ||
247+ avctx->sw_pix_fmt == AV_PIX_FMT_GBRP14)
248+ memcpy(pd.fmt_lut, (int [4]) { 2, 1, 0, 3 }, 4*sizeof(int));
249+ else
250+ ff_vk_set_perm(avctx->sw_pix_fmt, pd.fmt_lut, 1);
251+
252+ ff_vk_shader_update_push_const(&fv->s, exec, &fv->rct_search,
253+ VK_SHADER_STAGE_COMPUTE_BIT,
254+ 0, sizeof(pd), &pd);
255+
256+ vk->CmdDispatch(exec->buf, fv->ctx.num_h_slices, fv->ctx.num_v_slices, 1);
257+
258+ return 0;
259+}
260+
193261 static int vulkan_encode_ffv1_submit_frame(AVCodecContext *avctx,
194262 FFVkExecContext *exec,
195263 const AVFrame *pict)
@@ -366,6 +434,25 @@ static int vulkan_encode_ffv1_submit_frame(AVCodecContext *avctx, · 170 unmodified lines
366434 };
367435 }
368436
437+ if (fv->optimize_rct) {
438+ RET(run_rct_search(avctx, exec,
439+ src, src_views,
440+ slice_data_buf, slice_data_size));
441+
442+ buf_bar[nb_buf_bar++] = (VkBufferMemoryBarrier2) {
443+ .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
444+ .srcStageMask = slice_data_buf->stage,
445+ .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
446+ .srcAccessMask = slice_data_buf->access,
447+ .dstAccessMask = VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
448+ .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
449+ .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
450+ .buffer = slice_data_buf->buf,
451+ .size = slice_data_size*f->slice_count,
452+ .offset = 0,
453+ };
454+ }
455+
369456 vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) {
370457 .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
371458 .pImageMemoryBarriers = img_bar,
@@ -412,6 +499,7 @@ static int vulkan_encode_ffv1_submit_frame(AVCodecContext *avctx, · 40 unmodified lines
412499 .ec = f->ec,
413500 .ppi = fv->ppi,
414501 .chunks = fv->chunks,
502+ .rct_search = fv->optimize_rct,
415503 };
416504
417505 /* For some reason the C FFv1 encoder/decoder treats these differently */
@@ -920,6 +1008,103 @@ static void define_shared_code(AVCodecContext *avctx, FFVulkanShader *shd) · 502 unmodified lines
9201008 GLSLD(ff_source_ffv1_common_comp);
9211009 }
9221010
1011+static int init_rct_search_shader(AVCodecContext *avctx, FFVkSPIRVCompiler *spv)
1012+{
1013+ int err;
1014+ VulkanEncodeFFv1Context *fv = avctx->priv_data;
1015+ FFVulkanShader *shd = &fv->rct_search;
1016+ FFVulkanDescriptorSetBinding *desc_set;
1017+
1018+ uint8_t *spv_data;
1019+ size_t spv_len;
1020+ void *spv_opaque = NULL;
1021+
1022+ RET(ff_vk_shader_init(&fv->s, shd, "ffv1_rct_search",
1023+ VK_SHADER_STAGE_COMPUTE_BIT,
1024+ (const char *[]) { "GL_EXT_buffer_reference",
1025+ "GL_EXT_buffer_reference2",
1026+ "GL_EXT_null_initializer" }, 3,
1027+ 32, 32, 1,
1028+ 0));
1029+
1030+ /* Common codec header */
1031+ GLSLD(ff_source_common_comp);
1032+
1033+ GLSLC(0, layout(push_constant, scalar) uniform pushConstants { );
1034+ GLSLC(1, ivec4 fmt_lut; );
1035+ GLSLC(1, int rct_offset; );
1036+ GLSLC(1, uint8_t planar_rgb; );
1037+ GLSLC(1, uint8_t transparency; );
1038+ GLSLC(1, uint8_t key_frame; );
1039+ GLSLC(1, uint8_t force_pcm; );
1040+ GLSLC(1, uint8_t version; );
1041+ GLSLC(1, uint8_t micro_version; );
1042+ GLSLC(1, uint8_t padding[3]; );
1043+ GLSLC(0, }; );
1044+ ff_vk_shader_add_push_const(shd, 0, sizeof(FFv1VkResetParameters),
1045+ VK_SHADER_STAGE_COMPUTE_BIT);
1046+
1047+ av_bprintf(&shd->src, "#define MAX_QUANT_TABLES %i\n", MAX_QUANT_TABLES);
1048+ av_bprintf(&shd->src, "#define MAX_CONTEXT_INPUTS %i\n", MAX_CONTEXT_INPUTS);
1049+ av_bprintf(&shd->src, "#define MAX_QUANT_TABLE_SIZE %i\n", MAX_QUANT_TABLE_SIZE);
1050+
1051+ /* Never used */
1052+ desc_set = (FFVulkanDescriptorSetBinding []) {
1053+ {
1054+ .name = "rangecoder_static_buf",
1055+ .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
1056+ .stages = VK_SHADER_STAGE_COMPUTE_BIT,
1057+ .mem_layout = "scalar",
1058+ .buf_content = "uint8_t zero_one_state[512];",
1059+ },
1060+ {
1061+ .name = "quant_buf",
1062+ .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
1063+ .stages = VK_SHADER_STAGE_COMPUTE_BIT,
1064+ .mem_layout = "scalar",
1065+ .buf_content = "int16_t quant_table[MAX_QUANT_TABLES]"
1066+ "[MAX_CONTEXT_INPUTS][MAX_QUANT_TABLE_SIZE];",
1067+ },
1068+ };
1069+ RET(ff_vk_shader_add_descriptor_set(&fv->s, shd, desc_set, 2, 1, 1));
1070+
1071+ define_shared_code(avctx, shd);
1072+
1073+ desc_set = (FFVulkanDescriptorSetBinding []) {
1074+ {
1075+ .name = "slice_data_buf",
1076+ .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
1077+ .stages = VK_SHADER_STAGE_COMPUTE_BIT,
1078+ .buf_content = "SliceContext slice_ctx[1024];",
1079+ },
1080+ {
1081+ .name = "src",
1082+ .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
1083+ .dimensions = 2,
1084+ .mem_layout = ff_vk_shader_rep_fmt(fv->s.frames->sw_format,
1085+ fv->rep_fmt),
1086+ .elems = av_pix_fmt_count_planes(fv->s.frames->sw_format),
1087+ .mem_quali = "readonly",
1088+ .stages = VK_SHADER_STAGE_COMPUTE_BIT,
1089+ },
1090+ };
1091+ RET(ff_vk_shader_add_descriptor_set(&fv->s, shd, desc_set, 2, 0, 0));
1092+
1093+ GLSLD(ff_source_ffv1_rct_search_comp);
1094+
1095+ RET(spv->compile_shader(&fv->s, spv, shd, &spv_data, &spv_len, "main",
1096+ &spv_opaque));
1097+ RET(ff_vk_shader_link(&fv->s, shd, spv_data, spv_len, "main"));
1098+
1099+ RET(ff_vk_shader_register_exec(&fv->s, &fv->exec_pool, shd));
1100+
1101+fail:
1102+ if (spv_opaque)
1103+ spv->free_shader(spv, &spv_opaque);
1104+
1105+ return err;
1106+}
1107+
9231108 static int init_setup_shader(AVCodecContext *avctx, FFVkSPIRVCompiler *spv)
9241109 {
9251110 int err;
@@ -1417,6 +1602,17 @@ static av_cold int vulkan_encode_ffv1_init(AVCodecContext *avctx) · 491 unmodified lines
14171602 if (!fv->is_rgb && f->bits_per_raw_sample > 8)
14181603 fv->rep_fmt = FF_VK_REP_INT;
14191604
1605+ /* Init rct search shader */
1606+ fv->optimize_rct = fv->is_rgb && f->version >= 4 &&
1607+ !fv->force_pcm && fv->optimize_rct;
1608+ if (fv->optimize_rct) {
1609+ err = init_rct_search_shader(avctx, spv);
1610+ if (err < 0) {
1611+ spv->uninit(&spv);
1612+ return err;
1613+ }
1614+ }
1615+
14201616 /* Init setup shader */
14211617 err = init_setup_shader(avctx, spv);
14221618 if (err < 0) {
@@ -1528,6 +1724,7 @@ static av_cold int vulkan_encode_ffv1_close(AVCodecContext *avctx) · 105 unmodified lines
15281724 ff_vk_shader_free(&fv->s, &fv->enc);
15291725 ff_vk_shader_free(&fv->s, &fv->reset);
15301726 ff_vk_shader_free(&fv->s, &fv->setup);
1727+ ff_vk_shader_free(&fv->s, &fv->rct_search);
15311728
15321729 if (fv->exec_ctx_info) {
15331730 for (int i = 0; i < fv->async_depth; i++) {
@@ -1591,6 +1788,9 @@ static const AVOption vulkan_encode_ffv1_options[] = { · 57 unmodified lines
15911788 { "force_pcm", "Code all slices with no prediction", OFFSET(force_pcm), AV_OPT_TYPE_BOOL,
15921789 { .i64 = 0 }, 0, 1, VE },
15931790
1791+ { "rct_search", "Run a search for RCT parameters (level 4 only)", OFFSET(optimize_rct), AV_OPT_TYPE_BOOL,
1792+ { .i64 = 1 }, 0, 1, VE },
1793+
15941794 { "async_depth", "Internal parallelization depth", OFFSET(async_depth), AV_OPT_TYPE_INT,
15951795 { .i64 = 1 }, 1, INT_MAX, VE },
15961796
libavcodec/vulkan/Makefile+1−1
@@ -7,7 +7,7 @@ OBJS-$(CONFIG_FFV1_VULKAN_ENCODER) += vulkan/common.o \ · 6 unmodified lines
77 vulkan/rangecoder.o vulkan/ffv1_vlc.o \
88 vulkan/ffv1_common.o vulkan/ffv1_reset.o \
99 vulkan/ffv1_enc_rct.o vulkan/ffv1_enc_setup.o \
10− vulkan/ffv1_enc.o
10+ vulkan/ffv1_rct_search.o vulkan/ffv1_enc.o
1111
1212 OBJS-$(CONFIG_FFV1_VULKAN_HWACCEL) += vulkan/common.o \
1313 vulkan/rangecoder.o vulkan/ffv1_vlc.o \
libavcodec/vulkan/ffv1_enc_setup.comp+4−2
@@ -22,7 +22,7 @@ · 21 unmodified lines
2222
2323 uint8_t state[CONTEXT_SIZE];
2424
25−void init_slice(out SliceContext sc, const uint slice_idx)
25+void init_slice(inout SliceContext sc, const uint slice_idx)
2626 {
2727 /* Set coordinates */
2828 uvec2 img_size = imageSize(src[0]);
@@ -37,11 +37,13 @@ void init_slice(out SliceContext sc, const uint slice_idx) · 8 unmodified lines
3737
3838 sc.slice_pos = ivec2(sxs, sys);
3939 sc.slice_dim = ivec2(sxe - sxs, sye - sys);
40− sc.slice_rct_coef = ivec2(1, 1);
4140 sc.slice_coding_mode = int(force_pcm == 1);
4241 sc.slice_reset_contexts = sc.slice_coding_mode == 1;
4342 sc.quant_table_idx = u8vec3(context_model);
4443
44+ if ((rct_search == 0) || (sc.slice_coding_mode == 1))
45+ sc.slice_rct_coef = ivec2(1, 1);
46+
4547 rac_init(sc.c,
4648 OFFBUF(u8buf, out_data, slice_idx * slice_size_max),
4749 slice_size_max);
libavcodec/vulkan/ffv1_rct_search.comp+139−0
@@ -0,0 +1,139 @@
1+/*
2+ * FFv1 codec
3+ *
4+ * Copyright (c) 2024 Lynne <dev@lynne.ee>
5+ *
6+ * This file is part of FFmpeg.
7+ *
8+ * FFmpeg is free software; you can redistribute it and/or
9+ * modify it under the terms of the GNU Lesser General Public
10+ * License as published by the Free Software Foundation; either
11+ * version 2.1 of the License, or (at your option) any later version.
12+ *
13+ * FFmpeg is distributed in the hope that it will be useful,
14+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
15+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
16+ * Lesser General Public License for more details.
17+ *
18+ * You should have received a copy of the GNU Lesser General Public
19+ * License along with FFmpeg; if not, write to the Free Software
20+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
21+ */
22+
23+ivec3 load_components(ivec2 pos)
24+{
25+ ivec3 pix = ivec3(imageLoad(src[0], pos));
26+ if (planar_rgb != 0) {
27+ for (int i = 1; i < 3; i++)
28+ pix[i] = int(imageLoad(src[i], pos)[0]);
29+ }
30+
31+ return ivec3(pix[fmt_lut[0]], pix[fmt_lut[1]], pix[fmt_lut[2]]);
32+}
33+
34+#define NUM_CHECKS 15
35+const ivec2 rct_y_coeff[NUM_CHECKS] = {
36+ ivec2(0, 0), // 4G
37+
38+ ivec2(0, 1), // 3G + B
39+ ivec2(1, 0), // R + 3G
40+ ivec2(1, 1), // R + 2G + B
41+
42+ ivec2(0, 2), // 2G + 2B
43+ ivec2(2, 0), // 2R + 2G
44+ ivec2(2, 2), // 2R + 2B
45+
46+ ivec2(0, 3), // 1G + 3B
47+ ivec2(3, 0), // 3R + 1G
48+
49+ ivec2(0, 4), // 4B
50+ ivec2(4, 0), // 4R
51+
52+ ivec2(1, 2), // R + G + 2B
53+ ivec2(2, 1), // 2R + G + B
54+
55+ ivec2(3, 1), // 3R + B
56+ ivec2(1, 3), // R + 3B
57+};
58+
59+shared ivec3 pix_buf[gl_WorkGroupSize.x + 1][gl_WorkGroupSize.y + 1] = { };
60+
61+ivec3 transform_sample(ivec3 pix, ivec2 rct_coef)
62+{
63+ pix.b -= pix.g;
64+ pix.r -= pix.g;
65+ pix.g += (pix.r*rct_coef.x + pix.b*rct_coef.y) >> 2;
66+ pix.b += rct_offset;
67+ pix.r += rct_offset;
68+ return pix;
69+}
70+
71+uint get_dist(ivec3 cur)
72+{
73+ ivec3 LL = pix_buf[gl_LocalInvocationID.x + 0][gl_LocalInvocationID.y + 1];
74+ ivec3 TL = pix_buf[gl_LocalInvocationID.x + 0][gl_LocalInvocationID.y + 0];
75+ ivec3 TT = pix_buf[gl_LocalInvocationID.x + 1][gl_LocalInvocationID.y + 0];
76+
77+ ivec3 pred = ivec3(predict(LL.r, ivec2(TL.r, TT.r)),
78+ predict(LL.g, ivec2(TL.g, TT.g)),
79+ predict(LL.b, ivec2(TL.b, TT.b)));
80+
81+ uvec3 c = abs(pred - cur);
82+ return mid_pred(c.r, c.g, c.b);
83+}
84+
85+shared uint score_cols[gl_WorkGroupSize.y] = { };
86+shared uint score_mode[16] = { };
87+
88+void process(ivec2 pos)
89+{
90+ ivec3 pix = load_components(pos);
91+
92+ for (int i = 0; i < NUM_CHECKS; i++) {
93+ ivec3 tx_pix = transform_sample(pix, rct_y_coeff[i]);
94+ pix_buf[gl_LocalInvocationID.x + 1][gl_LocalInvocationID.y + 1] = tx_pix;
95+ memoryBarrierShared();
96+
97+ uint dist = get_dist(tx_pix);
98+ atomicAdd(score_mode[i], dist);
99+ }
100+}
101+
102+void coeff_search(inout SliceContext sc)
103+{
104+ uvec2 img_size = imageSize(src[0]);
105+ uint sxs = slice_coord(img_size.x, gl_WorkGroupID.x + 0,
106+ gl_NumWorkGroups.x, 0);
107+ uint sxe = slice_coord(img_size.x, gl_WorkGroupID.x + 1,
108+ gl_NumWorkGroups.x, 0);
109+ uint sys = slice_coord(img_size.y, gl_WorkGroupID.y + 0,
110+ gl_NumWorkGroups.y, 0);
111+ uint sye = slice_coord(img_size.y, gl_WorkGroupID.y + 1,
112+ gl_NumWorkGroups.y, 0);
113+
114+ for (uint y = sys + gl_LocalInvocationID.y; y < sye; y += gl_WorkGroupSize.y) {
115+ for (uint x = sxs + gl_LocalInvocationID.x; x < sxe; x += gl_WorkGroupSize.x) {
116+ process(ivec2(x, y));
117+ }
118+ }
119+
120+ if (gl_LocalInvocationID.x == 0 && gl_LocalInvocationID.y == 0) {
121+ uint min_score = 0xFFFFFFFF;
122+ uint min_idx = 3;
123+ for (int i = 0; i < NUM_CHECKS; i++) {
124+ if (score_mode[i] < min_score) {
125+ min_score = score_mode[i];
126+ min_idx = i;
127+ }
128+ }
129+ sc.slice_rct_coef = rct_y_coeff[min_idx];
130+ }
131+}
132+
133+void main(void)
134+{
135+ if (force_pcm == 1)
136+ return;
137+ const uint slice_idx = gl_WorkGroupID.y*gl_NumWorkGroups.x + gl_WorkGroupID.x;
138+ coeff_search(slice_ctx[slice_idx]);
139+}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t