rtx: pass normals to denoiser

This commit is contained in:
Ivan 'provod' Avdeev 2021-11-12 10:42:56 -08:00
parent fdad4fa016
commit 499adafaed
5 changed files with 65 additions and 4 deletions

View File

@ -10,6 +10,7 @@ layout(set = 0, binding = 1, rgba8) uniform image2D src_base_color;
layout(set = 0, binding = 2, rgba16f) uniform image2D src_diffuse_gi;
layout(set = 0, binding = 3, rgba16f) uniform image2D src_specular;
layout(set = 0, binding = 4, rgba16f) uniform image2D src_additive;
layout(set = 0, binding = 5, rgba16f) uniform image2D src_normals;
// Blatantly copypasted from https://www.shadertoy.com/view/XsGfWV
vec3 aces_tonemap(vec3 color){
@ -32,6 +33,12 @@ vec3 aces_tonemap(vec3 color){
float normpdf2(in float x2, in float sigma) { return 0.39894*exp(-0.5*x2/(sigma*sigma))/sigma; }
float normpdf(in float x, in float sigma) { return normpdf2(x*x, sigma); }
void readNormals(ivec2 uv, out vec3 geometry_normal, out vec3 shading_normal) {
const vec4 n = imageLoad(src_normals, uv);
geometry_normal = vec3(n.xy, sqrt(1. - n.x*n.x - n.y*n.y));
shading_normal = vec3(n.zw, sqrt(1. - n.z*n.z - n.w*n.w));
}
void main() {
ivec2 res = ivec2(imageSize(src_base_color));
ivec2 pix = ivec2(gl_GlobalInvocationID);
@ -49,6 +56,11 @@ void main() {
//imageStore(dest, pix, vec4(aces_tonemap(imageLoad(src_diffuse_gi, pix).rgb), 0.)); return;
//imageStore(dest, pix, vec4(aces_tonemap(imageLoad(src_specular, pix).rgb), 0.)); return;
vec3 geometry_normal, shading_normal;
readNormals(pix, geometry_normal, shading_normal);
//imageStore(dest, pix, vec4(.5 + geometry_normal * .5, 0.)); return;
// const uint mi = uint(material_index);
// imageStore(dest, pix, vec4(rand3_f01(uvec3(mi,mi+1,mi+2)), 0.));
// return;
@ -66,6 +78,14 @@ void main() {
if (c.a != material_index)
continue;
vec3 sample_geometry_normal, sample_shading_normal;
readNormals(p, sample_geometry_normal, sample_shading_normal);
if ( dot(sample_geometry_normal, geometry_normal) < .99 )
continue;
// TODO bilaterally filter shading normals
const float sigma = KERNEL_SIZE / 2.;
const float scale = normpdf(x, sigma) * normpdf(y, sigma);
colour += scale * imageLoad(src_diffuse_gi, p).rgb;

View File

@ -27,6 +27,7 @@ layout(set = 0, binding = 0, rgba8) uniform image2D out_image_base_color;
layout(set = 0, binding = 9, rgba16f) uniform image2D out_image_diffuse_gi;
layout(set = 0, binding = 10, rgba16f) uniform image2D out_image_specular;
layout(set = 0, binding = 11, rgba16f) uniform image2D out_image_additive;
layout(set = 0, binding = 12, rgba16f) uniform image2D out_image_normals;
layout(set = 0, binding = 1) uniform accelerationStructureEXT tlas;
layout(set = 0, binding = 2) uniform UBO {
@ -405,15 +406,21 @@ void main() {
material.baseColor = payload_opaque.base_color;
material.metalness = 0.f; // TODO
material.emissive = payload_opaque.emissive;
//material.roughness = uintToFloat01(xxhash32(uvec3(floor(payload_opaque.hit_pos_t.xyz/64.)))); //payload_opaque.roughness;
material.roughness = payload_opaque.roughness;
// material.roughness = uintToFloat01(xxhash32(uvec3(abs(floor(payload_opaque.hit_pos_t.xyz/64.)))));
// material.metalness = step(.5, xxhash32(uvec3(abs(floor(payload_opaque.hit_pos_t.xyz/32.)))));
const vec3 shadingNormal = payload_opaque.normal;
const vec3 geometryNormal = payload_opaque.geometry_normal;
if (bounce == 0) { //brdfType == SPECULAR_TYPE)
out_additive = payload_opaque.emissive + additive;
additive = vec3(0.);
out_material_index = float(payload_opaque.material_index);
imageStore(out_image_base_color, ivec2(gl_LaunchIDEXT.xy), vec4(payload_opaque.base_color, 0.));
imageStore(out_image_normals, ivec2(gl_LaunchIDEXT.xy), vec4(geometryNormal.xy, shadingNormal.xy));
payload_opaque.base_color = vec3(1.);
}
@ -428,8 +435,6 @@ void main() {
vec3 prev_throughput = throughput;
#endif
const vec3 shadingNormal = payload_opaque.normal;
const vec3 geometryNormal = payload_opaque.geometry_normal;
const vec3 V = -direction;
if (material.metalness == 1.0f && material.roughness == 0.0f) {
// Fast path for mirrors

View File

@ -12,6 +12,7 @@ enum {
DenoiserBinding_Source_DiffuseGI = 2,
DenoiserBinding_Source_Specular = 3,
DenoiserBinding_Source_Additive = 4,
DenoiserBinding_Source_Normals = 5,
DenoiserBinding_COUNT
};
@ -73,6 +74,13 @@ static void createLayouts( void ) {
.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
};
g_denoiser.desc_bindings[DenoiserBinding_Source_Normals] = (VkDescriptorSetLayoutBinding){
.binding = DenoiserBinding_Source_Normals,
.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
.descriptorCount = 1,
.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
};
VK_DescriptorsCreate(&g_denoiser.descriptors);
}
@ -139,6 +147,11 @@ void XVK_DenoiserDenoise( const xvk_denoiser_args_t* args ) {
.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
};
g_denoiser.desc_values[DenoiserBinding_Source_Normals].image = (VkDescriptorImageInfo){
.sampler = VK_NULL_HANDLE,
.imageView = args->src.normals_view,
.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
};
g_denoiser.desc_values[DenoiserBinding_DestImage].image = (VkDescriptorImageInfo){
.sampler = VK_NULL_HANDLE,

View File

@ -16,6 +16,7 @@ typedef struct {
VkImageView diffuse_gi_view;
VkImageView specular_view;
VkImageView additive_view;
VkImageView normals_view;
} src;
VkImageView dst_view;

View File

@ -74,7 +74,7 @@ enum {
RayDescBinding_Dest_ImageDiffuseGI = 9,
RayDescBinding_Dest_ImageSpecular = 10,
RayDescBinding_Dest_ImageAdditive = 11,
//RayDescBinding_Dest_ImageNormal = 12,
RayDescBinding_Dest_ImageNormals = 12,
RayDescBinding_COUNT
};
@ -85,6 +85,7 @@ typedef struct {
vk_image_t diffuse_gi;
vk_image_t specular;
vk_image_t additive;
vk_image_t normals;
} xvk_ray_frame_images_t;
static struct {
@ -651,6 +652,13 @@ static void updateDescriptors( VkCommandBuffer cmdbuf, const vk_ray_frame_render
.imageView = frame_dst->additive.view,
.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
};
g_rtx.desc_values[RayDescBinding_Dest_ImageNormals].image = (VkDescriptorImageInfo){
.sampler = VK_NULL_HANDLE,
.imageView = frame_dst->normals.view,
.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
};
VK_DescriptorsWrite(&g_rtx.descriptors);
}
@ -660,6 +668,7 @@ static qboolean rayTrace( VkCommandBuffer cmdbuf, const xvk_ray_frame_images_t *
X(diffuse_gi) \
X(specular) \
X(additive) \
X(normals) \
// 4. Barrier for TLAS build and dest image layout transfer
{
@ -1009,6 +1018,7 @@ LIST_GBUFFER_IMAGES(GBUFFER_READ_BARRIER)
.diffuse_gi_view = current_frame->diffuse_gi.view,
.specular_view = current_frame->specular.view,
.additive_view = current_frame->additive.view,
.normals_view = current_frame->normals.view,
},
.dst_view = current_frame->denoised.view,
};
@ -1073,6 +1083,12 @@ static void createLayouts( void ) {
.descriptorCount = 1,
.stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
};
g_rtx.desc_bindings[RayDescBinding_Dest_ImageNormals] = (VkDescriptorSetLayoutBinding){
.binding = RayDescBinding_Dest_ImageNormals,
.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
.descriptorCount = 1,
.stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
};
g_rtx.desc_bindings[RayDescBinding_Dest_ImageBaseColor] = (VkDescriptorSetLayoutBinding){
.binding = RayDescBinding_Dest_ImageBaseColor,
@ -1244,6 +1260,11 @@ qboolean VK_RayInit( void )
g_rtx.frames[i].additive = VK_ImageCreate(FRAME_WIDTH, FRAME_HEIGHT, VK_FORMAT_R16G16B16A16_SFLOAT,
VK_IMAGE_TILING_OPTIMAL, VK_IMAGE_USAGE_STORAGE_BIT);
SET_DEBUG_NAMEF(g_rtx.frames[i].additive.image, VK_OBJECT_TYPE_IMAGE, "rtx frames[%d] additive", i);
// TODO make sure this format and usage is suppported
g_rtx.frames[i].normals = VK_ImageCreate(FRAME_WIDTH, FRAME_HEIGHT, VK_FORMAT_R16G16B16A16_SNORM,
VK_IMAGE_TILING_OPTIMAL, VK_IMAGE_USAGE_STORAGE_BIT);
SET_DEBUG_NAMEF(g_rtx.frames[i].normals.image, VK_OBJECT_TYPE_IMAGE, "rtx frames[%d] normals", i);
}
if (vk_core.debug) {
@ -1264,6 +1285,7 @@ void VK_RayShutdown( void ) {
VK_ImageDestroy(&g_rtx.frames[i].diffuse_gi);
VK_ImageDestroy(&g_rtx.frames[i].specular);
VK_ImageDestroy(&g_rtx.frames[i].additive);
VK_ImageDestroy(&g_rtx.frames[i].normals);
}
vkDestroyPipeline(vk_core.device, g_rtx.pipeline, NULL);