GPU ray tracing in Unity

Original author: David Kuri
  • Transfer
Amazing times have come for ray tracing. NVIDIA implements AI-accelerated noise reduction , Microsoft announces native support in DirectX 12 , and Peter Shirley sells his books at a free price ( pay what you want ). It seems that ray tracing finally got a chance to be accepted at court. It may be too early to talk about the beginning of the revolution, but it is definitely worth starting to study and accumulate knowledge in this area.

In this article, we will write from scratch in Unity a very simple ray tracer using compute shaders. We will write scripts in C #, and shaders in HLSL. All code is uploaded to Bitbucket .

As a result, we can render something like this:


Ray Tracing Theory


I want to start with a brief overview of the basics of ray tracing theory. If you are familiar with it, you can safely skip this section.

Let's imagine how photographs appear in the real world - very simplified, but this will be enough to explain the rendering. It all starts with a light source emitting photons. The photon flies in a straight line until it collides with the surface, after which it is reflected or refracted, and then continues its journey, losing some of the energy absorbed by the surface. Sooner or later, some of the photons fall into the camera’s sensor, which in turn creates a finished image. In essence, the ray tracing procedure mimics these steps to create photorealistic images.

In practice, only a small fraction of the photons emitted by the light source will reach the camera. Therefore, using the Heimgoltz principle of reversibility, the calculations are performed in the reverse order: instead of emitting photons from light sources, rays are emitted into the scene from the camera, reflected or refracted, and finally reach the light source.

The ray tracer we are going to create is based on a 1980 article by Turner Whited . We can simulate sharp shadows and perfectly correct reflections. In addition, the tracer will serve as the basis for the implementation of more complex effects, such as refraction, diffuse global illumination, brilliant reflections and soft shadows.

The basics


Let's start by creating a new Unity project. Create a C # script RayTracingMaster.csand compute shader RayTracingShader.compute. Paste the following base code into your C # script:

using UnityEngine;
public class RayTracingMaster : MonoBehaviour
{
    public ComputeShader RayTracingShader;
    private RenderTexture _target;
    private void OnRenderImage(RenderTexture source, RenderTexture destination)
    {
        Render(destination);
    }
    private void Render(RenderTexture destination)
    {
        // Make sure we have a current render target
        InitRenderTexture();
        // Set the target and dispatch the compute shader
        RayTracingShader.SetTexture(0, "Result", _target);
        int threadGroupsX = Mathf.CeilToInt(Screen.width / 8.0f);
        int threadGroupsY = Mathf.CeilToInt(Screen.height / 8.0f);
        RayTracingShader.Dispatch(0, threadGroupsX, threadGroupsY, 1);
        // Blit the result texture to the screen
        Graphics.Blit(_target, destination);
    }
    private void InitRenderTexture()
    {
        if (_target == null || _target.width != Screen.width || _target.height != Screen.height)
        {
            // Release render texture if we already have one
            if (_target != null)
                _target.Release();
            // Get a render target for Ray Tracing
            _target = new RenderTexture(Screen.width, Screen.height, 0,
                RenderTextureFormat.ARGBFloat, RenderTextureReadWrite.Linear);
            _target.enableRandomWrite = true;
            _target.Create();
        }
    }
}

The function is OnRenderImageautomatically called by Unity after the camera finishes rendering. To render, we first need to create a render target with the appropriate dimensions and inform the compute shader about this. 0 is the index of the compute shader kernel function - we have only one.

Then we passshader. This means that we ask the GPU to deal with groups of threads that execute our shader code. Each group of threads consists of several threads, the number of which is set in the shader itself. The size and number of thread groups can be indicated in three dimensions, so you can simply apply compute shaders to tasks of any dimension. In our case, we need to create one stream per pixel of the target render. The default thread group size specified in the compute shader Unity template is equal [numthreads(8,8,1)], so we will stick to it and create one stream group for every 8 × 8 pixels. In the end, we write the result to the screen with Graphics.Blit.

Let's check the program. Add a component to the scene camera RayTracingMaster(this is important when callingOnRenderImage), assign a compute shader and run play mode. You should see the output of the compute shader Unity template in the form of a beautiful triangular fractal.

Camera


Now that we can display the images on the screen, let's generate the rays of the camera. Since Unity provides us with a fully functional camera, we can simply use the calculated matrices for this. Let's start by setting the matrices in the shader. Add the following lines to the script RayTracingMaster.cs:

private Camera _camera;
private void Awake()
{
    _camera = GetComponent();
}
private void SetShaderParameters()
{
    RayTracingShader.SetMatrix("_CameraToWorld", _camera.cameraToWorldMatrix);
    RayTracingShader.SetMatrix("_CameraInverseProjection", _camera.projectionMatrix.inverse);
}

Before rendering, call SetShaderParametersfrom OnRenderImage.

In the shader, we define the matrices, structure, Rayand function to construct. Keep in mind that in HLSL, unlike C #, a function or variable must be declared before they are used. For the center of each screen pixel, we calculate the source and direction of the beam, and display the latter as color. Here's what the whole shader looks like:

#pragma kernel CSMain
RWTexture2D Result;
float4x4 _CameraToWorld;
float4x4 _CameraInverseProjection;
struct Ray
{
    float3 origin;
    float3 direction;
};
Ray CreateRay(float3 origin, float3 direction)
{
    Ray ray;
    ray.origin = origin;
    ray.direction = direction;
    return ray;
}
Ray CreateCameraRay(float2 uv)
{
    // Transform the camera origin to world space
    float3 origin = mul(_CameraToWorld, float4(0.0f, 0.0f, 0.0f, 1.0f)).xyz;
    // Invert the perspective projection of the view-space position
    float3 direction = mul(_CameraInverseProjection, float4(uv, 0.0f, 1.0f)).xyz;
    // Transform the direction from camera to world space and normalize
    direction = mul(_CameraToWorld, float4(direction, 0.0f)).xyz;
    direction = normalize(direction);
    return CreateRay(origin, direction);
}
[numthreads(8,8,1)]
void CSMain (uint3 id : SV_DispatchThreadID)
{
    // Get the dimensions of the RenderTexture
    uint width, height;
    Result.GetDimensions(width, height);
    // Transform pixel to [-1,1] range
    float2 uv = float2((id.xy + float2(0.5f, 0.5f)) / float2(width, height) * 2.0f - 1.0f);
    // Get a ray for the UVs
    Ray ray = CreateCameraRay(uv);
    // Write some colors
    Result[id.xy] = float4(ray.direction * 0.5f + 0.5f, 1.0f);
}

Try to rotate the camera in the inspector. You will see that the “colored sky" behaves accordingly.

Now let's replace the colors with a real skybox. In my examples, I will use Cape Hill from the HDRI Haven website, but you can, of course, choose any other. Download and drag it to Unity. In the import settings, do not forget to increase the maximum resolution if the resolution of the downloaded file is greater than 2048. Now add to the script public Texture SkyboxTexture, assign a texture in the inspector and set it in the shader, adding SetShaderParametersthis line to the function :

RayTracingShader.SetTexture(0, "_SkyboxTexture", SkyboxTexture);

In the shader, define the texture and the corresponding sampler, as well as the constant that we will use soon:

Texture2D _SkyboxTexture;
SamplerState sampler_SkyboxTexture;
static const float PI = 3.14159265f;

Now, instead of recording the color of the direction, we sample the skybox. To do this, we transform the Cartesian direction vector into spherical coordinates and associate it with the texture coordinates. Replace the last part with the CSMainfollowing:

// Sample the skybox and write it
float theta = acos(ray.direction.y) / -PI;
float phi = atan2(ray.direction.x, -ray.direction.z) / -PI * 0.5f;
Result[id.xy] = _SkyboxTexture.SampleLevel(sampler_SkyboxTexture, float2(phi, theta), 0);

Trace


So far so good. Now we will begin the ray tracing itself. Mathematically, we can calculate the intersection between the beam and the geometry of the scene, and save the collision parameters (position, normal, and distance along the beam). If our beam collides with several objects, then we will choose the nearest one. Let's define a struct in the shader RayHit:

struct RayHit
{
    float3 position;
    float distance;
    float3 normal;
};
RayHit CreateRayHit()
{
    RayHit hit;
    hit.position = float3(0.0f, 0.0f, 0.0f);
    hit.distance = 1.#INF;
    hit.normal = float3(0.0f, 0.0f, 0.0f);
    return hit;
}

