Add an ambient occlusion pass

CMakeLists.txt

@@ -178,11 +178,14 @@ set( SOURCES
 	src/lightmap/glsl_rchit_bounce.h
 	src/lightmap/glsl_rchit_light.h
 	src/lightmap/glsl_rchit_sun.h
+	src/lightmap/glsl_rchit_ambient.h
 	src/lightmap/glsl_rgen_bounce.h
 	src/lightmap/glsl_rgen_light.h
+	src/lightmap/glsl_rgen_ambient.h
 	src/lightmap/glsl_rmiss_bounce.h
 	src/lightmap/glsl_rmiss_light.h
 	src/lightmap/glsl_rmiss_sun.h
+	src/lightmap/glsl_rmiss_ambient.h
 	src/lightmap/cpuraytracer.cpp
 	src/lightmap/cpuraytracer.h
 	src/math/mat.cpp

src/lightmap/glsl_rchit_ambient.h

@@ -0,0 +1,33 @@
+static const char* glsl_rchit_ambient = R"glsl(
+
+#version 460
+#extension GL_EXT_ray_tracing : require
+
+struct hitPayload
+{
+	vec3 hitPosition;
+	float hitAttenuation;
+	int hitSurfaceIndex;
+};
+
+struct SurfaceInfo
+{
+	vec3 Normal;
+	float EmissiveDistance;
+	vec3 EmissiveColor;
+	float EmissiveIntensity;
+	float Sky;
+	float Padding0, Padding1, Padding2;
+};
+
+layout(location = 0) rayPayloadInEXT hitPayload payload;
+
+layout(set = 0, binding = 5) buffer SurfaceIndexBuffer { int surfaceIndices[]; };
+layout(set = 0, binding = 6) buffer SurfaceBuffer { SurfaceInfo surfaces[]; };
+
+void main()
+{
+	payload.hitAttenuation = gl_HitTEXT;
+}
+
+)glsl";

src/lightmap/glsl_rgen_ambient.h

@@ -0,0 +1,111 @@
+static const char* glsl_rgen_ambient = R"glsl(
+
+#version 460
+#extension GL_EXT_ray_tracing : require
+
+struct hitPayload
+{
+	vec3 hitPosition;
+	float hitAttenuation;
+	int hitSurfaceIndex;
+};
+
+layout(location = 0) rayPayloadEXT hitPayload payload;
+
+layout(set = 0, binding = 0) uniform accelerationStructureEXT acc;
+layout(set = 0, binding = 1, rgba32f) uniform image2D startpositions;
+layout(set = 0, binding = 2, rgba32f) uniform image2D positions;
+layout(set = 0, binding = 3, rgba32f) uniform image2D outputs;
+
+layout(set = 0, binding = 4) uniform Uniforms
+{
+	uint SampleIndex;
+	uint SampleCount;
+	uint PassType;
+	uint Padding0;
+	vec3 SunDir;
+	float SampleDistance;
+	vec3 SunColor;
+	float SunIntensity;
+	vec3 HemisphereVec;
+	float Padding1;
+};
+
+struct SurfaceInfo
+{
+	vec3 Normal;
+	float EmissiveDistance;
+	vec3 EmissiveColor;
+	float EmissiveIntensity;
+	float Sky;
+	float Padding0, Padding1, Padding2;
+};
+
+layout(set = 0, binding = 6) buffer SurfaceBuffer { SurfaceInfo surfaces[]; };
+
+layout(push_constant) uniform PushConstants
+{
+	uint LightStart;
+	uint LightEnd;
+	ivec2 pushPadding;
+};
+
+vec2 Hammersley(uint i, uint N);
+float RadicalInverse_VdC(uint bits);
+
+void main()
+{
+	ivec2 texelPos = ivec2(gl_LaunchIDEXT.xy);
+
+	vec4 data0 = imageLoad(startpositions, texelPos);
+	vec4 incoming = imageLoad(outputs, texelPos);
+
+	if (PassType == 1)
+		incoming.rgb = vec3(1.0); // For debugging
+
+	vec3 origin = data0.xyz;
+	int surfaceIndex = int(data0.w);
+
+	if (surfaceIndex >= 0)
+	{
+		const float minDistance = 0.05;
+		const float aoDistance = 100;
+
+		vec3 N = surfaces[surfaceIndex].Normal;
+		vec3 up = abs(N.x) < abs(N.y) ? vec3(1.0, 0.0, 0.0) : vec3(0.0, 1.0, 0.0);
+		vec3 tangent = normalize(cross(up, N));
+		vec3 bitangent = cross(N, tangent);
+
+		float ambience = 0.0f;
+		for (uint i = 0; i < SampleCount; i++)
+		{
+			vec2 Xi = Hammersley(i, SampleCount);
+			vec3 H = normalize(vec3(Xi.x * 2.0f - 1.0f, Xi.y * 2.0f - 1.0f, RadicalInverse_VdC(i) + 0.01f));
+			vec3 L = H.x * tangent + H.y * bitangent + H.z * N;
+			traceRayEXT(acc, gl_RayFlagsOpaqueEXT, 0xff, 3, 0, 3, origin, minDistance, L, 32768, 0);
+			ambience += clamp(payload.hitAttenuation / aoDistance, 0.0, 1.0);
+		}
+		ambience /= float(SampleCount);
+
+		incoming.rgb = incoming.rgb * ambience;
+	}
+
+	imageStore(outputs, texelPos, incoming);
+}
+
+vec2 Hammersley(uint i, uint N)
+{
+	return vec2(float(i) / float(N), RadicalInverse_VdC(i));
+}
+
+float RadicalInverse_VdC(uint bits)
+{
+	bits = (bits << 16u) | (bits >> 16u);
+	bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
+	bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
+	bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
+	bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
+	return float(bits) * 2.3283064365386963e-10f; // / 0x100000000
+}
+
+)glsl";

src/lightmap/glsl_rmiss_ambient.h

@@ -0,0 +1,20 @@
+static const char* glsl_rmiss_ambient = R"glsl(
+
+#version 460
+#extension GL_EXT_ray_tracing : require
+
+struct hitPayload
+{
+	vec3 hitPosition;
+	float hitAttenuation;
+	int hitSurfaceIndex;
+};
+
+layout(location = 0) rayPayloadInEXT hitPayload payload;
+
+void main()
+{
+	payload.hitAttenuation = 100000.0;
+}
+
+)glsl";

src/lightmap/gpuraytracer.cpp

@@ -15,12 +15,15 @@
 #include <mutex>
 #include "glsl_rgen_bounce.h"
 #include "glsl_rgen_light.h"
+#include "glsl_rgen_ambient.h"
 #include "glsl_rchit_bounce.h"
 #include "glsl_rchit_light.h"
 #include "glsl_rchit_sun.h"
+#include "glsl_rchit_ambient.h"
 #include "glsl_rmiss_bounce.h"
 #include "glsl_rmiss_light.h"
 #include "glsl_rmiss_sun.h"
+#include "glsl_rmiss_ambient.h"
 
 extern bool VKDebug;
 
@@ -132,6 +135,11 @@ void GPURaytracer::Raytrace(LevelMesh* level)
 				}
 			}
 
+			uniforms.PassType = 0;
+			uniforms.SampleIndex = 0;
+			uniforms.SampleCount = ambientSampleCount;
+			RunTrace(uniforms, rgenAmbientRegion);
+
 			EndTracing();
 			DownloadTasks(tasks.data() + startTask, numTasks);
 		}
@@ -652,12 +660,15 @@ void GPURaytracer::CreateShaders()
 {
 	rgenBounce = CompileRayGenShader(glsl_rgen_bounce, "rgen.bounce");
 	rgenLight = CompileRayGenShader(glsl_rgen_light, "rgen.light");
+	rgenAmbient = CompileRayGenShader(glsl_rgen_ambient, "rgen.ambient");
 	rchitBounce = CompileClosestHitShader(glsl_rchit_bounce, "rchit.bounce");
 	rchitLight = CompileClosestHitShader(glsl_rchit_light, "rchit.light");
 	rchitSun = CompileClosestHitShader(glsl_rchit_sun, "rchit.sun");
+	rchitAmbient = CompileClosestHitShader(glsl_rchit_ambient, "rchit.ambient");
 	rmissBounce = CompileMissShader(glsl_rmiss_bounce, "rmiss.bounce");
 	rmissLight = CompileMissShader(glsl_rmiss_light, "rmiss.light");
 	rmissSun = CompileMissShader(glsl_rmiss_sun, "rmiss.sun");
+	rmissAmbient = CompileMissShader(glsl_rmiss_ambient, "rmiss.ambient");
 }
 
 std::unique_ptr<VulkanShader> GPURaytracer::CompileRayGenShader(const char* code, const char* name)
@@ -733,20 +744,26 @@ void GPURaytracer::CreatePipeline()
 	builder.setMaxPipelineRayRecursionDepth(1);
 	builder.addShader(VK_SHADER_STAGE_RAYGEN_BIT_KHR, rgenBounce.get());
 	builder.addShader(VK_SHADER_STAGE_RAYGEN_BIT_KHR, rgenLight.get());
+	builder.addShader(VK_SHADER_STAGE_RAYGEN_BIT_KHR, rgenAmbient.get());
 	builder.addShader(VK_SHADER_STAGE_MISS_BIT_KHR, rmissBounce.get());
 	builder.addShader(VK_SHADER_STAGE_MISS_BIT_KHR, rmissLight.get());
 	builder.addShader(VK_SHADER_STAGE_MISS_BIT_KHR, rmissSun.get());
+	builder.addShader(VK_SHADER_STAGE_MISS_BIT_KHR, rmissAmbient.get());
 	builder.addShader(VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, rchitBounce.get());
 	builder.addShader(VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, rchitLight.get());
 	builder.addShader(VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, rchitSun.get());
+	builder.addShader(VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, rchitAmbient.get());
 	builder.addRayGenGroup(0);
 	builder.addRayGenGroup(1);
-	builder.addMissGroup(2);
+	builder.addRayGenGroup(2);
 	builder.addMissGroup(3);
 	builder.addMissGroup(4);
-	builder.addTrianglesHitGroup(5);
-	builder.addTrianglesHitGroup(6);
+	builder.addMissGroup(5);
+	builder.addMissGroup(6);
 	builder.addTrianglesHitGroup(7);
+	builder.addTrianglesHitGroup(8);
+	builder.addTrianglesHitGroup(9);
+	builder.addTrianglesHitGroup(10);
 	pipeline = builder.create(device.get());
 	pipeline->SetDebugName("pipeline");
 
@@ -760,9 +777,9 @@ void GPURaytracer::CreatePipeline()
 			return value;
 	};
 
-	VkDeviceSize raygenCount = 2;
-	VkDeviceSize missCount = 3;
-	VkDeviceSize hitCount = 3;
+	VkDeviceSize raygenCount = 3;
+	VkDeviceSize missCount = 4;
+	VkDeviceSize hitCount = 4;
 
 	VkDeviceSize handleSize = rtProperties.shaderGroupHandleSize;
 	VkDeviceSize handleSizeAligned = align_up(handleSize, rtProperties.shaderGroupHandleAlignment);
@@ -770,11 +787,6 @@ void GPURaytracer::CreatePipeline()
 	VkDeviceSize rgenStride = align_up(handleSizeAligned, rtProperties.shaderGroupBaseAlignment);
 	VkDeviceSize rgenSize = rgenStride * raygenCount;
 
-	rgenBounceRegion.stride = rgenStride;
-	rgenBounceRegion.size = rgenStride;
-	rgenLightRegion.stride = rgenStride;
-	rgenLightRegion.size = rgenStride;
-
 	missRegion.stride = handleSizeAligned;
 	missRegion.size = align_up(missCount * handleSizeAligned, rtProperties.shaderGroupBaseAlignment);
 
@@ -823,8 +835,14 @@ void GPURaytracer::CreatePipeline()
 	info.buffer = shaderBindingTable->buffer;
 	VkDeviceAddress sbtAddress = vkGetBufferDeviceAddress(device->device, &info);
 
-	rgenBounceRegion.deviceAddress = sbtAddress + rgenOffset;
-	rgenLightRegion.deviceAddress = sbtAddress + rgenOffset + rgenStride;
+	int i = 0;
+	for (VkStridedDeviceAddressRegionKHR* region : { &rgenBounceRegion, &rgenLightRegion, &rgenAmbientRegion })
+	{
+		region->stride = rgenStride;
+		region->size = rgenStride;
+		region->deviceAddress = sbtAddress + rgenOffset + rgenStride * i;
+		i++;
+	}
 	missRegion.deviceAddress = sbtAddress + missOffset;
 	hitRegion.deviceAddress = sbtAddress + hitOffset;
 }

src/lightmap/gpuraytracer.h

@@ -92,6 +92,7 @@ private:
 
 	const int coverageSampleCount = 256;
 	const int bounceSampleCount = 2048;
+	const int ambientSampleCount = 2048;
 	int rayTraceImageSize = 1024;
 
 	LevelMesh* mesh = nullptr;
@@ -120,9 +121,9 @@ private:
 	std::unique_ptr<VulkanBuffer> tlAccelStructBuffer;
 	std::unique_ptr<VulkanAccelerationStructure> tlAccelStruct;
 
-	std::unique_ptr<VulkanShader> rgenBounce, rgenLight;
-	std::unique_ptr<VulkanShader> rmissBounce, rmissLight, rmissSun;
-	std::unique_ptr<VulkanShader> rchitBounce, rchitLight, rchitSun;
+	std::unique_ptr<VulkanShader> rgenBounce, rgenLight, rgenAmbient;
+	std::unique_ptr<VulkanShader> rmissBounce, rmissLight, rmissSun, rmissAmbient;
+	std::unique_ptr<VulkanShader> rchitBounce, rchitLight, rchitSun, rchitAmbient;
 
 	std::unique_ptr<VulkanDescriptorSetLayout> descriptorSetLayout;
 
@@ -131,7 +132,7 @@ private:
 	std::unique_ptr<VulkanBuffer> shaderBindingTable;
 	std::unique_ptr<VulkanBuffer> sbtTransferBuffer;
 
-	VkStridedDeviceAddressRegionKHR rgenBounceRegion = {}, rgenLightRegion = {};
+	VkStridedDeviceAddressRegionKHR rgenBounceRegion = {}, rgenLightRegion = {}, rgenAmbientRegion = {};
 	VkStridedDeviceAddressRegionKHR missRegion = {};
 	VkStridedDeviceAddressRegionKHR hitRegion = {};
 	VkStridedDeviceAddressRegionKHR callRegion = {};