| @@ -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)*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 | + |
| 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 = (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 | + |
| 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 | +}; |