Usually scenes are made up of many triangles, but we will start with a simple one: from the intersection of an infinite plane of the earth and several spheres!

Ground plane


Calculation of the intersection of a line with an infinite plane for $ y = 0 $- a fairly simple task. However, we only consider collisions in the positive direction of the beam and discard all collisions that are no closer than a potential previous collision.

By default, parameters in HLSL are passed by value, not by reference, so we will only work with a copy and will not be able to pass changes to the calling function. We pass RayHit bestHitwith the qualifier inoutto be able to modify the original struct. Here's what the shader code looks like:

void IntersectGroundPlane(Ray ray, inout RayHit bestHit)
{
    // Calculate distance along the ray where the ground plane is intersected
    float t = -ray.origin.y / ray.direction.y;
    if (t > 0 && t < bestHit.distance)
    {
        bestHit.distance = t;
        bestHit.position = ray.origin + t * ray.direction;
        bestHit.normal = float3(0.0f, 1.0f, 0.0f);
    }
}

To use it, let's add a wireframe function Trace(how much we will extend it):

RayHit Trace(Ray ray)
{
    RayHit bestHit = CreateRayHit();
    IntersectGroundPlane(ray, bestHit);
    return bestHit;
}

In addition, we need a basic shader function. Here we pass again Raywith inout- we will change it later when we talk about reflections. For debugging purposes, we will return normal in a collision with geometry, and otherwise return to the skybox sampling code:

