History

lavfi: add nlmeans_vulkan filter

Lynne committed 3 years ago

Commit 160a4151 parent dfff3876 files · +1278 −0Browse files

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

Unified · Split

4 files
  1. libavfilter/Makefile +2 −0
  2. libavfilter/allfilters.c +1 −0
  3. libavfilter/vf_nlmeans_vulkan.c +1122 −0
  4. libavfilter/vulkan/prefix_sum.comp +151 −0
libavfilter/Makefile+2−0
@@ -390,6 +390,8 @@ OBJS-$(CONFIG_MULTIPLY_FILTER) += vf_multiply.o · 389 unmodified lines
390OBJS-$(CONFIG_NEGATE_FILTER) += vf_negate.o390OBJS-$(CONFIG_NEGATE_FILTER) += vf_negate.o
391OBJS-$(CONFIG_NLMEANS_FILTER) += vf_nlmeans.o391OBJS-$(CONFIG_NLMEANS_FILTER) += vf_nlmeans.o
392OBJS-$(CONFIG_NLMEANS_OPENCL_FILTER) += vf_nlmeans_opencl.o opencl.o opencl/nlmeans.o392OBJS-$(CONFIG_NLMEANS_OPENCL_FILTER) += vf_nlmeans_opencl.o opencl.o opencl/nlmeans.o
393OBJS-$(CONFIG_NLMEANS_VULKAN_FILTER) += vf_nlmeans_vulkan.o vulkan.o vulkan_filter.o \
394 vulkan/prefix_sum.o
393OBJS-$(CONFIG_NNEDI_FILTER) += vf_nnedi.o395OBJS-$(CONFIG_NNEDI_FILTER) += vf_nnedi.o
394OBJS-$(CONFIG_NOFORMAT_FILTER) += vf_format.o396OBJS-$(CONFIG_NOFORMAT_FILTER) += vf_format.o
395OBJS-$(CONFIG_NOISE_FILTER) += vf_noise.o397OBJS-$(CONFIG_NOISE_FILTER) += vf_noise.o
libavfilter/allfilters.c+1−0
@@ -368,6 +368,7 @@ extern const AVFilter ff_vf_multiply; · 367 unmodified lines
368extern const AVFilter ff_vf_negate;368extern const AVFilter ff_vf_negate;
369extern const AVFilter ff_vf_nlmeans;369extern const AVFilter ff_vf_nlmeans;
370extern const AVFilter ff_vf_nlmeans_opencl;370extern const AVFilter ff_vf_nlmeans_opencl;
371extern const AVFilter ff_vf_nlmeans_vulkan;
371extern const AVFilter ff_vf_nnedi;372extern const AVFilter ff_vf_nnedi;
372extern const AVFilter ff_vf_noformat;373extern const AVFilter ff_vf_noformat;
373extern const AVFilter ff_vf_noise;374extern const AVFilter ff_vf_noise;
libavfilter/vf_nlmeans_vulkan.c+1122−0
@@ -0,0 +1,1122 @@
1/*
2 * Copyright (c) Lynne
3 *
4 * This file is part of FFmpeg.
5 *
6 * FFmpeg is free software; you can redistribute it and/or
7 * modify it under the terms of the GNU Lesser General Public
8 * License as published by the Free Software Foundation; either
9 * version 2.1 of the License, or (at your option) any later version.
10 *
11 * FFmpeg is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
14 * Lesser General Public License for more details.
15 *
16 * You should have received a copy of the GNU Lesser General Public
17 * License along with FFmpeg; if not, write to the Free Software
18 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19 */
20
21#include "libavutil/random_seed.h"
22#include "libavutil/opt.h"
23#include "vulkan_filter.h"
24#include "vulkan_spirv.h"
25#include "internal.h"
26
27#define TYPE_NAME "vec4"
28#define TYPE_ELEMS 4
29#define TYPE_SIZE (TYPE_ELEMS*4)
30
31typedef struct NLMeansVulkanContext {
32 FFVulkanContext vkctx;
33
34 int initialized;
35 FFVkExecPool e;
36 FFVkQueueFamilyCtx qf;
37 VkSampler sampler;
38
39 AVBufferPool *integral_buf_pool;
40 AVBufferPool *state_buf_pool;
41 AVBufferPool *ws_buf_pool;
42
43 int pl_weights_rows;
44 FFVulkanPipeline pl_weights;
45 FFVkSPIRVShader shd_weights;
46
47 FFVulkanPipeline pl_denoise;
48 FFVkSPIRVShader shd_denoise;
49
50 int *xoffsets;
51 int *yoffsets;
52 int nb_offsets;
53 float strength[4];
54 int patch[4];
55
56 struct nlmeans_opts {
57 int r;
58 double s;
59 double sc[4];
60 int p;
61 int pc[4];
62 int t;
63 } opts;
64} NLMeansVulkanContext;
65
66extern const char *ff_source_prefix_sum_comp;
67
68static void insert_first(FFVkSPIRVShader *shd, int r, int horiz, int plane, int comp)
69{
70 GLSLF(2, s1 = texture(input_img[%i], ivec2(x + %i, y + %i))[%i];
71 ,plane, horiz ? r : 0, !horiz ? r : 0, comp);
72
73 if (TYPE_ELEMS == 4) {
74 GLSLF(2, s2[0] = texture(input_img[%i], ivec2(x + %i + xoffs[0], y + %i + yoffs[0]))[%i];
75 ,plane, horiz ? r : 0, !horiz ? r : 0, comp);
76 GLSLF(2, s2[1] = texture(input_img[%i], ivec2(x + %i + xoffs[1], y + %i + yoffs[1]))[%i];
77 ,plane, horiz ? r : 0, !horiz ? r : 0, comp);
78 GLSLF(2, s2[2] = texture(input_img[%i], ivec2(x + %i + xoffs[2], y + %i + yoffs[2]))[%i];
79 ,plane, horiz ? r : 0, !horiz ? r : 0, comp);
80 GLSLF(2, s2[3] = texture(input_img[%i], ivec2(x + %i + xoffs[3], y + %i + yoffs[3]))[%i];
81 ,plane, horiz ? r : 0, !horiz ? r : 0, comp);
82 } else {
83 for (int i = 0; i < 16; i++) {
84 GLSLF(2, s2[%i][%i] = texture(input_img[%i], ivec2(x + %i + xoffs[%i], y + %i + yoffs[%i]))[%i];
85 ,i / 4, i % 4, plane, horiz ? r : 0, i, !horiz ? r : 0, i, comp);
86 }
87 }
88
89 GLSLC(2, s2 = (s1 - s2) * (s1 - s2); );
90}
91
92static void insert_horizontal_pass(FFVkSPIRVShader *shd, int nb_rows, int first, int plane, int comp)
93{
94 GLSLF(1, x = int(gl_GlobalInvocationID.x) * %i; ,nb_rows);
95 if (!first) {
96 GLSLC(1, controlBarrier(gl_ScopeWorkgroup, gl_ScopeWorkgroup,
97 gl_StorageSemanticsBuffer,
98 gl_SemanticsAcquireRelease |
99 gl_SemanticsMakeAvailable |
100 gl_SemanticsMakeVisible); );
101 }
102 GLSLC(1, for (y = 0; y < height[0]; y++) { );
103 GLSLC(2, offset = uint64_t(int_stride)*y*T_ALIGN; );
104 GLSLC(2, dst = DataBuffer(uint64_t(integral_data) + offset); );
105 GLSLC(0, );
106 if (first) {
107 for (int r = 0; r < nb_rows; r++) {
108 insert_first(shd, r, 1, plane, comp);
109 GLSLF(2, dst.v[x + %i] = s2; ,r);
110 GLSLC(0, );
111 }
112 }
113 GLSLC(2, barrier(); );
114 GLSLC(2, prefix_sum(dst, 1, dst, 1); );
115 GLSLC(1, } );
116 GLSLC(0, );
117}
118
119static void insert_vertical_pass(FFVkSPIRVShader *shd, int nb_rows, int first, int plane, int comp)
120{
121 GLSLF(1, y = int(gl_GlobalInvocationID.x) * %i; ,nb_rows);
122 if (!first) {
123 GLSLC(1, controlBarrier(gl_ScopeWorkgroup, gl_ScopeWorkgroup,
124 gl_StorageSemanticsBuffer,
125 gl_SemanticsAcquireRelease |
126 gl_SemanticsMakeAvailable |
127 gl_SemanticsMakeVisible); );
128 }
129 GLSLC(1, for (x = 0; x < width[0]; x++) { );
130 GLSLC(2, dst = DataBuffer(uint64_t(integral_data) + x*T_ALIGN); );
131
132 for (int r = 0; r < nb_rows; r++) {
133 if (first) {
134 insert_first(shd, r, 0, plane, comp);
135 GLSLF(2, integral_data.v[(y + %i)*int_stride + x] = s2; ,r);
136 GLSLC(0, );
137 }
138 }
139
140 GLSLC(2, barrier(); );
141 GLSLC(2, prefix_sum(dst, int_stride, dst, int_stride); );
142 GLSLC(1, } );
143 GLSLC(0, );
144}
145
146static void insert_weights_pass(FFVkSPIRVShader *shd, int nb_rows, int vert,
147 int t, int dst_comp, int plane, int comp)
148{
149 GLSLF(1, p = patch_size[%i]; ,dst_comp);
150 GLSLC(0, );
151 GLSLC(1, controlBarrier(gl_ScopeWorkgroup, gl_ScopeWorkgroup,
152 gl_StorageSemanticsBuffer,
153 gl_SemanticsAcquireRelease |
154 gl_SemanticsMakeAvailable |
155 gl_SemanticsMakeVisible); );
156 GLSLC(1, barrier(); );
157 if (!vert) {
158 GLSLC(1, for (y = 0; y < height[0]; y++) { );
159 GLSLF(2, if (gl_GlobalInvocationID.x*%i >= width[%i]) ,nb_rows, plane);
160 GLSLC(3, break; );
161 GLSLF(2, for (r = 0; r < %i; r++) { ,nb_rows);
162 GLSLF(3, x = int(gl_GlobalInvocationID.x) * %i + r; ,nb_rows);
163 } else {
164 GLSLC(1, for (x = 0; x < width[0]; x++) { );
165 GLSLF(2, if (gl_GlobalInvocationID.x*%i >= height[%i]) ,nb_rows, plane);
166 GLSLC(3, break; );
167 GLSLF(2, for (r = 0; r < %i; r++) { ,nb_rows);
168 GLSLF(3, y = int(gl_GlobalInvocationID.x) * %i + r; ,nb_rows);
169 }
170 GLSLC(0, );
171 GLSLC(3, a = DTYPE(0); );
172 GLSLC(3, b = DTYPE(0); );
173 GLSLC(3, c = DTYPE(0); );
174 GLSLC(3, d = DTYPE(0); );
175 GLSLC(0, );
176 GLSLC(3, lt = ((x - p) < 0) || ((y - p) < 0); );
177 GLSLC(0, );
178 if (TYPE_ELEMS == 4) {
179 GLSLF(3, src[0] = texture(input_img[%i], ivec2(x + xoffs[0], y + yoffs[0]))[%i]; ,plane, comp);
180 GLSLF(3, src[1] = texture(input_img[%i], ivec2(x + xoffs[1], y + yoffs[1]))[%i]; ,plane, comp);
181 GLSLF(3, src[2] = texture(input_img[%i], ivec2(x + xoffs[2], y + yoffs[2]))[%i]; ,plane, comp);
182 GLSLF(3, src[3] = texture(input_img[%i], ivec2(x + xoffs[3], y + yoffs[3]))[%i]; ,plane, comp);
183 } else {
184 for (int i = 0; i < 16; i++)
185 GLSLF(3, src[%i][%i] = texture(input_img[%i], ivec2(x + xoffs[%i], y + yoffs[%i]))[%i];
186 ,i / 4, i % 4, plane, i, i, comp);
187
188 }
189 GLSLC(0, );
190 GLSLC(3, if (lt == false) { );
191 GLSLC(4, a = integral_data.v[(y - p)*int_stride + x - p]; );
192 GLSLC(4, c = integral_data.v[(y - p)*int_stride + x + p]; );
193 GLSLC(4, b = integral_data.v[(y + p)*int_stride + x - p]; );
194 GLSLC(4, d = integral_data.v[(y + p)*int_stride + x + p]; );
195 GLSLC(3, } );
196 GLSLC(0, );
197 GLSLC(3, patch_diff = d + a - b - c; );
198 if (TYPE_ELEMS == 4) {
199 GLSLF(3, w = exp(patch_diff * strength[%i]); ,dst_comp);
200 GLSLC(3, w_sum = w[0] + w[1] + w[2] + w[3]; );
201 GLSLC(3, sum = dot(w, src*255); );
202 } else {
203 for (int i = 0; i < 4; i++)
204 GLSLF(3, w[%i] = exp(patch_diff[%i] * strength[%i]); ,i,i,dst_comp);
205 for (int i = 0; i < 4; i++)
206 GLSLF(3, w_sum %s w[%i][0] + w[%i][1] + w[%i][2] + w[%i][3];
207 ,!i ? "=" : "+=", i, i, i, i);
208 for (int i = 0; i < 4; i++)
209 GLSLF(3, sum %s dot(w[%i], src[%i]*255);
210 ,!i ? "=" : "+=", i, i);
211 }
212 GLSLC(0, );
213 if (t > 1) {
214 GLSLF(3, atomicAdd(weights_%i[y*ws_stride[%i] + x], w_sum); ,dst_comp, dst_comp);
215 GLSLF(3, atomicAdd(sums_%i[y*ws_stride[%i] + x], sum); ,dst_comp, dst_comp);
216 } else {
217 GLSLF(3, weights_%i[y*ws_stride[%i] + x] += w_sum; ,dst_comp, dst_comp);
218 GLSLF(3, sums_%i[y*ws_stride[%i] + x] += sum; ,dst_comp, dst_comp);
219 }
220 GLSLC(2, } );
221 GLSLC(1, } );
222}
223
224typedef struct HorizontalPushData {
225 VkDeviceAddress integral_data;
226 VkDeviceAddress state_data;
227 int32_t xoffs[TYPE_ELEMS];
228 int32_t yoffs[TYPE_ELEMS];
229 uint32_t width[4];
230 uint32_t height[4];
231 uint32_t ws_stride[4];
232 int32_t patch_size[4];
233 float strength[4];
234 uint32_t int_stride;
235} HorizontalPushData;
236
237static av_cold int init_weights_pipeline(FFVulkanContext *vkctx, FFVkExecPool *exec,
238 FFVulkanPipeline *pl, FFVkSPIRVShader *shd,
239 VkSampler sampler, FFVkSPIRVCompiler *spv,
240 int width, int height, int t,
241 const AVPixFmtDescriptor *desc,
242 int planes, int *nb_rows)
243{
244 int err;
245 uint8_t *spv_data;
246 size_t spv_len;
247 void *spv_opaque = NULL;
248 FFVulkanDescriptorSetBinding *desc_set;
249 int max_dim = FFMAX(width, height);
250 uint32_t max_wg = vkctx->props.properties.limits.maxComputeWorkGroupSize[0];
251 int max_shm = vkctx->props.properties.limits.maxComputeSharedMemorySize;
252 int wg_size, wg_rows;
253
254 /* Round the max workgroup size to the previous power of two */
255 max_wg = 1 << (31 - ff_clz(max_wg));
256 wg_size = max_wg;
257 wg_rows = 1;
258
259 if (max_wg > max_dim) {
260 wg_size = max_wg / (max_wg / max_dim);
261 } else if (max_wg < max_dim) {
262 /* First, make it fit */
263 while (wg_size*wg_rows < max_dim)
264 wg_rows++;
265
266 /* Second, make sure there's enough shared memory */
267 while ((wg_size * TYPE_SIZE + TYPE_SIZE + 2*4) > max_shm) {
268 wg_size >>= 1;
269 wg_rows++;
270 }
271 }
272
273 RET(ff_vk_shader_init(pl, shd, "nlmeans_weights", VK_SHADER_STAGE_COMPUTE_BIT, 0));
274 ff_vk_shader_set_compute_sizes(shd, wg_size, 1, 1);
275 *nb_rows = wg_rows;
276
277 if (t > 1)
278 GLSLC(0, #extension GL_EXT_shader_atomic_float : require );
279 GLSLC(0, #extension GL_ARB_gpu_shader_int64 : require );
280 GLSLC(0, #pragma use_vulkan_memory_model );
281 GLSLC(0, #extension GL_KHR_memory_scope_semantics : enable );
282 GLSLC(0, );
283 GLSLF(0, #define N_ROWS %i ,*nb_rows);
284 GLSLC(0, #define WG_SIZE (gl_WorkGroupSize.x) );
285 GLSLF(0, #define LG_WG_SIZE %i ,ff_log2(shd->local_size[0]));
286 GLSLC(0, #define PARTITION_SIZE (N_ROWS*WG_SIZE) );
287 GLSLF(0, #define DTYPE %s ,TYPE_NAME);
288 GLSLF(0, #define T_ALIGN %i ,TYPE_SIZE);
289 GLSLC(0, );
290 GLSLC(0, layout(buffer_reference, buffer_reference_align = T_ALIGN) coherent buffer DataBuffer { );
291 GLSLC(1, DTYPE v[]; );
292 GLSLC(0, }; );
293 GLSLC(0, );
294 GLSLC(0, layout(buffer_reference) buffer StateData; );
295 GLSLC(0, );
296 GLSLC(0, layout(push_constant, std430) uniform pushConstants { );
297 GLSLC(1, coherent DataBuffer integral_data; );
298 GLSLC(1, StateData state; );
299 GLSLF(1, uint xoffs[%i]; ,TYPE_ELEMS);
300 GLSLF(1, uint yoffs[%i]; ,TYPE_ELEMS);
301 GLSLC(1, uvec4 width; );
302 GLSLC(1, uvec4 height; );
303 GLSLC(1, uvec4 ws_stride; );
304 GLSLC(1, ivec4 patch_size; );
305 GLSLC(1, vec4 strength; );
306 GLSLC(1, uint int_stride; );
307 GLSLC(0, }; );
308 GLSLC(0, );
309
310 ff_vk_add_push_constant(pl, 0, sizeof(HorizontalPushData), VK_SHADER_STAGE_COMPUTE_BIT);
311
312 desc_set = (FFVulkanDescriptorSetBinding []) {
313 {
314 .name = "input_img",
315 .type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
316 .dimensions = 2,
317 .elems = planes,
318 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
319 .samplers = DUP_SAMPLER(sampler),
320 },
321 {
322 .name = "weights_buffer_0",
323 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
324 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
325 .buf_content = "float weights_0[];",
326 },
327 {
328 .name = "sums_buffer_0",
329 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
330 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
331 .buf_content = "float sums_0[];",
332 },
333 {
334 .name = "weights_buffer_1",
335 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
336 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
337 .buf_content = "float weights_1[];",
338 },
339 {
340 .name = "sums_buffer_1",
341 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
342 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
343 .buf_content = "float sums_1[];",
344 },
345 {
346 .name = "weights_buffer_2",
347 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
348 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
349 .buf_content = "float weights_2[];",
350 },
351 {
352 .name = "sums_buffer_2",
353 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
354 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
355 .buf_content = "float sums_2[];",
356 },
357 {
358 .name = "weights_buffer_3",
359 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
360 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
361 .buf_content = "float weights_3[];",
362 },
363 {
364 .name = "sums_buffer_3",
365 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
366 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
367 .buf_content = "float sums_3[];",
368 },
369 };
370 RET(ff_vk_pipeline_descriptor_set_add(vkctx, pl, shd, desc_set, 1 + 2*desc->nb_components, 0, 0));
371
372 GLSLD( ff_source_prefix_sum_comp );
373 GLSLC(0, );
374 GLSLC(0, void main() );
375 GLSLC(0, { );
376 GLSLC(1, uint64_t offset; );
377 GLSLC(1, DataBuffer dst; );
378 GLSLC(1, float s1; );
379 GLSLC(1, DTYPE s2; );
380 GLSLC(1, int r; );
381 GLSLC(1, int x; );
382 GLSLC(1, int y; );
383 GLSLC(1, int p; );
384 GLSLC(0, );
385 GLSLC(1, DTYPE a; );
386 GLSLC(1, DTYPE b; );
387 GLSLC(1, DTYPE c; );
388 GLSLC(1, DTYPE d; );
389 GLSLC(0, );
390 GLSLC(1, DTYPE patch_diff; );
391 if (TYPE_ELEMS == 4) {
392 GLSLC(1, vec4 src; );
393 GLSLC(1, vec4 w; );
394 } else {
395 GLSLC(1, vec4 src[4]; );
396 GLSLC(1, vec4 w[4]; );
397 }
398 GLSLC(1, float w_sum; );
399 GLSLC(1, float sum; );
400 GLSLC(0, );
401 GLSLC(1, bool lt; );
402 GLSLC(1, bool gt; );
403 GLSLC(0, );
404
405 for (int i = 0; i < desc->nb_components; i++) {
406 int off = desc->comp[i].offset / (FFALIGN(desc->comp[i].depth, 8)/8);
407 if (width > height) {
408 insert_horizontal_pass(shd, *nb_rows, 1, desc->comp[i].plane, off);
409 insert_vertical_pass(shd, *nb_rows, 0, desc->comp[i].plane, off);
410 insert_weights_pass(shd, *nb_rows, 0, t, i, desc->comp[i].plane, off);
411 } else {
412 insert_vertical_pass(shd, *nb_rows, 1, desc->comp[i].plane, off);
413 insert_horizontal_pass(shd, *nb_rows, 0, desc->comp[i].plane, off);
414 insert_weights_pass(shd, *nb_rows, 1, t, i, desc->comp[i].plane, off);
415 }
416 }
417
418 GLSLC(0, } );
419
420 RET(spv->compile_shader(spv, vkctx, shd, &spv_data, &spv_len, "main", &spv_opaque));
421 RET(ff_vk_shader_create(vkctx, shd, spv_data, spv_len, "main"));
422
423 RET(ff_vk_init_compute_pipeline(vkctx, pl, shd));
424 RET(ff_vk_exec_pipeline_register(vkctx, exec, pl));
425
426 return 0;
427
428fail:
429 if (spv_opaque)
430 spv->free_shader(spv, &spv_opaque);
431
432 return err;
433}
434
435typedef struct DenoisePushData {
436 uint32_t ws_stride[4];
437} DenoisePushData;
438
439static av_cold int init_denoise_pipeline(FFVulkanContext *vkctx, FFVkExecPool *exec,
440 FFVulkanPipeline *pl, FFVkSPIRVShader *shd,
441 VkSampler sampler, FFVkSPIRVCompiler *spv,
442 const AVPixFmtDescriptor *desc, int planes)
443{
444 int err;
445 uint8_t *spv_data;
446 size_t spv_len;
447 void *spv_opaque = NULL;
448 FFVulkanDescriptorSetBinding *desc_set;
449
450 RET(ff_vk_shader_init(pl, shd, "nlmeans_denoise",
451 VK_SHADER_STAGE_COMPUTE_BIT, 0));
452
453 ff_vk_shader_set_compute_sizes(shd, 32, 32, 1);
454
455 GLSLC(0, layout(push_constant, std430) uniform pushConstants { );
456 GLSLC(1, uvec4 ws_stride; );
457 GLSLC(0, }; );
458
459 ff_vk_add_push_constant(pl, 0, sizeof(DenoisePushData), VK_SHADER_STAGE_COMPUTE_BIT);
460
461 desc_set = (FFVulkanDescriptorSetBinding []) {
462 {
463 .name = "input_img",
464 .type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
465 .dimensions = 2,
466 .elems = planes,
467 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
468 .samplers = DUP_SAMPLER(sampler),
469 },
470 {
471 .name = "output_img",
472 .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
473 .mem_layout = ff_vk_shader_rep_fmt(vkctx->output_format),
474 .mem_quali = "writeonly",
475 .dimensions = 2,
476 .elems = planes,
477 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
478 },
479 {
480 .name = "weights_buffer_0",
481 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
482 .mem_quali = "readonly",
483 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
484 .buf_content = "float weights_0[];",
485 },
486 {
487 .name = "sums_buffer_0",
488 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
489 .mem_quali = "readonly",
490 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
491 .buf_content = "float sums_0[];",
492 },
493 {
494 .name = "weights_buffer_1",
495 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
496 .mem_quali = "readonly",
497 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
498 .buf_content = "float weights_1[];",
499 },
500 {
501 .name = "sums_buffer_1",
502 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
503 .mem_quali = "readonly",
504 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
505 .buf_content = "float sums_1[];",
506 },
507 {
508 .name = "weights_buffer_2",
509 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
510 .mem_quali = "readonly",
511 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
512 .buf_content = "float weights_2[];",
513 },
514 {
515 .name = "sums_buffer_2",
516 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
517 .mem_quali = "readonly",
518 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
519 .buf_content = "float sums_2[];",
520 },
521 {
522 .name = "weights_buffer_3",
523 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
524 .mem_quali = "readonly",
525 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
526 .buf_content = "float weights_3[];",
527 },
528 {
529 .name = "sums_buffer_3",
530 .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
531 .mem_quali = "readonly",
532 .stages = VK_SHADER_STAGE_COMPUTE_BIT,
533 .buf_content = "float sums_3[];",
534 },
535 };
536 RET(ff_vk_pipeline_descriptor_set_add(vkctx, pl, shd, desc_set, 2 + 2*desc->nb_components, 0, 0));
537
538 GLSLC(0, void main() );
539 GLSLC(0, { );
540 GLSLC(1, ivec2 size; );
541 GLSLC(1, const ivec2 pos = ivec2(gl_GlobalInvocationID.xy); );
542 GLSLC(0, );
543 GLSLC(1, float w_sum; );
544 GLSLC(1, float sum; );
545 GLSLC(1, vec4 src; );
546 GLSLC(1, vec4 r; );
547 GLSLC(0, );
548
549 for (int i = 0; i < planes; i++) {
550 GLSLF(1, src = texture(input_img[%i], pos); ,i);
551 for (int c = 0; c < desc->nb_components; c++) {
552 if (desc->comp[c].plane == i) {
553 int off = desc->comp[c].offset / (FFALIGN(desc->comp[c].depth, 8)/8);
554 GLSLF(1, w_sum = weights_%i[pos.y*ws_stride[%i] + pos.x]; ,c, c);
555 GLSLF(1, sum = sums_%i[pos.y*ws_stride[%i] + pos.x]; ,c, c);
556 GLSLF(1, r[%i] = (sum + src[%i]*255) / (1.0 + w_sum) / 255; ,off, off);
557 GLSLC(0, );
558 }
559 }
560 GLSLF(1, imageStore(output_img[%i], pos, r); ,i);
561 GLSLC(0, );
562 }
563
564 GLSLC(0, } );
565
566 RET(spv->compile_shader(spv, vkctx, shd, &spv_data, &spv_len, "main", &spv_opaque));
567 RET(ff_vk_shader_create(vkctx, shd, spv_data, spv_len, "main"));
568
569 RET(ff_vk_init_compute_pipeline(vkctx, pl, shd));
570 RET(ff_vk_exec_pipeline_register(vkctx, exec, pl));
571
572 return 0;
573
574fail:
575 if (spv_opaque)
576 spv->free_shader(spv, &spv_opaque);
577
578 return err;
579}
580
581static av_cold int init_filter(AVFilterContext *ctx)
582{
583 int rad, err;
584 int xcnt = 0, ycnt = 0;
585 NLMeansVulkanContext *s = ctx->priv;
586 FFVulkanContext *vkctx = &s->vkctx;
587 const int planes = av_pix_fmt_count_planes(s->vkctx.output_format);
588 FFVkSPIRVCompiler *spv;
589
590 const AVPixFmtDescriptor *desc;
591 desc = av_pix_fmt_desc_get(vkctx->output_format);
592 if (!desc)
593 return AVERROR(EINVAL);
594
595 if (!(s->opts.r & 1)) {
596 s->opts.r |= 1;
597 av_log(ctx, AV_LOG_WARNING, "Research size should be odd, setting to %i",
598 s->opts.r);
599 }
600
601 if (!(s->opts.p & 1)) {
602 s->opts.p |= 1;
603 av_log(ctx, AV_LOG_WARNING, "Patch size should be odd, setting to %i",
604 s->opts.p);
605 }
606
607 for (int i = 0; i < 4; i++) {
608 double str = (s->opts.sc[i] > 1.0) ? s->opts.sc[i] : s->opts.s;
609 int ps = (s->opts.pc[i] ? s->opts.pc[i] : s->opts.p);
610 str = 10.0f*str;
611 str *= -str;
612 str = 255.0*255.0 / str;
613 s->strength[i] = str;
614 if (!(ps & 1)) {
615 ps |= 1;
616 av_log(ctx, AV_LOG_WARNING, "Patch size should be odd, setting to %i",
617 ps);
618 }
619 s->patch[i] = ps / 2;
620 }
621
622 rad = s->opts.r/2;
623 s->nb_offsets = (2*rad + 1)*(2*rad + 1) - 1;
624 s->xoffsets = av_malloc(s->nb_offsets*sizeof(*s->xoffsets));
625 s->yoffsets = av_malloc(s->nb_offsets*sizeof(*s->yoffsets));
626 s->nb_offsets = 0;
627
628 for (int x = -rad; x <= rad; x++) {
629 for (int y = -rad; y <= rad; y++) {
630 if (!x && !y)
631 continue;
632
633 s->xoffsets[xcnt++] = x;
634 s->yoffsets[ycnt++] = y;
635 s->nb_offsets++;
636 }
637 }
638
639 s->opts.t = FFMIN(s->opts.t, (FFALIGN(s->nb_offsets, TYPE_ELEMS) / TYPE_ELEMS));
640 if (!vkctx->atomic_float_feats.shaderBufferFloat32AtomicAdd) {
641 av_log(ctx, AV_LOG_WARNING, "Device doesn't support atomic float adds, "
642 "disabling dispatch parallelism\n");
643 s->opts.t = 1;
644 }
645
646 if (!vkctx->feats_12.vulkanMemoryModel) {
647 av_log(ctx, AV_LOG_ERROR, "Device doesn't support the Vulkan memory model!");
648 return AVERROR(EINVAL);;
649 }
650
651 spv = ff_vk_spirv_init();
652 if (!spv) {
653 av_log(ctx, AV_LOG_ERROR, "Unable to initialize SPIR-V compiler!\n");
654 return AVERROR_EXTERNAL;
655 }
656
657 ff_vk_qf_init(vkctx, &s->qf, VK_QUEUE_COMPUTE_BIT);
658 RET(ff_vk_exec_pool_init(vkctx, &s->qf, &s->e, 1, 0, 0, 0, NULL));
659 RET(ff_vk_init_sampler(vkctx, &s->sampler, 1, VK_FILTER_NEAREST));
660
661 RET(init_weights_pipeline(vkctx, &s->e, &s->pl_weights, &s->shd_weights, s->sampler,
662 spv, s->vkctx.output_width, s->vkctx.output_height,
663 s->opts.t, desc, planes, &s->pl_weights_rows));
664
665 RET(init_denoise_pipeline(vkctx, &s->e, &s->pl_denoise, &s->shd_denoise, s->sampler,
666 spv, desc, planes));
667
668 av_log(ctx, AV_LOG_VERBOSE, "Filter initialized, %i x/y offsets, %i dispatches, %i parallel\n",
669 s->nb_offsets, (FFALIGN(s->nb_offsets, TYPE_ELEMS) / TYPE_ELEMS) + 1, s->opts.t);
670
671 s->initialized = 1;
672
673 return 0;
674
675fail:
676 if (spv)
677 spv->uninit(&spv);
678
679 return err;
680}
681
682static int denoise_pass(NLMeansVulkanContext *s, FFVkExecContext *exec,
683 FFVkBuffer *ws_vk, uint32_t ws_stride[4])
684{
685 FFVulkanContext *vkctx = &s->vkctx;
686 FFVulkanFunctions *vk = &vkctx->vkfn;
687 VkBufferMemoryBarrier2 buf_bar[8];
688 int nb_buf_bar = 0;
689
690 /* Denoise pass pipeline */
691 ff_vk_exec_bind_pipeline(vkctx, exec, &s->pl_denoise);
692
693 /* Push data */
694 ff_vk_update_push_exec(vkctx, exec, &s->pl_denoise, VK_SHADER_STAGE_COMPUTE_BIT,
695 0, sizeof(DenoisePushData), &(DenoisePushData) {
696 { ws_stride[0], ws_stride[1], ws_stride[2], ws_stride[3] },
697 });
698
699 buf_bar[nb_buf_bar++] = (VkBufferMemoryBarrier2) {
700 .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
701 .srcStageMask = ws_vk->stage,
702 .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
703 .srcAccessMask = ws_vk->access,
704 .dstAccessMask = VK_ACCESS_2_SHADER_STORAGE_READ_BIT,
705 .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
706 .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
707 .buffer = ws_vk->buf,
708 .size = ws_vk->size,
709 .offset = 0,
710 };
711
712 vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) {
713 .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
714 .pBufferMemoryBarriers = buf_bar,
715 .bufferMemoryBarrierCount = nb_buf_bar,
716 });
717 ws_vk->stage = buf_bar[0].dstStageMask;
718 ws_vk->access = buf_bar[0].dstAccessMask;
719
720 /* End of denoise pass */
721 vk->CmdDispatch(exec->buf,
722 FFALIGN(vkctx->output_width, s->pl_denoise.wg_size[0])/s->pl_denoise.wg_size[0],
723 FFALIGN(vkctx->output_height, s->pl_denoise.wg_size[1])/s->pl_denoise.wg_size[1],
724 1);
725
726 return 0;
727}
728
729static int nlmeans_vulkan_filter_frame(AVFilterLink *link, AVFrame *in)
730{
731 int err;
732 AVFrame *out = NULL;
733 AVFilterContext *ctx = link->dst;
734 NLMeansVulkanContext *s = ctx->priv;
735 AVFilterLink *outlink = ctx->outputs[0];
736 FFVulkanContext *vkctx = &s->vkctx;
737 FFVulkanFunctions *vk = &vkctx->vkfn;
738
739 const AVPixFmtDescriptor *desc;
740 int plane_widths[4];
741 int plane_heights[4];
742
743 /* Integral */
744 AVBufferRef *state_buf;
745 FFVkBuffer *state_vk;
746 AVBufferRef *integral_buf;
747 FFVkBuffer *integral_vk;
748 uint32_t int_stride;
749 size_t int_size;
750 size_t state_size;
751 int t_offset = 0;
752
753 /* Weights/sums */
754 AVBufferRef *ws_buf;
755 FFVkBuffer *ws_vk;
756 VkDeviceAddress weights_addr[4];
757 VkDeviceAddress sums_addr[4];
758 uint32_t ws_stride[4];
759 size_t ws_size[4];
760 size_t ws_total_size = 0;
761
762 FFVkExecContext *exec;
763 VkImageView in_views[AV_NUM_DATA_POINTERS];
764 VkImageView out_views[AV_NUM_DATA_POINTERS];
765 VkImageMemoryBarrier2 img_bar[8];
766 int nb_img_bar = 0;
767 VkBufferMemoryBarrier2 buf_bar[8];
768 int nb_buf_bar = 0;
769
770 if (!s->initialized)
771 RET(init_filter(ctx));
772
773 desc = av_pix_fmt_desc_get(vkctx->output_format);
774 if (!desc)
775 return AVERROR(EINVAL);
776
777 /* Integral image */
778 int_stride = s->pl_weights.wg_size[0]*s->pl_weights_rows;
779 int_size = int_stride * int_stride * TYPE_SIZE;
780 state_size = int_stride * 3 *TYPE_SIZE;
781
782 /* Plane dimensions */
783 for (int i = 0; i < desc->nb_components; i++) {
784 plane_widths[i] = !i || (i == 3) ? vkctx->output_width : AV_CEIL_RSHIFT(vkctx->output_width, desc->log2_chroma_w);
785 plane_heights[i] = !i || (i == 3) ? vkctx->output_height : AV_CEIL_RSHIFT(vkctx->output_height, desc->log2_chroma_w);
786 plane_widths[i] = FFALIGN(plane_widths[i], s->pl_denoise.wg_size[0]);
787 plane_heights[i] = FFALIGN(plane_heights[i], s->pl_denoise.wg_size[1]);
788
789 ws_stride[i] = plane_widths[i];
790 ws_size[i] = ws_stride[i] * plane_heights[i] * sizeof(float);
791 ws_total_size += ws_size[i];
792 }
793
794 /* Buffers */
795 err = ff_vk_get_pooled_buffer(&s->vkctx, &s->integral_buf_pool, &integral_buf,
796 VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
797 VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
798 NULL,
799 s->opts.t * int_size,
800 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
801 if (err < 0)
802 return err;
803 integral_vk = (FFVkBuffer *)integral_buf->data;
804
805 err = ff_vk_get_pooled_buffer(&s->vkctx, &s->state_buf_pool, &state_buf,
806 VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
807 VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
808 NULL,
809 s->opts.t * state_size,
810 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
811 if (err < 0)
812 return err;
813 state_vk = (FFVkBuffer *)state_buf->data;
814
815 err = ff_vk_get_pooled_buffer(&s->vkctx, &s->ws_buf_pool, &ws_buf,
816 VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
817 VK_BUFFER_USAGE_TRANSFER_DST_BIT |
818 VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
819 NULL,
820 ws_total_size * 2,
821 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
822 if (err < 0)
823 return err;
824 ws_vk = (FFVkBuffer *)ws_buf->data;
825
826 weights_addr[0] = ws_vk->address;
827 sums_addr[0] = ws_vk->address + ws_total_size;
828 for (int i = 1; i < desc->nb_components; i++) {
829 weights_addr[i] = weights_addr[i - 1] + ws_size[i - 1];
830 sums_addr[i] = sums_addr[i - 1] + ws_size[i - 1];
831 }
832
833 /* Output frame */
834 out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
835 if (!out) {
836 err = AVERROR(ENOMEM);
837 goto fail;
838 }
839
840 /* Execution context */
841 exec = ff_vk_exec_get(&s->e);
842 ff_vk_exec_start(vkctx, exec);
843
844 /* Dependencies */
845 RET(ff_vk_exec_add_dep_frame(vkctx, exec, in,
846 VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
847 VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT));
848 RET(ff_vk_exec_add_dep_frame(vkctx, exec, out,
849 VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
850 VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT));
851 RET(ff_vk_exec_add_dep_buf(vkctx, exec, &integral_buf, 1, 0));
852 RET(ff_vk_exec_add_dep_buf(vkctx, exec, &state_buf, 1, 0));
853 RET(ff_vk_exec_add_dep_buf(vkctx, exec, &ws_buf, 1, 0));
854
855 /* Input frame prep */
856 RET(ff_vk_create_imageviews(vkctx, exec, in_views, in));
857 ff_vk_update_descriptor_img_array(vkctx, &s->pl_weights, exec, in, in_views, 0, 0,
858 VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
859 s->sampler);
860 ff_vk_frame_barrier(vkctx, exec, in, img_bar, &nb_img_bar,
861 VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
862 VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
863 VK_ACCESS_SHADER_READ_BIT,
864 VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
865 VK_QUEUE_FAMILY_IGNORED);
866
867 /* Output frame prep */
868 RET(ff_vk_create_imageviews(vkctx, exec, out_views, out));
869 ff_vk_frame_barrier(vkctx, exec, out, img_bar, &nb_img_bar,
870 VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
871 VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
872 VK_ACCESS_SHADER_WRITE_BIT,
873 VK_IMAGE_LAYOUT_GENERAL,
874 VK_QUEUE_FAMILY_IGNORED);
875
876 buf_bar[nb_buf_bar++] = (VkBufferMemoryBarrier2) {
877 .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
878 .srcStageMask = ws_vk->stage,
879 .dstStageMask = VK_PIPELINE_STAGE_2_TRANSFER_BIT,
880 .srcAccessMask = ws_vk->access,
881 .dstAccessMask = VK_ACCESS_2_TRANSFER_WRITE_BIT,
882 .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
883 .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
884 .buffer = ws_vk->buf,
885 .size = ws_vk->size,
886 .offset = 0,
887 };
888
889 vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) {
890 .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
891 .pImageMemoryBarriers = img_bar,
892 .imageMemoryBarrierCount = nb_img_bar,
893 .pBufferMemoryBarriers = buf_bar,
894 .bufferMemoryBarrierCount = nb_buf_bar,
895 });
896 ws_vk->stage = buf_bar[0].dstStageMask;
897 ws_vk->access = buf_bar[0].dstAccessMask;
898
899 /* Weights/sums buffer zeroing */
900 vk->CmdFillBuffer(exec->buf, ws_vk->buf, 0, ws_vk->size, 0x0);
901
902 buf_bar[nb_buf_bar++] = (VkBufferMemoryBarrier2) {
903 .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
904 .srcStageMask = ws_vk->stage,
905 .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
906 .srcAccessMask = ws_vk->access,
907 .dstAccessMask = VK_ACCESS_2_SHADER_STORAGE_READ_BIT |
908 VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
909 .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
910 .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
911 .buffer = ws_vk->buf,
912 .size = ws_vk->size,
913 .offset = 0,
914 };
915
916 vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) {
917 .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
918 .pBufferMemoryBarriers = buf_bar,
919 .bufferMemoryBarrierCount = nb_buf_bar,
920 });
921 ws_vk->stage = buf_bar[0].dstStageMask;
922 ws_vk->access = buf_bar[0].dstAccessMask;
923
924 /* Update weights descriptors */
925 ff_vk_update_descriptor_img_array(vkctx, &s->pl_weights, exec, in, in_views, 0, 0,
926 VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
927 s->sampler);
928 for (int i = 0; i < desc->nb_components; i++) {
929 RET(ff_vk_set_descriptor_buffer(&s->vkctx, &s->pl_weights, exec, 0, 1 + i*2 + 0, 0,
930 weights_addr[i], ws_size[i],
931 VK_FORMAT_UNDEFINED));
932 RET(ff_vk_set_descriptor_buffer(&s->vkctx, &s->pl_weights, exec, 0, 1 + i*2 + 1, 0,
933 sums_addr[i], ws_size[i],
934 VK_FORMAT_UNDEFINED));
935 }
936
937 /* Update denoise descriptors */
938 ff_vk_update_descriptor_img_array(vkctx, &s->pl_denoise, exec, in, in_views, 0, 0,
939 VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
940 s->sampler);
941 ff_vk_update_descriptor_img_array(vkctx, &s->pl_denoise, exec, out, out_views, 0, 1,
942 VK_IMAGE_LAYOUT_GENERAL, s->sampler);
943 for (int i = 0; i < desc->nb_components; i++) {
944 RET(ff_vk_set_descriptor_buffer(&s->vkctx, &s->pl_denoise, exec, 0, 2 + i*2 + 0, 0,
945 weights_addr[i], ws_size[i],
946 VK_FORMAT_UNDEFINED));
947 RET(ff_vk_set_descriptor_buffer(&s->vkctx, &s->pl_denoise, exec, 0, 2 + i*2 + 1, 0,
948 sums_addr[i], ws_size[i],
949 VK_FORMAT_UNDEFINED));
950 }
951
952 /* Weights pipeline */
953 ff_vk_exec_bind_pipeline(vkctx, exec, &s->pl_weights);
954
955 for (int i = 0; i < s->nb_offsets; i += TYPE_ELEMS) {
956 int *xoffs = s->xoffsets + i;
957 int *yoffs = s->yoffsets + i;
958 HorizontalPushData pd = {
959 integral_vk->address + t_offset*int_size,
960 state_vk->address + t_offset*state_size,
961 { 0 },
962 { 0 },
963 { plane_widths[0], plane_widths[1], plane_widths[2], plane_widths[3] },
964 { plane_heights[0], plane_heights[1], plane_heights[2], plane_heights[3] },
965 { ws_stride[0], ws_stride[1], ws_stride[2], ws_stride[3] },
966 { s->patch[0], s->patch[1], s->patch[2], s->patch[3] },
967 { s->strength[0], s->strength[1], s->strength[2], s->strength[2], },
968 int_stride,
969 };
970
971 memcpy(pd.xoffs, xoffs, sizeof(pd.xoffs));
972 memcpy(pd.yoffs, yoffs, sizeof(pd.yoffs));
973
974 /* Put a barrier once we run out of parallelism buffers */
975 if (!t_offset) {
976 nb_buf_bar = 0;
977 /* Buffer prep/sync */
978 buf_bar[nb_buf_bar++] = (VkBufferMemoryBarrier2) {
979 .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
980 .srcStageMask = integral_vk->stage,
981 .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
982 .srcAccessMask = integral_vk->access,
983 .dstAccessMask = VK_ACCESS_2_SHADER_STORAGE_READ_BIT |
984 VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
985 .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
986 .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
987 .buffer = integral_vk->buf,
988 .size = integral_vk->size,
989 .offset = 0,
990 };
991 buf_bar[nb_buf_bar++] = (VkBufferMemoryBarrier2) {
992 .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
993 .srcStageMask = state_vk->stage,
994 .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
995 .srcAccessMask = state_vk->access,
996 .dstAccessMask = VK_ACCESS_2_SHADER_STORAGE_READ_BIT |
997 VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
998 .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
999 .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
1000 .buffer = state_vk->buf,
1001 .size = state_vk->size,
1002 .offset = 0,
1003 };
1004
1005 vk->CmdPipelineBarrier2(exec->buf, &(VkDependencyInfo) {
1006 .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
1007 .pBufferMemoryBarriers = buf_bar,
1008 .bufferMemoryBarrierCount = nb_buf_bar,
1009 });
1010 integral_vk->stage = buf_bar[0].dstStageMask;
1011 integral_vk->access = buf_bar[0].dstAccessMask;
1012 state_vk->stage = buf_bar[1].dstStageMask;
1013 state_vk->access = buf_bar[1].dstAccessMask;
1014 }
1015 t_offset = (t_offset + 1) % s->opts.t;
1016
1017 /* Push data */
1018 ff_vk_update_push_exec(vkctx, exec, &s->pl_weights, VK_SHADER_STAGE_COMPUTE_BIT,
1019 0, sizeof(pd), &pd);
1020
1021 /* End of horizontal pass */
1022 vk->CmdDispatch(exec->buf, 1, 1, 1);
1023 }
1024
1025 RET(denoise_pass(s, exec, ws_vk, ws_stride));
1026
1027 err = ff_vk_exec_submit(vkctx, exec);
1028 if (err < 0)
1029 return err;
1030
1031 err = av_frame_copy_props(out, in);
1032 if (err < 0)
1033 goto fail;
1034
1035 av_frame_free(&in);
1036
1037 return ff_filter_frame(outlink, out);
1038
1039fail:
1040 av_frame_free(&in);
1041 av_frame_free(&out);
1042 return err;
1043}
1044
1045static void nlmeans_vulkan_uninit(AVFilterContext *avctx)
1046{
1047 NLMeansVulkanContext *s = avctx->priv;
1048 FFVulkanContext *vkctx = &s->vkctx;
1049 FFVulkanFunctions *vk = &vkctx->vkfn;
1050
1051 ff_vk_exec_pool_free(vkctx, &s->e);
1052 ff_vk_pipeline_free(vkctx, &s->pl_weights);
1053 ff_vk_shader_free(vkctx, &s->shd_weights);
1054 ff_vk_pipeline_free(vkctx, &s->pl_denoise);
1055 ff_vk_shader_free(vkctx, &s->shd_denoise);
1056
1057 av_buffer_pool_uninit(&s->integral_buf_pool);
1058 av_buffer_pool_uninit(&s->state_buf_pool);
1059 av_buffer_pool_uninit(&s->ws_buf_pool);
1060
1061 if (s->sampler)
1062 vk->DestroySampler(vkctx->hwctx->act_dev, s->sampler,
1063 vkctx->hwctx->alloc);
1064
1065 ff_vk_uninit(&s->vkctx);
1066
1067 s->initialized = 0;
1068}
1069
1070#define OFFSET(x) offsetof(NLMeansVulkanContext, x)
1071#define FLAGS (AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM)
1072static const AVOption nlmeans_vulkan_options[] = {
1073 { "s", "denoising strength for all components", OFFSET(opts.s), AV_OPT_TYPE_DOUBLE, { .dbl = 1.0 }, 1.0, 100.0, FLAGS },
1074 { "p", "patch size for all components", OFFSET(opts.p), AV_OPT_TYPE_INT, { .i64 = 3*2+1 }, 0, 99, FLAGS },
1075 { "r", "research window radius", OFFSET(opts.r), AV_OPT_TYPE_INT, { .i64 = 7*2+1 }, 0, 99, FLAGS },
1076 { "t", "parallelism", OFFSET(opts.t), AV_OPT_TYPE_INT, { .i64 = 36 }, 1, 168, FLAGS },
1077
1078 { "s1", "denoising strength for component 1", OFFSET(opts.sc[0]), AV_OPT_TYPE_DOUBLE, { .dbl = 1.0 }, 1.0, 100.0, FLAGS },
1079 { "s2", "denoising strength for component 2", OFFSET(opts.sc[1]), AV_OPT_TYPE_DOUBLE, { .dbl = 1.0 }, 1.0, 100.0, FLAGS },
1080 { "s3", "denoising strength for component 3", OFFSET(opts.sc[2]), AV_OPT_TYPE_DOUBLE, { .dbl = 1.0 }, 1.0, 100.0, FLAGS },
1081 { "s4", "denoising strength for component 4", OFFSET(opts.sc[3]), AV_OPT_TYPE_DOUBLE, { .dbl = 1.0 }, 1.0, 100.0, FLAGS },
1082
1083 { "p1", "patch size for component 1", OFFSET(opts.pc[0]), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 99, FLAGS },
1084 { "p2", "patch size for component 2", OFFSET(opts.pc[1]), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 99, FLAGS },
1085 { "p3", "patch size for component 3", OFFSET(opts.pc[2]), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 99, FLAGS },
1086 { "p4", "patch size for component 4", OFFSET(opts.pc[3]), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 99, FLAGS },
1087
1088 { NULL }
1089};
1090
1091AVFILTER_DEFINE_CLASS(nlmeans_vulkan);
1092
1093static const AVFilterPad nlmeans_vulkan_inputs[] = {
1094 {
1095 .name = "default",
1096 .type = AVMEDIA_TYPE_VIDEO,
1097 .filter_frame = &nlmeans_vulkan_filter_frame,
1098 .config_props = &ff_vk_filter_config_input,
1099 },
1100};
1101
1102static const AVFilterPad nlmeans_vulkan_outputs[] = {
1103 {
1104 .name = "default",
1105 .type = AVMEDIA_TYPE_VIDEO,
1106 .config_props = &ff_vk_filter_config_output,
1107 },
1108};
1109
1110const AVFilter ff_vf_nlmeans_vulkan = {
1111 .name = "nlmeans_vulkan",
1112 .description = NULL_IF_CONFIG_SMALL("Non-local means denoiser (Vulkan)"),
1113 .priv_size = sizeof(NLMeansVulkanContext),
1114 .init = &ff_vk_filter_init,
1115 .uninit = &nlmeans_vulkan_uninit,
1116 FILTER_INPUTS(nlmeans_vulkan_inputs),
1117 FILTER_OUTPUTS(nlmeans_vulkan_outputs),
1118 FILTER_SINGLE_PIXFMT(AV_PIX_FMT_VULKAN),
1119 .priv_class = &nlmeans_vulkan_class,
1120 .flags = AVFILTER_FLAG_HWDEVICE,
1121 .flags_internal = FF_FILTER_FLAG_HWFRAME_AWARE,
1122};
libavfilter/vulkan/prefix_sum.comp+151−0
@@ -0,0 +1,151 @@
1#extension GL_EXT_buffer_reference : require
2#extension GL_EXT_buffer_reference2 : require
3
4#define ACQUIRE gl_StorageSemanticsBuffer, gl_SemanticsAcquire
5#define RELEASE gl_StorageSemanticsBuffer, gl_SemanticsRelease
6
7// These correspond to X, A, P respectively in the prefix sum paper.
8#define FLAG_NOT_READY 0u
9#define FLAG_AGGREGATE_READY 1u
10#define FLAG_PREFIX_READY 2u
11
12layout(buffer_reference, buffer_reference_align = T_ALIGN) nonprivate buffer StateData {
13 DTYPE aggregate;
14 DTYPE prefix;
15 uint flag;
16};
17
18shared DTYPE sh_scratch[WG_SIZE];
19shared DTYPE sh_prefix;
20shared uint sh_part_ix;
21shared uint sh_flag;
22
23void prefix_sum(DataBuffer dst, uint dst_stride, DataBuffer src, uint src_stride)
24{
25 DTYPE local[N_ROWS];
26 // Determine partition to process by atomic counter (described in Section 4.4 of prefix sum paper).
27 if (gl_GlobalInvocationID.x == 0)
28 sh_part_ix = gl_WorkGroupID.x;
29// sh_part_ix = atomicAdd(part_counter, 1);
30
31 barrier();
32 uint part_ix = sh_part_ix;
33
34 uint ix = part_ix * PARTITION_SIZE + gl_LocalInvocationID.x * N_ROWS;
35
36 // TODO: gate buffer read? (evaluate whether shader check or CPU-side padding is better)
37 local[0] = src.v[ix*src_stride];
38 for (uint i = 1; i < N_ROWS; i++)
39 local[i] = local[i - 1] + src.v[(ix + i)*src_stride];
40
41 DTYPE agg = local[N_ROWS - 1];
42 sh_scratch[gl_LocalInvocationID.x] = agg;
43 for (uint i = 0; i < LG_WG_SIZE; i++) {
44 barrier();
45 if (gl_LocalInvocationID.x >= (1u << i))
46 agg += sh_scratch[gl_LocalInvocationID.x - (1u << i)];
47 barrier();
48
49 sh_scratch[gl_LocalInvocationID.x] = agg;
50 }
51
52 // Publish aggregate for this partition
53 if (gl_LocalInvocationID.x == WG_SIZE - 1) {
54 state[part_ix].aggregate = agg;
55 if (part_ix == 0)
56 state[0].prefix = agg;
57 }
58
59 // Write flag with release semantics
60 if (gl_LocalInvocationID.x == WG_SIZE - 1) {
61 uint flag = part_ix == 0 ? FLAG_PREFIX_READY : FLAG_AGGREGATE_READY;
62 atomicStore(state[part_ix].flag, flag, gl_ScopeDevice, RELEASE);
63 }
64
65 DTYPE exclusive = DTYPE(0);
66 if (part_ix != 0) {
67 // step 4 of paper: decoupled lookback
68 uint look_back_ix = part_ix - 1;
69
70 DTYPE their_agg;
71 uint their_ix = 0;
72 while (true) {
73 // Read flag with acquire semantics.
74 if (gl_LocalInvocationID.x == WG_SIZE - 1)
75 sh_flag = atomicLoad(state[look_back_ix].flag, gl_ScopeDevice, ACQUIRE);
76
77 // The flag load is done only in the last thread. However, because the
78 // translation of memoryBarrierBuffer to Metal requires uniform control
79 // flow, we broadcast it to all threads.
80 barrier();
81
82 uint flag = sh_flag;
83 barrier();
84
85 if (flag == FLAG_PREFIX_READY) {
86 if (gl_LocalInvocationID.x == WG_SIZE - 1) {
87 DTYPE their_prefix = state[look_back_ix].prefix;
88 exclusive = their_prefix + exclusive;
89 }
90 break;
91 } else if (flag == FLAG_AGGREGATE_READY) {
92 if (gl_LocalInvocationID.x == WG_SIZE - 1) {
93 their_agg = state[look_back_ix].aggregate;
94 exclusive = their_agg + exclusive;
95 }
96 look_back_ix--;
97 their_ix = 0;
98 continue;
99 } // else spins
100
101 if (gl_LocalInvocationID.x == WG_SIZE - 1) {
102 // Unfortunately there's no guarantee of forward progress of other
103 // workgroups, so compute a bit of the aggregate before trying again.
104 // In the worst case, spinning stops when the aggregate is complete.
105 DTYPE m = src.v[(look_back_ix * PARTITION_SIZE + their_ix)*src_stride];
106 if (their_ix == 0)
107 their_agg = m;
108 else
109 their_agg += m;
110
111 their_ix++;
112 if (their_ix == PARTITION_SIZE) {
113 exclusive = their_agg + exclusive;
114 if (look_back_ix == 0) {
115 sh_flag = FLAG_PREFIX_READY;
116 } else {
117 look_back_ix--;
118 their_ix = 0;
119 }
120 }
121 }
122 barrier();
123 flag = sh_flag;
124 barrier();
125 if (flag == FLAG_PREFIX_READY)
126 break;
127 }
128
129 // step 5 of paper: compute inclusive prefix
130 if (gl_LocalInvocationID.x == WG_SIZE - 1) {
131 DTYPE inclusive_prefix = exclusive + agg;
132 sh_prefix = exclusive;
133 state[part_ix].prefix = inclusive_prefix;
134 }
135
136 if (gl_LocalInvocationID.x == WG_SIZE - 1)
137 atomicStore(state[part_ix].flag, FLAG_PREFIX_READY, gl_ScopeDevice, RELEASE);
138 }
139
140 barrier();
141 if (part_ix != 0)
142 exclusive = sh_prefix;
143
144 DTYPE row = exclusive;
145 if (gl_LocalInvocationID.x > 0)
146 row += sh_scratch[gl_LocalInvocationID.x - 1];
147
148 // note - may overwrite
149 for (uint i = 0; i < N_ROWS; i++)
150 dst.v[(ix + i)*dst_stride] = row + local[i];
151}

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