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

libavfilter/Makefile +2 −0
libavfilter/allfilters.c +1 −0
libavfilter/vf_nlmeans_vulkan.c +1122 −0
libavfilter/vulkan/prefix_sum.comp +151 −0

@@ -390,6 +390,8 @@ OBJS-$(CONFIG_MULTIPLY_FILTER) += vf_multiply.o · 389 unmodified lines
390	390	OBJS-$(CONFIG_NEGATE_FILTER) += vf_negate.o
391	391	OBJS-$(CONFIG_NLMEANS_FILTER) += vf_nlmeans.o
392	392	OBJS-$(CONFIG_NLMEANS_OPENCL_FILTER) += vf_nlmeans_opencl.o opencl.o opencl/nlmeans.o
	393	+OBJS-$(CONFIG_NLMEANS_VULKAN_FILTER) += vf_nlmeans_vulkan.o vulkan.o vulkan_filter.o \
	394	+ vulkan/prefix_sum.o
393	395	OBJS-$(CONFIG_NNEDI_FILTER) += vf_nnedi.o
394	396	OBJS-$(CONFIG_NOFORMAT_FILTER) += vf_format.o
395	397	OBJS-$(CONFIG_NOISE_FILTER) += vf_noise.o

libavfilter/allfilters.c+1−0

@@ -368,6 +368,7 @@ extern const AVFilter ff_vf_multiply; · 367 unmodified lines
368	368	extern const AVFilter ff_vf_negate;
369	369	extern const AVFilter ff_vf_nlmeans;
370	370	extern const AVFilter ff_vf_nlmeans_opencl;
	371	+extern const AVFilter ff_vf_nlmeans_vulkan;
371	372	extern const AVFilter ff_vf_nnedi;
372	373	extern const AVFilter ff_vf_noformat;
373	374	extern 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	+
	31	+typedef 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	+
	66	+extern const char *ff_source_prefix_sum_comp;
	67	+
	68	+static 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	+
	92	+static 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)yT_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	+
	119	+static 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	+
	146	+static 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	+
	224	+typedef 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	+
	237	+static 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	+
	428	+fail:
	429	+ if (spv_opaque)
	430	+ spv->free_shader(spv, &spv_opaque);
	431	+
	432	+ return err;
	433	+}
	434	+
	435	+typedef struct DenoisePushData {
	436	+ uint32_t ws_stride[4];
	437	+} DenoisePushData;
	438	+
	439	+static 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	+
	574	+fail:
	575	+ if (spv_opaque)
	576	+ spv->free_shader(spv, &spv_opaque);
	577	+
	578	+ return err;
	579	+}
	580	+
	581	+static 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 = (2rad + 1)(2*rad + 1) - 1;
	624	+ s->xoffsets = av_malloc(s->nb_offsetssizeof(s->xoffsets));
	625	+ s->yoffsets = av_malloc(s->nb_offsetssizeof(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	+
	675	+fail:
	676	+ if (spv)
	677	+ spv->uninit(&spv);
	678	+
	679	+ return err;
	680	+}
	681	+
	682	+static 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	+
	729	+static 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	+
	1039	+fail:
	1040	+ av_frame_free(&in);
	1041	+ av_frame_free(&out);
	1042	+ return err;
	1043	+}
	1044	+
	1045	+static 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)
	1072	+static 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	+
	1091	+AVFILTER_DEFINE_CLASS(nlmeans_vulkan);
	1092	+
	1093	+static 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	+
	1102	+static 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	+
	1110	+const 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	+
	12	+layout(buffer_reference, buffer_reference_align = T_ALIGN) nonprivate buffer StateData {
	13	+ DTYPE aggregate;
	14	+ DTYPE prefix;
	15	+ uint flag;
	16	+};
	17	+
	18	+shared DTYPE sh_scratch[WG_SIZE];
	19	+shared DTYPE sh_prefix;
	20	+shared uint sh_part_ix;
	21	+shared uint sh_flag;
	22	+
	23	+void 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	+}

lavfi: add nlmeans_vulkan filter

Keyboard shortcuts

General

Issues and pull requests

Issues, pull requests and Inbox

Repository code