1
0
mirror of https://github.com/godotengine/godot.git synced 2025-11-10 13:00:37 +00:00

Replace OIDN denoiser with a JNLM denoiser compute shader implementation.

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

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;
}
}