float3 Shade(inout Ray ray, RayHit hit)
{
    if (hit.distance < 1.#INF)
    {
        // Return the normal
        return hit.normal * 0.5f + 0.5f;
    }
    else
    {
        // Sample the skybox and write it
        float theta = acos(ray.direction.y) / -PI;
        float phi = atan2(ray.direction.x, -ray.direction.z) / -PI * 0.5f;
        return _SkyboxTexture.SampleLevel(sampler_SkyboxTexture, float2(phi, theta), 0).xyz;
    }
}

We will use both functions in CSMain. Delete the skybox sampling code if you haven’t done so already, and add the following lines for ray tracing and shadowing the collision:

// Trace and shade
RayHit hit = Trace(ray);
float3 result = Shade(ray, hit);
Result[id.xy] = float4(result, 1);

Sphere


A plane is not the most interesting object in the world, so let's add a sphere right away. Mathematical calculations of the intersection of a line and a sphere can be found on Wikipedia . This time we will have only two options for beam collisions: the entry point p1 - p2and the exit point p1 + p2. First we will check the entry point, and use the exit point if the other does not fit. In our case, the sphere is defined as the value float4consisting of the position (xyz) and radius (w). Here's what the code looks like:

void IntersectSphere(Ray ray, inout RayHit bestHit, float4 sphere)
{
    // Calculate distance along the ray where the sphere is intersected
    float3 d = ray.origin - sphere.xyz;
    float p1 = -dot(ray.direction, d);
    float p2sqr = p1 * p1 - dot(d, d) + sphere.w * sphere.w;
    if (p2sqr < 0)
        return;
    float p2 = sqrt(p2sqr);
    float t = p1 - p2 > 0 ? p1 - p2 : p1 + p2;
    if (t > 0 && t < bestHit.distance)
    {
        bestHit.distance = t;
        bestHit.position = ray.origin + t * ray.direction;
        bestHit.normal = normalize(bestHit.position - sphere.xyz);
    }
}

To add a sphere, we simply call this function from Trace, for example, like this:

// Add a floating unit sphere
IntersectSphere(ray, bestHit, float4(0, 3.0f, 0, 1.0f));

Smoothing


The approach used has one problem: we check only the center of each pixel, so distortions (ugly "ladders") will be noticeable as a result. To get around this problem, we will not trace one, but several rays per pixel. Each ray receives a random offset within the region of the pixel. In order to maintain an acceptable level of frame rate, we will perform progressive sampling, that is, trace one beam per pixel per frame and over time average the value if the camera does not move. Each time you move the camera (or change any other parameters - visibility, stage geometry or lighting), we will have to start all over again.

Let's create a very simple image effect shader, which we will use to add a few results. Call this shaderAddShaderand check that the first line is Shader "Hidden/AddShader". After Cull Off ZWrite Off ZTest Alwaysadd Blend SrcAlpha OneMinusSrcAlphato enable alpha blending. Then replace the function with the fragfollowing lines:

float _Sample;
float4 frag (v2f i) : SV_Target
{
    return float4(tex2D(_MainTex, i.uv).rgb, 1.0f / (_Sample + 1.0f));
}

Now this shader will simply render the first sample with opacity $ 1 $next with opacity $ \ frac {1} {2} $then $ \ frac {1} {3} $and so on, averaging all samples with equal weight.

In the script, we need to read the samples and apply the image effect shader:

private uint _currentSample = 0;
private Material _addMaterial;

Also, when rebuilding the target render, InitRenderTexturewe need to perform a reset _currentSamples = 0and add a function Updatethat recognizes the change in camera transformations:

private void Update()
{
    if (transform.hasChanged)
    {
        _currentSample = 0;
        transform.hasChanged = false;
    }
}

To use our shader, we need to initialize the material, tell it about the current sample and use it to insert on the screen in the function Render:

// Blit the result texture to the screen
if (_addMaterial == null)
    _addMaterial = new Material(Shader.Find("Hidden/AddShader"));
_addMaterial.SetFloat("_Sample", _currentSample);
Graphics.Blit(_target, destination, _addMaterial);
_currentSample++;

So, we are already doing progressive sampling, but still use the center of the pixel. In compute shader, we set it float2 _PixelOffsetand use it CSMaininstead of a hard-coded offset float2(0.5f, 0.5f). Go back to the script and create a random offset by adding to the SetShaderParametersfollowing line:

RayTracingShader.SetVector("_PixelOffset", new Vector2(Random.value, Random.value));

If you move the camera, you can see that the image is still visible distortion, but they quickly disappear if you stand still for a couple of frames. Here is a comparison of what we did:



Reflection


The foundation for our ray tracer is ready, so we can get down to tricky things that actually distinguish ray tracing from other rendering techniques. The first on this list are perfect reflections. The idea is simple: when we collide with the surface, we reflect the ray in accordance with the law of reflection, which you may remember from school (angle of incidence = angle of reflection), reduce its energy and repeat the process until the ray either collides with the sky, or he will not run out of energy after a given number of reflections.

In the shader, add a variable to the ray float3 energyand initialize it in the function CreateRayas ray.energy = float3(1.0f, 1.0f, 1.0f). Initially, the beam will have maximum values ​​in all color channels, which will decrease with each reflection.

We will perform a maximum of 8 traces (the initial ray plus 7 reflections), and add the results in function calls Shade, but multiplied by the ray energy. For example, imagine that a ray was reflected once and lost$ \ frac {3} {4} $your energy. Then it continues to move and collides with the sky, so we only transfer to the pixel$ \ frac {1} {4} $energy of the sky. Change CSMainas follows, replacing previous calls Traceand Shade:

// Trace and shade
float3 result = float3(0, 0, 0);
for (int i = 0; i < 8; i++)
{
    RayHit hit = Trace(ray);
    result += ray.energy * Shade(ray, hit);
    if (!any(ray.energy))
        break;
}

Our function Shadenow also performs energy renewal and reflection generation, which is why it becomes important here inout. To renew energy, we perform element-wise multiplication by the reflected color of the surface. For example, for gold, the specular reflection coefficient is approximately equal float3(1.0f, 0.78f, 0.34f), that is, it reflects 100% red, 78% green and only 34% blue, giving the reflection a characteristic golden hue. Be careful, none of these values ​​should exceed 1, because otherwise the energy from us will be created from nowhere. In addition, reflectivity is often lower than you might think. For example, see some of the values ​​on slide 64 in the article by Physics and Math of Shading by Naty Hoffman.

HLSL has a built-in function for reflecting a beam with a given normal, and this is convenient. Due to the inaccuracy of floating point numbers, it may happen that the reflected ray is blocked by the surface from which it is reflected. To avoid this, we will slightly shift the position along the normal direction. Here's what the new feature looks like Shade:

float3 Shade(inout Ray ray, RayHit hit)
{
    if (hit.distance < 1.#INF)
    {
        float3 specular = float3(0.6f, 0.6f, 0.6f);
        // Reflect the ray and multiply energy with specular reflection
        ray.origin = hit.position + hit.normal * 0.001f;
        ray.direction = reflect(ray.direction, hit.normal);
        ray.energy *= specular;
        // Return nothing
        return float3(0.0f, 0.0f, 0.0f);
    }
    else
    {
        // Erase the ray's energy - the sky doesn't reflect anything
        ray.energy = 0.0f;
        // Sample the skybox and write it
        float theta = acos(ray.direction.y) / -PI;
        float phi = atan2(ray.direction.x, -ray.direction.z) / -PI * 0.5f;
        return _SkyboxTexture.SampleLevel(sampler_SkyboxTexture, float2(phi, theta), 0).xyz;
    }
}

You can try to slightly increase the brightness of the skybox by multiplying it by a factor greater than 1. Now experiment with the function Trace. Put several spheres in the loop and the result will be like this:


Directional light source


So, we can trace mirror reflections, which allows us to render smooth metal surfaces, but for non-metallic surfaces we need one more property: diffuse reflection. In short, metals only reflect incident light with a hint of their reflection color, while non-metals allow light to refract on the surface, scatter and leave it in a random direction, colored in the color of its albedo. In the case of an ideal Lambert surface , which is usually used, the probability is proportional to the cosine of the angle between the above direction and the surface normal. This topic is discussed in more detail here .

To get started with diffuse lighting, let's add toRayTracingMasterpublic Light DirectionalLightand set the directional light source in the scene. It may also be necessary to recognize changes in the transformations of the light source into functions Update, as we did with the transformations of the camera. Now add the SetShaderParametersfollowing lines to the function :

Vector3 l = DirectionalLight.transform.forward;
RayTracingShader.SetVector("_DirectionalLight", new Vector4(l.x, l.y, l.z, DirectionalLight.intensity));

Define in the shader float4 _DirectionalLight. In the function, Shadedefine the albedo color immediately after the specular color:

float3 albedo = float3(0.8f, 0.8f, 0.8f);

Replace the returned black values ​​with simple diffuse shading:

// Return a diffuse-shaded color
return saturate(dot(hit.normal, _DirectionalLight.xyz) * -1) * _DirectionalLight.w * albedo;

Do not forget that the scalar product is defined as $ a \ cdot b = || a || \ || b ||  \ cos \ theta $. Since both of our vectors (normal and direction of light) have a unit length, we need exactly the scalar product: the cosine of the angle. Ray and light have opposite directions, therefore, in direct lighting, the scalar product returns not 1, but -1. To take this into account, we must change sign. Finally, we saturate this value (for example, limit it to the interval$ [0,1] $) to avoid negative energy.

In order for the directional light source to cast shadows, we need to trace the shadow beam. It starts from the position of the surface under consideration (also with a very small displacement to avoid self-shadowing), and indicates in the direction from which the light came. If something blocks his path to infinity, then we will not use diffuse lighting. Add these lines above the diffuse color return:

// Shadow test ray
bool shadow = false;
Ray shadowRay = CreateRay(hit.position + hit.normal * 0.001f, -1 * _DirectionalLight.xyz);
RayHit shadowHit = Trace(shadowRay);
if (shadowHit.distance != 1.#INF)
{
    return float3(0.0f, 0.0f, 0.0f);
}

Now we can trace glossy plastic spheres with sharp shadows! If you set 0.04 for specular and 0.8 for albedo, then we get the following results:


Scene and materials


Let's now begin to create more complex and colorful scenes! Instead of hard tasking everything in the shader, we set the scene in C # for more versatility.

To begin with, we will expand the structure RayHitin the shader. Instead of setting global properties of materials in a function Shade, we will define them for each object and store them in RayHit. Add to struct float3 albedoand float3 specularinitialize them with values float3(0.0f, 0.0f, 0.0f)in CreateRayHit. Also change the function Shadeso that it uses hitthese values ​​instead of hard-coded values. <

To understand in general what the sphere is in the CPU and GPU, we will define struct Spherein the shader and in the script in C #. From the shader side, it looks like this:

struct Sphere
{
    float3 position;
    float radius;
    float3 albedo;
    float3 specular;
};

Copy this structure into a C # script.

In the shader, we need to make the function IntersectSpherework with our struct, not with float4. This is easy to do:

void IntersectSphere(Ray ray, inout RayHit bestHit, Sphere sphere)
{
    // Calculate distance along the ray where the sphere is intersected
    float3 d = ray.origin - sphere.position;
    float p1 = -dot(ray.direction, d);
    float p2sqr = p1 * p1 - dot(d, d) + sphere.radius * sphere.radius;
    if (p2sqr < 0)
        return;
    float p2 = sqrt(p2sqr);
    float t = p1 - p2 > 0 ? p1 - p2 : p1 + p2;
    if (t > 0 && t < bestHit.distance)
    {
        bestHit.distance = t;
        bestHit.position = ray.origin + t * ray.direction;
        bestHit.normal = normalize(bestHit.position - sphere.position);
        bestHit.albedo = sphere.albedo;
        bestHit.specular = sphere.specular;
    }
}

Also set bestHit.albedoand bestHit.specularin function IntersectGroundPlaneto adjust its material.

Then determine . At this point, the CPU will store all the areas that make up the scene. Remove all hard-coded spheres from the function and add the following lines:StructuredBuffer _SpheresTrace

// Trace spheres
uint numSpheres, stride;
_Spheres.GetDimensions(numSpheres, stride);
for (uint i = 0; i < numSpheres; i++)
    IntersectSphere(ray, bestHit, _Spheres[i]);

Now we will breathe a little life into the scene. Let's add general parameters to the C # script to control the location of the spheres and the compute buffer:

public Vector2 SphereRadius = new Vector2(3.0f, 8.0f);
public uint SpheresMax = 100;
public float SpherePlacementRadius = 100.0f;
private ComputeBuffer _sphereBuffer;

We will configure the scene in OnEnableand release the buffer in OnDisable. Thus, each time the component is turned on, a random scene will be generated. The function SetUpScenewill try to position the spheres in a certain radius and discard those that intersect existing ones. Half of the spheres are metallic (black albedo, color specular), the other half non-metallic (color albedo, 4% specular):

private void OnEnable()
{
    _currentSample = 0;
    SetUpScene();
}
private void OnDisable()
{
    if (_sphereBuffer != null)
        _sphereBuffer.Release();
}
private void SetUpScene()
{
    List spheres = new List();
    // Add a number of random spheres
    for (int i = 0; i < SpheresMax; i++)
    {
        Sphere sphere = new Sphere();
        // Radius and radius
        sphere.radius = SphereRadius.x + Random.value * (SphereRadius.y - SphereRadius.x);
        Vector2 randomPos = Random.insideUnitCircle * SpherePlacementRadius;
        sphere.position = new Vector3(randomPos.x, sphere.radius, randomPos.y);
        // Reject spheres that are intersecting others
        foreach (Sphere other in spheres)
        {
            float minDist = sphere.radius + other.radius;
            if (Vector3.SqrMagnitude(sphere.position - other.position) < minDist * minDist)
                goto SkipSphere;
        }
        // Albedo and specular color
        Color color = Random.ColorHSV();
        bool metal = Random.value < 0.5f;
        sphere.albedo = metal ? Vector3.zero : new Vector3(color.r, color.g, color.b);
        sphere.specular = metal ? new Vector3(color.r, color.g, color.b) : Vector3.one * 0.04f;
        // Add the sphere to the list
        spheres.Add(sphere);
    SkipSphere:
        continue;
    }
    // Assign to compute buffer
    _sphereBuffer = new ComputeBuffer(spheres.Count, 40);
    _sphereBuffer.SetData(spheres);
}

A magic number of 40 V new ComputeBuffer(spheres.Count, 40)is a step in our buffer, i.e. size of one sphere in memory in bytes. To calculate it, calculate the number of float in the struct Sphereand multiply it by the float byte size (4 bytes). Finally, we set the shader buffer in the function SetShaderParameters:

RayTracingShader.SetBuffer(0, "_Spheres", _sphereBuffer);

results


Congratulations, we did it! Now we have a ready-made Whited ray tracer on the GPU, capable of rendering many spheres with mirror reflections, simple diffused lighting and sharp shadows. Full source code is uploaded to Bitbucket . Experiment with the placement parameters of the spheres and observe the beautiful views:



What's next?


Today we have achieved a lot, but much more can be realized: diffused global lighting, soft shadows, partially transparent materials with refractions and, obviously, the use of triangles instead of spheres. In the next article, we will expand our Whited ray tracer into a path tracer to learn some of the above.

Also popular now: