Sara/behaviour-tree-test
1
0
Fork 0
Code Issues Pull requests Projects Releases 2 Packages Wiki Activity Actions

Replace OIDN denoiser with a JNLM denoiser compute shader implementation.

Browse source
This commit is contained in:
Dario 2023-09-11 14:59:50 -03:00
parent fcbc50ec14
commit 1b2b726502
8 changed files with 304 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

132
modules/lightmapper_rd/lightmapper_rd.cpp
View file

@ -614,24 +614,28 @@ void LightmapperRD::_raster_geometry(RenderingDevice *rd, Size2i atlas_size, int
}
}
LightmapperRD::BakeError LightmapperRD::_dilate(RenderingDevice *rd, Ref<RDShaderFile> &compute_shader, RID &compute_base_uniform_set, PushConstant &push_constant, RID &source_light_tex, RID &dest_light_tex, const Size2i &atlas_size, int atlas_slices) {
static Vector<RD::Uniform> dilate_or_denoise_common_uniforms(RID &p_source_light_tex, RID &p_dest_light_tex) {
Vector<RD::Uniform> uniforms;
{
{
RD::Uniform u;
u.uniform_type = RD::UNIFORM_TYPE_IMAGE;
u.binding = 0;
u.append_id(dest_light_tex);
uniforms.push_back(u);
}
{
RD::Uniform u;
u.uniform_type = RD::UNIFORM_TYPE_TEXTURE;
u.binding = 1;
u.append_id(source_light_tex);
uniforms.push_back(u);
}
RD::Uniform u;
u.uniform_type = RD::UNIFORM_TYPE_IMAGE;
u.binding = 0;
u.append_id(p_dest_light_tex);
uniforms.push_back(u);
}
{
RD::Uniform u;
u.uniform_type = RD::UNIFORM_TYPE_TEXTURE;
u.binding = 1;
u.append_id(p_source_light_tex);
uniforms.push_back(u);
}
return uniforms;
}
LightmapperRD::BakeError LightmapperRD::_dilate(RenderingDevice *rd, Ref<RDShaderFile> &compute_shader, RID &compute_base_uniform_set, PushConstant &push_constant, RID &source_light_tex, RID &dest_light_tex, const Size2i &atlas_size, int atlas_slices) {
Vector<RD::Uniform> uniforms = dilate_or_denoise_common_uniforms(source_light_tex, dest_light_tex);
RID compute_shader_dilate = rd->shader_create_from_spirv(compute_shader->get_spirv_stages("dilate"));
ERR_FAIL_COND_V(compute_shader_dilate.is_null(), BAKE_ERROR_LIGHTMAP_CANT_PRE_BAKE_MESHES); //internal check, should not happen
@ -667,7 +671,77 @@ LightmapperRD::BakeError LightmapperRD::_dilate(RenderingDevice *rd, Ref<RDShade
return BAKE_OK;
}
LightmapperRD::BakeError LightmapperRD::bake(BakeQuality p_quality, bool p_use_denoiser, int p_bounces, float p_bias, int p_max_texture_size, bool p_bake_sh, GenerateProbes p_generate_probes, const Ref<Image> &p_environment_panorama, const Basis &p_environment_transform, BakeStepFunc p_step_function, void *p_bake_userdata, float p_exposure_normalization) {
LightmapperRD::BakeError LightmapperRD::_denoise(RenderingDevice *p_rd, Ref<RDShaderFile> &p_compute_shader, const RID &p_compute_base_uniform_set, PushConstant &p_push_constant, RID p_source_light_tex, RID p_source_normal_tex, RID p_dest_light_tex, float p_denoiser_strength, const Size2i &p_atlas_size, int p_atlas_slices, bool p_bake_sh, BakeStepFunc p_step_function) {
RID denoise_params_buffer = p_rd->uniform_buffer_create(sizeof(DenoiseParams));
DenoiseParams denoise_params;
denoise_params.spatial_bandwidth = 5.0f;
denoise_params.light_bandwidth = p_denoiser_strength;
denoise_params.albedo_bandwidth = 1.0f;
denoise_params.normal_bandwidth = 0.1f;
denoise_params.filter_strength = 10.0f;
p_rd->buffer_update(denoise_params_buffer, 0, sizeof(DenoiseParams), &denoise_params);
Vector<RD::Uniform> uniforms = dilate_or_denoise_common_uniforms(p_source_light_tex, p_dest_light_tex);
{
RD::Uniform u;
u.uniform_type = RD::UNIFORM_TYPE_TEXTURE;
u.binding = 2;
u.append_id(p_source_normal_tex);
uniforms.push_back(u);
}
{
RD::Uniform u;
u.uniform_type = RD::UNIFORM_TYPE_UNIFORM_BUFFER;
u.binding = 3;
u.append_id(denoise_params_buffer);
uniforms.push_back(u);
}
RID compute_shader_denoise = p_rd->shader_create_from_spirv(p_compute_shader->get_spirv_stages("denoise"));
ERR_FAIL_COND_V(compute_shader_denoise.is_null(), BAKE_ERROR_LIGHTMAP_CANT_PRE_BAKE_MESHES);
RID compute_shader_denoise_pipeline = p_rd->compute_pipeline_create(compute_shader_denoise);
RID denoise_uniform_set = p_rd->uniform_set_create(uniforms, compute_shader_denoise, 1);
// We denoise in fixed size regions and synchronize execution to avoid GPU timeouts.
// We use a region with 1/4 the amount of pixels if we're denoising SH lightmaps, as
// all four of them are denoised in the shader in one dispatch.
const int max_region_size = p_bake_sh ? 512 : 1024;
int x_regions = (p_atlas_size.width - 1) / max_region_size + 1;
int y_regions = (p_atlas_size.height - 1) / max_region_size + 1;
for (int s = 0; s < p_atlas_slices; s++) {
p_push_constant.atlas_slice = s;
for (int i = 0; i < x_regions; i++) {
for (int j = 0; j < y_regions; j++) {
int x = i * max_region_size;
int y = j * max_region_size;
int w = MIN((i + 1) * max_region_size, p_atlas_size.width) - x;
int h = MIN((j + 1) * max_region_size, p_atlas_size.height) - y;
p_push_constant.region_ofs[0] = x;
p_push_constant.region_ofs[1] = y;
RD::ComputeListID compute_list = p_rd->compute_list_begin();
p_rd->compute_list_bind_compute_pipeline(compute_list, compute_shader_denoise_pipeline);
p_rd->compute_list_bind_uniform_set(compute_list, p_compute_base_uniform_set, 0);
p_rd->compute_list_bind_uniform_set(compute_list, denoise_uniform_set, 1);
p_rd->compute_list_set_push_constant(compute_list, &p_push_constant, sizeof(PushConstant));
p_rd->compute_list_dispatch(compute_list, (w - 1) / 8 + 1, (h - 1) / 8 + 1, 1);
p_rd->compute_list_end();
p_rd->submit();
p_rd->sync();
}
}
}
p_rd->free(compute_shader_denoise);
p_rd->free(denoise_params_buffer);
return BAKE_OK;
}
LightmapperRD::BakeError LightmapperRD::bake(BakeQuality p_quality, bool p_use_denoiser, float p_denoiser_strength, int p_bounces, float p_bias, int p_max_texture_size, bool p_bake_sh, GenerateProbes p_generate_probes, const Ref<Image> &p_environment_panorama, const Basis &p_environment_transform, BakeStepFunc p_step_function, void *p_bake_userdata, float p_exposure_normalization) {
if (p_step_function) {
p_step_function(0.0, RTR("Begin Bake"), p_bake_userdata, true);
}
@ -1434,27 +1508,11 @@ LightmapperRD::BakeError LightmapperRD::bake(BakeQuality p_quality, bool p_use_d
p_step_function(0.8, RTR("Denoising"), p_bake_userdata, true);
}
Ref<LightmapDenoiser> denoiser = LightmapDenoiser::create();
if (denoiser.is_valid()) {
for (int i = 0; i < atlas_slices * (p_bake_sh ? 4 : 1); i++) {
Vector<uint8_t> s = rd->texture_get_data(light_accum_tex, i);
Ref<Image> img = Image::create_from_data(atlas_size.width, atlas_size.height, false, Image::FORMAT_RGBAH, s);
Ref<Image> denoised = denoiser->denoise_image(img);
if (denoised != img) {
denoised->convert(Image::FORMAT_RGBAH);
Vector<uint8_t> ds = denoised->get_data();
denoised.unref(); //avoid copy on write
{ //restore alpha
uint32_t count = s.size() / 2; //uint16s
const uint16_t *src = (const uint16_t *)s.ptr();
uint16_t *dst = (uint16_t *)ds.ptrw();
for (uint32_t j = 0; j < count; j += 4) {
dst[j + 3] = src[j + 3];
}
}
rd->texture_update(light_accum_tex, i, ds);
}
{
SWAP(light_accum_tex, light_accum_tex2);
BakeError error = _denoise(rd, compute_shader, compute_base_uniform_set, push_constant, light_accum_tex2, normal_tex, light_accum_tex, p_denoiser_strength, atlas_size, atlas_slices, p_bake_sh, p_step_function);
if (unlikely(error != BAKE_OK)) {
return error;
}
}

13
modules/lightmapper_rd/lightmapper_rd.h
View file

@ -229,11 +229,22 @@ class LightmapperRD : public Lightmapper {
Vector<Ref<Image>> bake_textures;
Vector<Color> probe_values;
struct DenoiseParams {
float spatial_bandwidth;
float light_bandwidth;
float albedo_bandwidth;
float normal_bandwidth;
float filter_strength;
float pad[3];
};
BakeError _blit_meshes_into_atlas(int p_max_texture_size, Vector<Ref<Image>> &albedo_images, Vector<Ref<Image>> &emission_images, AABB &bounds, Size2i &atlas_size, int &atlas_slices, BakeStepFunc p_step_function, void *p_bake_userdata);
void _create_acceleration_structures(RenderingDevice *rd, Size2i atlas_size, int atlas_slices, AABB &bounds, int grid_size, Vector<Probe> &probe_positions, GenerateProbes p_generate_probes, Vector<int> &slice_triangle_count, Vector<int> &slice_seam_count, RID &vertex_buffer, RID &triangle_buffer, RID &lights_buffer, RID &triangle_cell_indices_buffer, RID &probe_positions_buffer, RID &grid_texture, RID &seams_buffer, BakeStepFunc p_step_function, void *p_bake_userdata);
void _raster_geometry(RenderingDevice *rd, Size2i atlas_size, int atlas_slices, int grid_size, AABB bounds, float p_bias, Vector<int> slice_triangle_count, RID position_tex, RID unocclude_tex, RID normal_tex, RID raster_depth_buffer, RID rasterize_shader, RID raster_base_uniform);
BakeError _dilate(RenderingDevice *rd, Ref<RDShaderFile> &compute_shader, RID &compute_base_uniform_set, PushConstant &push_constant, RID &source_light_tex, RID &dest_light_tex, const Size2i &atlas_size, int atlas_slices);
BakeError _denoise(RenderingDevice *p_rd, Ref<RDShaderFile> &p_compute_shader, const RID &p_compute_base_uniform_set, PushConstant &p_push_constant, RID p_source_light_tex, RID p_source_normal_tex, RID p_dest_light_tex, float p_denoiser_strength, const Size2i &p_atlas_size, int p_atlas_slices, bool p_bake_sh, BakeStepFunc p_step_function);
public:
virtual void add_mesh(const MeshData &p_mesh) override;
@ -241,7 +252,7 @@ public:
virtual void add_omni_light(bool p_static, const Vector3 &p_position, const Color &p_color, float p_energy, float p_range, float p_attenuation, float p_size, float p_shadow_blur) override;
virtual void add_spot_light(bool p_static, const Vector3 &p_position, const Vector3 p_direction, const Color &p_color, float p_energy, float p_range, float p_attenuation, float p_spot_angle, float p_spot_attenuation, float p_size, float p_shadow_blur) override;
virtual void add_probe(const Vector3 &p_position) override;
virtual BakeError bake(BakeQuality p_quality, bool p_use_denoiser, int p_bounces, float p_bias, int p_max_texture_size, bool p_bake_sh, GenerateProbes p_generate_probes, const Ref<Image> &p_environment_panorama, const Basis &p_environment_transform, BakeStepFunc p_step_function = nullptr, void *p_bake_userdata = nullptr, float p_exposure_normalization = 1.0) override;
virtual BakeError bake(BakeQuality p_quality, bool p_use_denoiser, float p_denoiser_strength, int p_bounces, float p_bias, int p_max_texture_size, bool p_bake_sh, GenerateProbes p_generate_probes, const Ref<Image> &p_environment_panorama, const Basis &p_environment_transform, BakeStepFunc p_step_function = nullptr, void *p_bake_userdata = nullptr, float p_exposure_normalization = 1.0) override;
int get_bake_texture_count() const override;
Ref<Image> get_bake_texture(int p_index) const override;

165
modules/lightmapper_rd/lm_compute.glsl
View file

@ -5,6 +5,7 @@ secondary = "#define MODE_BOUNCE_LIGHT";
dilate = "#define MODE_DILATE";
unocclude = "#define MODE_UNOCCLUDE";
light_probes = "#define MODE_LIGHT_PROBES";
denoise = "#define MODE_DENOISE";
#[compute]
@ -65,11 +66,24 @@ layout(set = 1, binding = 6) uniform texture2D environment;
layout(rgba32f, set = 1, binding = 5) uniform restrict writeonly image2DArray primary_dynamic;
#endif
#ifdef MODE_DILATE
#if defined(MODE_DILATE) || defined(MODE_DENOISE)
layout(rgba16f, set = 1, binding = 0) uniform restrict writeonly image2DArray dest_light;
layout(set = 1, binding = 1) uniform texture2DArray source_light;
#endif
#ifdef MODE_DENOISE
layout(set = 1, binding = 2) uniform texture2DArray source_normal;
layout(set = 1, binding = 3) uniform DenoiseParams {
float spatial_bandwidth;
float light_bandwidth;
float albedo_bandwidth;
float normal_bandwidth;
float filter_strength;
}
denoise_params;
#endif
layout(push_constant, std430) uniform Params {
ivec2 atlas_size; // x used for light probe mode total probes
uint ray_count;
@ -735,4 +749,153 @@ void main() {
imageStore(dest_light, ivec3(atlas_pos, params.atlas_slice), c);
#endif
#ifdef MODE_DENOISE
// Joint Non-local means (JNLM) denoiser.
//
// Based on YoctoImageDenoiser's JNLM implementation with corrections from "Nonlinearly Weighted First-order Regression for Denoising Monte Carlo Renderings".
//
// <https://github.com/ManuelPrandini/YoctoImageDenoiser/blob/06e19489dd64e47792acffde536393802ba48607/libs/yocto_extension/yocto_extension.cpp#L207>
// <https://benedikt-bitterli.me/nfor/nfor.pdf>
//
// MIT License
//
// Copyright (c) 2020 ManuelPrandini
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//
// Most of the constants below have been hand-picked to fit the common scenarios lightmaps
// are generated with, but they can be altered freely to experiment and achieve better results.
// Half the size of the patch window around each pixel that is weighted to compute the denoised pixel.
// A value of 1 represents a 3x3 window, a value of 2 a 5x5 window, etc.
const int HALF_PATCH_WINDOW = 4;
// Half the size of the search window around each pixel that is denoised and weighted to compute the denoised pixel.
const int HALF_SEARCH_WINDOW = 10;
// For all of the following sigma values, smaller values will give less weight to pixels that have a bigger distance
// in the feature being evaluated. Therefore, smaller values are likely to cause more noise to appear, but will also
// cause less features to be erased in the process.
// Controls how much the spatial distance of the pixels influences the denoising weight.
const float SIGMA_SPATIAL = denoise_params.spatial_bandwidth;
// Controls how much the light color distance of the pixels influences the denoising weight.
const float SIGMA_LIGHT = denoise_params.light_bandwidth;
// Controls how much the albedo color distance of the pixels influences the denoising weight.
const float SIGMA_ALBEDO = denoise_params.albedo_bandwidth;
// Controls how much the normal vector distance of the pixels influences the denoising weight.
const float SIGMA_NORMAL = denoise_params.normal_bandwidth;
// Strength of the filter. The original paper recommends values around 10 to 15 times the Sigma parameter.
const float FILTER_VALUE = denoise_params.filter_strength * SIGMA_LIGHT;
// Formula constants.
const int PATCH_WINDOW_DIMENSION = (HALF_PATCH_WINDOW * 2 + 1);
const int PATCH_WINDOW_DIMENSION_SQUARE = (PATCH_WINDOW_DIMENSION * PATCH_WINDOW_DIMENSION);
const float TWO_SIGMA_SPATIAL_SQUARE = 2.0f * SIGMA_SPATIAL * SIGMA_SPATIAL;
const float TWO_SIGMA_LIGHT_SQUARE = 2.0f * SIGMA_LIGHT * SIGMA_LIGHT;
const float TWO_SIGMA_ALBEDO_SQUARE = 2.0f * SIGMA_ALBEDO * SIGMA_ALBEDO;
const float TWO_SIGMA_NORMAL_SQUARE = 2.0f * SIGMA_NORMAL * SIGMA_NORMAL;
const float FILTER_SQUARE_TWO_SIGMA_LIGHT_SQUARE = FILTER_VALUE * FILTER_VALUE * TWO_SIGMA_LIGHT_SQUARE;
const float EPSILON = 1e-6f;
#ifdef USE_SH_LIGHTMAPS
const uint slice_count = 4;
const uint slice_base = params.atlas_slice * slice_count;
#else
const uint slice_count = 1;
const uint slice_base = params.atlas_slice;
#endif
for (uint i = 0; i < slice_count; i++) {
uint lightmap_slice = slice_base + i;
vec3 denoised_rgb = vec3(0.0f);
vec4 input_light = texelFetch(sampler2DArray(source_light, linear_sampler), ivec3(atlas_pos, lightmap_slice), 0);
vec3 input_albedo = texelFetch(sampler2DArray(albedo_tex, linear_sampler), ivec3(atlas_pos, params.atlas_slice), 0).rgb;
vec3 input_normal = texelFetch(sampler2DArray(source_normal, linear_sampler), ivec3(atlas_pos, params.atlas_slice), 0).xyz;
if (length(input_normal) > EPSILON) {
// Compute the denoised pixel if the normal is valid.
float sum_weights = 0.0f;
vec3 input_rgb = input_light.rgb;
for (int search_y = -HALF_SEARCH_WINDOW; search_y <= HALF_SEARCH_WINDOW; search_y++) {
for (int search_x = -HALF_SEARCH_WINDOW; search_x <= HALF_SEARCH_WINDOW; search_x++) {
ivec2 search_pos = atlas_pos + ivec2(search_x, search_y);
vec3 search_rgb = texelFetch(sampler2DArray(source_light, linear_sampler), ivec3(search_pos, lightmap_slice), 0).rgb;
vec3 search_albedo = texelFetch(sampler2DArray(albedo_tex, linear_sampler), ivec3(search_pos, params.atlas_slice), 0).rgb;
vec3 search_normal = texelFetch(sampler2DArray(source_normal, linear_sampler), ivec3(search_pos, params.atlas_slice), 0).xyz;
float patch_square_dist = 0.0f;
for (int offset_y = -HALF_PATCH_WINDOW; offset_y <= HALF_PATCH_WINDOW; offset_y++) {
for (int offset_x = -HALF_PATCH_WINDOW; offset_x <= HALF_PATCH_WINDOW; offset_x++) {
ivec2 offset_input_pos = atlas_pos + ivec2(offset_x, offset_y);
ivec2 offset_search_pos = search_pos + ivec2(offset_x, offset_y);
vec3 offset_input_rgb = texelFetch(sampler2DArray(source_light, linear_sampler), ivec3(offset_input_pos, lightmap_slice), 0).rgb;
vec3 offset_search_rgb = texelFetch(sampler2DArray(source_light, linear_sampler), ivec3(offset_search_pos, lightmap_slice), 0).rgb;
vec3 offset_delta_rgb = offset_input_rgb - offset_search_rgb;
patch_square_dist += dot(offset_delta_rgb, offset_delta_rgb) - TWO_SIGMA_LIGHT_SQUARE;
}
}
patch_square_dist = max(0.0f, patch_square_dist / (3.0f * PATCH_WINDOW_DIMENSION_SQUARE));
float weight = 1.0f;
// Ignore weight if search position is out of bounds.
weight *= step(0, search_pos.x) * step(search_pos.x, params.atlas_size.x - 1);
weight *= step(0, search_pos.y) * step(search_pos.y, params.atlas_size.y - 1);
// Ignore weight if normal is zero length.
weight *= step(EPSILON, length(search_normal));
// Weight with pixel distance.
vec2 pixel_delta = vec2(search_x, search_y);
float pixel_square_dist = dot(pixel_delta, pixel_delta);
weight *= exp(-pixel_square_dist / TWO_SIGMA_SPATIAL_SQUARE);
// Weight with patch.
weight *= exp(-patch_square_dist / FILTER_SQUARE_TWO_SIGMA_LIGHT_SQUARE);
// Weight with albedo.
vec3 albedo_delta = input_albedo - search_albedo;
float albedo_square_dist = dot(albedo_delta, albedo_delta);
weight *= exp(-albedo_square_dist / TWO_SIGMA_ALBEDO_SQUARE);
// Weight with normal.
vec3 normal_delta = input_normal - search_normal;
float normal_square_dist = dot(normal_delta, normal_delta);
weight *= exp(-normal_square_dist / TWO_SIGMA_NORMAL_SQUARE);
denoised_rgb += weight * search_rgb;
sum_weights += weight;
}
}
denoised_rgb /= sum_weights;
} else {
// Ignore pixels where the normal is empty, just copy the light color.
denoised_rgb = input_light.rgb;
}
imageStore(dest_light, ivec3(atlas_pos, lightmap_slice), vec4(denoised_rgb, input_light.a));
}
#endif
}