Vulkan examples for ray traced shadows and reflections using VK_NV_ray_tracing

After adding a basic Nvidia RTX ray tracing example last week, I spent some more time with Vulkan and the VK_NV_ray_tracing extension. The result are two new, more advanced examples that I just uploaded to my Vulkan C++ example repository. As with the basic example I tried to keep them as straight forward as possible with all the relevant code parts put into one source file, so that following and building upon is as easy as possible.

Unlike most of the other advanced Vulkan examples, the most interesting parts can be found inside the ray tracing shaders. Most of the C++ code is just setup to get the shaders up and running and passing data to the ray tracing hardware.

Ray traced shadows

C ++ Source, GLSL Shaders

This example adds an additional miss shader named shadow.rmiss to the pipeline thats sets the shadowed flag payload to false if the ray does not hit any of the scene geometry:

#version 460
#extension GL_NV_ray_tracing : require

layout(location = 2) rayPayloadInNV bool shadowed;

void main()
{
	shadowed = false;
}

The actual shadow casting is triggered at the end of the closest hit shader closesthit.rchit:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

#version 460
#extension GL_NV_ray_tracing : require
#extension GL_EXT_nonuniform_qualifier : enable

layout(location = 0) rayPayloadInNV vec3 hitValue;
layout(location = 2) rayPayloadNV bool shadowed;

...

void main()
{
    ...

	// Shadow casting
	float tmin = 0.001;
	float tmax = 100.0;
	vec3 origin = gl_WorldRayOriginNV + gl_WorldRayDirectionNV * gl_HitTNV;
	shadowed = true;  
	// Offset indices to match shadow hit/miss index
	traceNV(
        topLevelAS, 
        gl_RayFlagsTerminateOnFirstHitNV | gl_RayFlagsOpaqueNV | gl_RayFlagsSkipClosestHitShaderNV, 
        0xFF, 
        1, 
        0, 
        1, 
        origin, tmin, lightVector, tmax, 2);
	if (shadowed) {
		hitValue *= 0.3;
	}
}

Note that we need to pass the proper offsets (line 24 and 26) to traceNV to make it use the correct shaders that have were passed to vkCreateRayTracingPipelinesNV. Since our shadow miss shader is the second shader of that given type, we need to offset by 1.

With this setup, if the ray misses our geometry, the shadowed ray payload will be set to false by the shadow miss shader and can be used to darken the hit color to darken shadowed parts of the scene.

Ray traced reflections

C ++ Source, GLSL Shaders

While doing real time reflections with rasterization is hard and involves lots of tricks, esp. if you want recursion and multi-directional reflections, they’re a perfect match for ray tracing. This example renders an old-school looking scene with multiple reflective objects and multiple levels of recursion for creating perfect reflections.

Reflections are traced within a loop inside the ray generation shaders rgen.rgen:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

#version 460
#extension GL_NV_ray_tracing : require

...

struct RayPayload {
	vec3 color;
	float distance;
	vec3 normal;
	float reflector;
};

layout(location = 0) rayPayloadNV RayPayload rayPayload;

// Max. number of recursion is passed via a specialization constant
layout (constant_id = 0) const int MAX_RECURSION = 0;

void main() 
{
    ...

	vec3 color = vec3(0.0);

	for (int i = 0; i < MAX_RECURSION; i++) {
		traceNV(topLevelAS, rayFlags, cullMask, 0, 0, 0, origin.xyz, tmin, direction.xyz, tmax, 0);
		vec3 hitColor = rayPayload.color;

		if (rayPayload.distance < 0.0f) {
			color += hitColor;
			break;
		} else if (rayPayload.reflector == 1.0f) {
			const vec4 hitPos = origin + direction * rayPayload.distance;
			origin.xyz = hitPos.xyz + rayPayload.normal * 0.001f;
			direction.xyz = reflect(direction.xyz, rayPayload.normal);
		} else {
			color += hitColor;
			break;
		}
	}

	imageStore(image, ivec2(gl_LaunchIDNV.xy), vec4(color, 0.0));
}

If the ray hits a reflective surface (line 31), we reflect it’s direction using the normal of the current hit point, so the next loop iteration traces in the reflected direction.

For this shader we need some additional info on the ray hit, so instead of a single hitValue like in the previous sample, the ray payload (line 6-10) has been expanded to include some additional data that is filled in the closes hit shaders closeshit.rchit :

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

#version 460
#extension GL_NV_ray_tracing : require
#extension GL_EXT_nonuniform_qualifier : enable

struct RayPayload {
	vec3 color;
	float distance;
	vec3 normal;
	float reflector;
};

layout(location = 0) rayPayloadInNV RayPayload rayPayload;

...

void main()
{
    ...

	rayPayload.color = v0.color * vec3(dot_product);
	rayPayload.distance = gl_RayTmaxNV;
	rayPayload.normal = normal;

	// Objects with full white vertex color are treated as reflectors
	rayPayload.reflector = ((v0.color.r == 1.0f) && (v0.color.g == 1.0f) && (v0.color.b == 1.0f)) ? 1.0f : 0.0f; 
}

Note the use of the gl_RayTmaxNV GLSL built-in. In a closest hit shader, this value stores the closest distance to the intersected primitive. This value is used in the above ray generation shader. Details on the VK_NV_ray_tracing built-ins can be found in the official extension spec.

To keep things simple, we treat vertices with full white color as reflectors. In a more complex example you’d pass that info e.g. via a material buffer storing material IDs.

And this is how it looks in motion: