raytracing direct lighting

VZout · Dec 3, 2019 · d033c71 · d033c71
1 parent 823f458
commit d033c71
Show file tree

Hide file tree

Showing 9 changed files with 209 additions and 46 deletions.
diff --git a/src/engine_registry.cpp b/src/engine_registry.cpp
@@ -292,7 +292,7 @@ REGISTER(root_signatures::generate_brdf_lut, RootSignatureRegistry)({
 REGISTER(root_signatures::raytracing, RootSignatureRegistry)({
   .m_parameters = []() -> decltype(RootSignatureDesc::m_parameters)
   {
-	  decltype(RootSignatureDesc::m_parameters) params(7);
+	  decltype(RootSignatureDesc::m_parameters) params(11);
 	  params[0].binding = 0; // acceleration structure
 	  params[0].descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_NV;
 	  params[0].descriptorCount = 1;
@@ -328,6 +328,26 @@ REGISTER(root_signatures::raytracing, RootSignatureRegistry)({
 	  params[6].descriptorCount = 1;
 	  params[6].stageFlags = VK_SHADER_STAGE_CLOSEST_HIT_BIT_NV;
 	  params[6].pImmutableSamplers = nullptr;
+	  params[7].binding = 7; // materials
+	  params[7].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+	  params[7].descriptorCount = 1;
+	  params[7].stageFlags = VK_SHADER_STAGE_CLOSEST_HIT_BIT_NV;
+	  params[7].pImmutableSamplers = nullptr;
+	  params[8].binding = 8; // materials
+	  params[8].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
+	  params[8].descriptorCount = gfx::settings::max_num_rtx_textures;
+	  params[8].stageFlags = VK_SHADER_STAGE_CLOSEST_HIT_BIT_NV;
+	  params[8].pImmutableSamplers = nullptr;
+	  params[9].binding = 9; // skybox
+	  params[9].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
+	  params[9].descriptorCount = gfx::settings::max_num_rtx_textures;
+	  params[9].stageFlags = VK_SHADER_STAGE_MISS_BIT_NV;
+	  params[9].pImmutableSamplers = nullptr;
+	  params[10].binding = 10; // pbrlut
+	  params[10].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
+	  params[10].descriptorCount = gfx::settings::max_num_rtx_textures;
+	  params[10].stageFlags = VK_SHADER_STAGE_CLOSEST_HIT_BIT_NV;
+	  params[10].pImmutableSamplers = nullptr;
 	  return params;
   }(),
 });

diff --git a/src/graphics/gfx_settings.hpp b/src/graphics/gfx_settings.hpp
@@ -36,4 +36,5 @@ namespace gfx::settings
 	static const std::uint32_t max_lights = 50;
 	static const std::uint32_t max_render_batch_size = 700;
 	static const std::uint32_t max_num_rtx_materials = 2000;
+	static const std::uint32_t max_num_rtx_textures = 500;
 }
diff --git a/src/graphics/vk_texture_pool.cpp b/src/graphics/vk_texture_pool.cpp
@@ -120,3 +120,15 @@ std::vector<gfx::StagingTexture*> gfx::VkTexturePool::GetTextures(std::vector<st
 
 	return textures;
 }
+
+std::vector<gfx::StagingTexture*> gfx::VkTexturePool::GetAllTexturesPadded(std::uint32_t num)
+{
+	std::vector<gfx::StagingTexture*> all_textures;
+	for (auto& texture : m_staged_textures)
+	{
+		all_textures.push_back(texture.second);
+	}
+	all_textures.resize(num, m_staged_textures.begin()->second);
+
+	return all_textures;
+}
diff --git a/src/graphics/vk_texture_pool.hpp b/src/graphics/vk_texture_pool.hpp
@@ -26,6 +26,7 @@ namespace gfx
 		void Stage(gfx::CommandList* command_list) final;
 		void PostStage() final;
 		std::vector<gfx::StagingTexture*> GetTextures(std::vector<std::uint32_t> texture_handles) final;
+		std::vector<gfx::StagingTexture*> GetAllTexturesPadded(std::uint32_t num);
 
 	private:
 		void Load_Impl(TextureData const & data, std::uint32_t id, bool mipmap, bool srgb) final;

diff --git a/src/render_tasks/vk_build_acceleration_structures_task.hpp b/src/render_tasks/vk_build_acceleration_structures_task.hpp
@@ -33,14 +33,23 @@ namespace tasks
 		std::uint32_t m_index_offset;
 	};
 
+	struct RaytracingMaterial
+	{
+		std::uint32_t m_albedo_texture;
+		std::uint32_t m_normal_texture;
+		std::uint32_t m_roughness_texture;
+	};
+
 	struct BuildASData
 	{
 		gfx::AccelerationStructure* m_tlas;
 		std::vector<gfx::AccelerationStructure*> m_blasses;
 		gfx::GPUBuffer* m_scratch_buffer;
 
 		std::vector<RaytracingOffset> m_offsets;
+		std::vector<RaytracingMaterial> m_materials;
 		gfx::GPUBuffer* m_offsets_buffer;
+		gfx::GPUBuffer* m_materials_buffer;
 
 		bool m_should_build = true;
 	};
@@ -54,8 +63,11 @@ namespace tasks
 
 			const auto scratch_size = 65536*20;
 
-			data.m_offsets.resize(gfx::settings::max_num_rtx_materials, RaytracingOffset{0, 0});
+			data.m_offsets.resize(gfx::settings::max_num_rtx_materials);
+			data.m_materials.resize(gfx::settings::max_num_rtx_materials);
+
 			data.m_offsets_buffer = new gfx::GPUBuffer(rs.GetContext(), std::nullopt, gfx::settings::max_num_rtx_materials * sizeof(RaytracingOffset), gfx::enums::BufferUsageFlag::RAYTRACING_STORAGE);
+			data.m_materials_buffer = new gfx::GPUBuffer(rs.GetContext(), std::nullopt, gfx::settings::max_num_rtx_materials * sizeof(RaytracingMaterial), gfx::enums::BufferUsageFlag::RAYTRACING_STORAGE);
 			//data.m_scratch_buffer = new gfx::GPUBuffer(rs.GetContext(), std::nullopt, scratch_size, gfx::enums::BufferUsageFlag::RAYTRACING);
 			if (!resize)
 			{
@@ -110,6 +122,12 @@ namespace tasks
 						offset.m_index_offset = geom_desc.m_indices_offset;
 						data.m_offsets[material_id] = offset;
 
+						RaytracingMaterial material;
+						material.m_albedo_texture = batch.m_material_handles[i].m_albedo_texture_handle;
+						material.m_normal_texture = batch.m_material_handles[i].m_normal_texture_handle;
+						material.m_roughness_texture = batch.m_material_handles[i].m_roughness_texture_handle;
+						data.m_materials[material_id] = material;
+
 						gfx::InstanceDesc new_instance;
 						new_instance.m_transform = (glm::mat3x4)(model_mat);
 						new_instance.m_blas = blas;
@@ -126,6 +144,11 @@ namespace tasks
 			data.m_offsets_buffer->Update(data.m_offsets.data(), data.m_offsets.size() * sizeof(RaytracingOffset));
 			data.m_offsets_buffer->Unmap();
 
+			// Fill the materials buffer
+			data.m_materials_buffer->Map();
+			data.m_materials_buffer->Update(data.m_materials.data(), data.m_materials.size() * sizeof(RaytracingMaterial));
+			data.m_materials_buffer->Unmap();
+
 			// Create TLAS
 			data.m_tlas = new gfx::AccelerationStructure(context);
 			data.m_tlas->CreateTopLevel(cmd_list, blas_instances);
@@ -145,6 +168,7 @@ namespace tasks
 				}
 
 				delete data.m_tlas;
+				delete data.m_materials_buffer;
 				delete data.m_offsets_buffer;
 			}
 		}

diff --git a/src/render_tasks/vk_raytracing_task.hpp b/src/render_tasks/vk_raytracing_task.hpp
@@ -32,7 +32,11 @@ namespace tasks
 	struct RaytracingData
 	{
 		std::uint32_t m_tlas_set;
+		std::uint32_t m_skybox_set;
+		std::uint32_t m_brdf_set;
 		std::uint32_t m_offsets_set;
+		std::uint32_t m_materials_set;
+		std::uint32_t m_textures_set;
 		std::uint32_t m_uav_target_set;
 		gfx::DescriptorHeap* m_gbuffer_heap;
 
@@ -65,7 +69,7 @@ namespace tasks
 
 			gfx::DescriptorHeap::Desc descriptor_heap_desc = {};
 			descriptor_heap_desc.m_versions = 1;
-			descriptor_heap_desc.m_num_descriptors = 3;
+			descriptor_heap_desc.m_num_descriptors = 7;
 			data.m_gbuffer_heap = new gfx::DescriptorHeap(rs.GetContext(), descriptor_heap_desc);
 			data.m_uav_target_set = data.m_gbuffer_heap->CreateUAVSetFromRT(render_target, 0, data.m_root_sig, 1, 0, input_sampler_desc);
 
@@ -81,6 +85,34 @@ namespace tasks
 
 			data.m_hitgroup_shader_table = new gfx::ShaderTable(rs.GetContext(), 1);
 			data.m_hitgroup_shader_table->AddShaderRecord(data.m_pipeline, 2, 1);*/
+
+			gfx::SamplerDesc skybox_sampler_desc
+			{
+				.m_filter = gfx::enums::TextureFilter::FILTER_ANISOTROPIC,
+				.m_address_mode = gfx::enums::TextureAddressMode::TAM_WRAP,
+				.m_border_color = gfx::enums::BorderColor::BORDER_WHITE,
+			};
+
+			gfx::SamplerDesc lut_sampler_desc
+			{
+				.m_filter = gfx::enums::TextureFilter::FILTER_LINEAR,
+				.m_address_mode = gfx::enums::TextureAddressMode::TAM_CLAMP,
+				.m_border_color = gfx::enums::BorderColor::BORDER_WHITE,
+			};
+
+			// Skybox
+			if (fg.HasTask<GenerateCubemapData>())
+			{
+				auto skybox_rt = fg.GetPredecessorRenderTarget<GenerateCubemapData>();
+				data.m_skybox_set = data.m_gbuffer_heap->CreateSRVSetFromRT(skybox_rt, data.m_root_sig, 9, 0, false, skybox_sampler_desc);
+			}
+
+			// BRDF Lut
+			if (fg.HasTask<GenerateBRDFLutData>())
+			{
+				auto brdf_rt = fg.GetPredecessorRenderTarget<GenerateBRDFLutData>();
+				data.m_brdf_set = data.m_gbuffer_heap->CreateSRVSetFromRT(brdf_rt, data.m_root_sig, 10, 0, false, lut_sampler_desc);
+			}
 		}
 
 		inline void ExecuteRaytracingTask(Renderer& rs, fg::FrameGraph& fg, sg::SceneGraph& sg, fg::RenderTaskHandle handle)
@@ -90,6 +122,7 @@ namespace tasks
 			auto render_target = fg.GetRenderTarget(handle);
 
 			auto model_pool = static_cast<gfx::VkModelPool*>(rs.GetModelPool());
+			auto texture_pool = static_cast<gfx::VkTexturePool*>(rs.GetTexturePool());
 			auto camera_pool = static_cast<gfx::VkConstantBufferPool*>(sg.GetInverseCameraConstantBufferPool());
 			auto camera_handle = sg.m_inverse_camera_cb_handles[0].m_value;
 			auto light_pool = static_cast<gfx::VkConstantBufferPool*>(sg.GetLightConstantBufferPool());
@@ -100,6 +133,8 @@ namespace tasks
 				auto as_build_data = fg.GetPredecessorData<BuildASData>();
 				data.m_tlas_set = data.m_gbuffer_heap->CreateSRVFromAS(as_build_data.m_tlas, data.m_root_sig, 0, 0);
 				data.m_offsets_set = data.m_gbuffer_heap->CreateSRVFromCB(as_build_data.m_offsets_buffer, data.m_root_sig, 6, 0, false);
+				data.m_materials_set = data.m_gbuffer_heap->CreateSRVFromCB(as_build_data.m_materials_buffer, data.m_root_sig, 7, 0, false);
+				data.m_textures_set = data.m_gbuffer_heap->CreateSRVSetFromTexture(texture_pool->GetAllTexturesPadded(gfx::settings::max_num_rtx_textures), data.m_root_sig, 8, 0);
 
 				data.m_first_execute = false;
 			}
@@ -114,6 +149,10 @@ namespace tasks
 				{ model_pool->m_heap, model_pool->m_big_vb_desc_set_id },
 				{ model_pool->m_heap, model_pool->m_big_ib_desc_set_id },
 				{ data.m_gbuffer_heap, data.m_offsets_set },
+				{ data.m_gbuffer_heap, data.m_materials_set },
+				{ data.m_gbuffer_heap, data.m_textures_set },
+				{ data.m_gbuffer_heap, data.m_skybox_set },
+				{ data.m_gbuffer_heap, data.m_brdf_set },
 			};
 
 			cmd_list->BindPipelineState(data.m_pipeline);

diff --git a/src/shaders/composition.comp b/src/shaders/composition.comp
@@ -120,12 +120,7 @@ void main()
 	float subsurface_power = 0;
 	float thickness = decoded_ao_thick.g;
 
-
-#ifdef PBR_PLUS
     float roughness = clamp(albedo_roughness.a, MIN_PERCEPTUAL_ROUGHNESS, 1.f);
-#else
-    float roughness = albedo_roughness.a;
-#endif
     float metallic = normal_metallic.a;
 
     mat4 inv_view = inverse(camera.view);
@@ -146,7 +141,6 @@ void main()
 	float NdotV = max(dot(N, V), MIN_N_DOT_V);
     vec2 dfg = texture(t_brdf_lut, vec2(NdotV, roughness)).gr; // lut
 
-#ifdef PBR_PLUS
     vec3 diffuse_color = ComputeDiffuseColor(albedo, metallic);
 
     vec3 F0;
@@ -197,7 +191,6 @@ void main()
     float ibl_luminance = 1; // TODO figure this IBLluminance out.
 
     vec3 ibl_color = (diff + spec) * ibl_luminance;
-#endif
 
     uint num_lights = 1; //Light count is stored in 30 upper-bits of first light
     vec3 lighting = vec3(0);
@@ -213,7 +206,6 @@ void main()
 		const vec3 pos_to_light = light_pos - world_pos;
         vec3 L = normalize(pos_to_light);
 
-#ifdef PBR_PLUS
 		float sqr_falloff = lights.lights[i].m_radius * lights.lights[i].m_radius;
 		float sqr_inv_radius = sqr_falloff > 0.0f ? (1 / sqr_falloff) : 0;
 
@@ -248,27 +240,9 @@ void main()
 			lighting += BRDF(L, V, N, geometric_normal, metallic, roughness, diffuse_color, light_color, F0, energy_compensation, attenuation, 1.0f, anisotropy, anisotropic_t, anisotropic_b, clear_coat, cc_roughness);
 		}
 #endif
-
-#else
-        lighting += BRDF(L, V, N, metallic, roughness, albedo, light_color);
-#endif
     }
 
-    #ifdef PBR_PLUS
     vec3 color = ibl_color + lighting;
-    #else
-    vec3 diffuse = irradiance * albedo;
-    vec3 F = F_SchlickRoughness(max(dot(N, V), 0.0), metallic, albedo, roughness);
-    vec3 reflection = GetEnvReflection(R, roughness);
-    vec3 specular = reflection * (F * dfg.x + dfg.y);
-
-    vec3 kD = 1.0 - F;
-    kD *= 1.0 - metallic;
-    vec3 color = (kD * diffuse + specular) + lighting;
-
-    color *= ao;
-    #endif
-
 
     // Skybox
     if (world_pos.z == 0)