WebGPU Engine from Scratch Part 8: Physically Based Lighting (PBR)

I wasn’t expecting to do this so soon but the whole roughness texture stuff really got to me so I decided to see what it would take to improve the lighting to PBR (Physically based rendering). Luckily, not very hard. While we do need to look at equations it’s mostly a matter of plugging them in, no aligned buffer packing or anything crazy.

Most of the materials I looked through started with the rendering equation. To be honest, I don’t really care for the rendering equation because it makes everything look way more complicated than it needs to be. I understand it’s fundamental value though but I’m not going to bother printing it because we aren’t going to look at. Suffice to say, all that matters is what it means: the color of a pixel is the result of all incoming light and the amount of light reflected back to the viewer which involves the dot product of the light direction and the normal.

This is mostly what we’ve been doing so far, but PBR is a bit more strict in that it respects energy conservation. That is, all light has to be accounted for either absorbed, reflected or transmitted. It is impossible for there to be more light given off than there was coming in. For the classic Blinn-Phong this isn’t true because we had factors we could change for the specular, tuning them artistically. PBR doesn’t have that, it gives a set of physical properties (roughness and reflectivity) we use calculate both the diffuse and specular light.

The other part is that it’s based on micro-facets. This means that instead of a flat plane the surface is modeled as having little bumps that reflect light in different directions. What’s nice is that this is not really different than calculating the light reflected by a flat plane but since each microfacet is smaller than a pixel we just use statistical scaling. If 30% of the micro facets of a pixel face the viewer then you essentially multiply the value by 0.3.

BDRF

The BDRF is the Bidirectional Radiance Function and it describes how light gets reflected off the surface. There are multiple BDRFs of which Blinn-Phong is one. In fact, in practice it’s split up the same way into diffuse and specular parts.

Here s and d are scaling values and R_s and R_d are functions that describe the diffuse and specular parts respectively. The key to energy conservation here is the constraint that s + d = 1. They cannot add up to more that the total light. Note: the Cook-Torrence model also contains a Fresnel term. You only need to apply Fresnel once, so the s term will be removed in this case.

Diffuse

Normal Lambertian diffuse seems to be the most common and it’s exactly what we’ve been doing.

$diffuse = Albedo \star (L \cdot N) * LightColor$

Albedo is the surface color
L is the vector towards the light
N is the surface normal
LightColor is the color of the light

That is the diffuse light is just the incoming light times the dot product of the surface normal and the vector pointing toward the light. (note that the dot product will get factored out later).

Specular

This is where it gets a little more interesting. We’ll be using the Cook-Torrence model because that seems to be one of the most common but there are others.

$cookTorrence = \frac{DGF}{4 (V \cdot N)(L \cdot N)}$

V is the direction toward the viewer
N is the normal of the surface
L is the direction to the light
D, G and F are functions

D (the normal distribution function)

This function describes how the normals of microfacets are oriented on the surface. What we want to know is what ratio of them point toward the viewer. For this there are several but a common one I saw was GGX (Trowbridge-Reitz).

$GGX = \frac{\alpha^2}{\pi ((N \cdot H)^2(\alpha^2 - 1) + 1)^2}$

alpha is roughness (yay! we found it)
H is the half vector (same as Blinn Phong)
N is still the surface normal

I honestly have no idea how this one works under the hood but as long as it works I’m not going to question it.

G (the geometry function)

This function describes how the light bounces between the surface bump, specifically to create cases where there is self-shadowing. So if a ray comes in and bounces between two peaks it might not reflect back the viewer as might be expected if we didn’t consider the bouncing cases. The most common one here seemed to be Shlick-GGX.

$GGX_{schlick} = \frac{N \cdot X}{(N \cdot X)(1 - k) + k}$

k is alpha/2 in some renderers, where alpha is roughness. Might also see k = (roughness + 1)^2 / 8 .
N is the surface normal
X is either the vector to the light or to the viewer.

We actually have to apply this twice once where X = toLight and the other where X = toCamera.

$G = GGX_{schlick}(L) \star GGX_{schlick}(V)$

L is the vector toward the light
V is the vector toward the viewer

F (Fresnel)

This is the amount of light that is reflected versus refracted. We use the Schlick approximation for it.

$F_{schlick} = F_0 + (1 - F_0)(1 - (V \cdot H))^5$

F_0 is the base reflectance (ie how reflective it is when looking straight down the normal vector). It is an RGB vector.
V is the vector toward the viewer
H is the half vector between the view direction and the light direction

The F_0 value is something you can look up for common materials. You can also find the Index of Refraction (IOR) for the material and use the equation:

$F_0 = \frac{(r-1)^2}{(r+1)^2}$

r is the index of refraction

I also saw that it’s common to just use a hardcoded value of 0.04 for all dielectrics (non-metals). My guess is because the range is so small and 2 orders of magnitude less than those of metals.

The last part of the BDRF is figuring out s and d from the original BDRF equation. By definition, d = 1 - s. And s will be the Fresnel value shown above. With that, we have all equations and all variables filled so we can implement it in the shader.

Implementation

struct VertexOut {
    @builtin(position) frag_position : vec4,
    @location(0) world_position: vec4,
    @location(1) uv : vec2,
    @location(2) normal : vec3
};
struct Scene {
    view_matrix: mat4x4,
    projection_matrix: mat4x4,
    model_matrix: mat4x4,
    normal_matrix: mat3x3,
    camera_position: vec3
}
struct Light {
    light_type: u32,
    position: vec3,
    direction: vec3,
    color: vec4
}
struct LightCount {
    count: u32
}
struct SpecularMaterial {
    use_specular_map: u32,
    roughness: f32
}

@group(0) @binding(0) var scene : Scene;

@group(1) @binding(0) var albedo_sampler: sampler;
@group(1) @binding(1) var albedo_map: texture_2d;
@group(1) @binding(2) var specular: SpecularMaterial;
@group(1) @binding(3) var roughness_sampler: sampler;
@group(1) @binding(4) var roughness_map: texture_2d;

@group(2) @binding(0) var lights: array;
@group(2) @binding(1) var light_count: LightCount;

//schlick
fn get_fresnel(f0: vec3, to_view: vec3, half_vector: vec3) -> vec3 {
   return f0 + (vec3(1.0) - f0) * pow(1 - max(dot(to_view, half_vector), 0), 5.0);
}

//lambert
fn get_diffuse(surface_albedo: vec3, light_color: vec3, normal: vec3, to_light: vec3) -> vec3 {
    return surface_albedo * max(dot(normal, to_light), 0.0) * light_color;
}

//GGX
fn get_normal_distribution(roughness: f32, normal: vec3, half_vector: vec3) -> f32 {
    var roughness_squared = pow(roughness, 2.0);
    var n_dot_h = max(dot(normal, half_vector), 0.0);
    var denominator = pow(n_dot_h * n_dot_h *  (roughness_squared - 1.0) + 1.0, 2.0);
    return roughness_squared / ((3.1415 * denominator) + 0.00001);
}

//GGX-schlick
fn get_geometry(to_view: vec3, to_light: vec3, normal: vec3, roughness: f32) -> f32 {
    var k = roughness / 2; //might be different in other renderers
    //var k = pow(roughness + 1, 2.0) / 8; 
    var n_dot_l = max(dot(normal, to_light), 0.0);
    var geometry_light = n_dot_l / ((n_dot_l * (1.0 - k) + k) + 0.00001);
    var n_dot_v = max(dot(normal, to_view), 0.0);
    var geometry_view = n_dot_v / ((n_dot_v * (1.0 - k) + k) + 0.00001);

    return geometry_light * geometry_view;
}

//cook-torrence
fn get_specular(fresnel: vec3, to_view: vec3, to_light: vec3, normal: vec3, half_vector: vec3, roughness: f32) -> vec3 {
    var d = get_normal_distribution(roughness, normal, half_vector);
    var g = get_geometry(to_view, to_light, normal, roughness);
    var f = fresnel;
    var v_dot_n = max(dot(to_view, normal), 0.0);
    var l_dot_n = max(dot(to_light, normal), 0.0);
    return (d * g * f) / (4 * v_dot_n * l_dot_n);
}

fn get_bdrf(surface_albedo: vec3, f0: vec3, roughness: f32, normal: vec3, light_color: vec3, light_pos: vec3, camera_pos: vec3, frag_pos: vec3) -> vec3{
    var to_view = normalize(camera_pos - frag_pos);
    var to_light = normalize(light_pos - frag_pos);
    var half_vector = normalize(to_view + to_light);
    var fresnel = get_fresnel(f0, to_view, half_vector);
    var ks = fresnel;
    var kd = vec3(1.0) - ks;

    var diffuse = get_diffuse(surface_albedo, light_color, normal, to_light);
    var specular = get_specular(fresnel, to_view, to_light, normal, half_vector, roughness);

    return kd * diffuse + specular;
}

@vertex
fn vertex_main(@location(0) position: vec3, @location(1) uv: vec2, @location(2) normal: vec3) -> VertexOut
{
    var output : VertexOut;
    output.frag_position =  scene.projection_matrix * scene.view_matrix * scene.model_matrix * vec4(position, 1.0);
    output.world_position = scene.model_matrix * vec4(position, 1.0);
    output.uv = uv;
    output.normal = scene.normal_matrix * normal;
    return output;
}
@fragment
fn fragment_main(frag_data: VertexOut) -> @location(0) vec4
{   
    var surface_albedo = textureSample(albedo_map, albedo_sampler, frag_data.uv).rgb;
    var roughness_from_map = textureSample(roughness_map, roughness_sampler, frag_data.uv).x;
    var roughness = mix(specular.roughness, roughness_from_map, f32(specular.use_specular_map));
    var f0 = vec3(0.04, 0.04, 0.04);
    var total_color = vec3(0.0);
    var normal = normalize(frag_data.normal);

    for(var i: u32 = 0; i < light_count.count; i++){
        var light = lights[i];
        total_color += get_bdrf(surface_albedo, f0, roughness, normal, light.color.rgb, light.position, scene.camera_position, frag_data.world_position.xyz);
    }

    return vec4(total_color, 1.0);
}

A few things here. I changed the names of some things, “specular” is now “roughness” and the base texture is now call “albedo” which is the more common name to use. Anywhere you saw a dot product, it’s max(dot(X, Y), 0.0). This is prevent negative dot products which might cause weird bugs. I also add 0.00001 to each denominator to prevent accidental divide by zero.

As for the overall result. It seems decent to me.

Light attenuation

One thing we didn’t really apply is light attenuation. That is, the light get less bright the further away you are. We can add that term.

@fragment
fn fragment_main(frag_data: VertexOut) -> @location(0) vec4
{   
    var surface_albedo = textureSample(albedo_map, albedo_sampler, frag_data.uv).rgb;
    var roughness_from_map = textureSample(roughness_map, roughness_sampler, frag_data.uv).x;
    var roughness = mix(specular.roughness, roughness_from_map, f32(specular.use_specular_map));
    var f0 = vec3(0.04, 0.04, 0.04);
    var total_color = vec3(0.0);
    var normal = normalize(frag_data.normal);

    for(var i: u32 = 0; i < light_count.count; i++){
        var light = lights[i];
+       var light_distance = length(light.position - frag_data.world_position.xyz);
+       var attenuation = 1.0 / pow(light_distance, 2.0);
+       var radiance = light.color.rgb * attenuation;
+       total_color += get_bdrf(surface_albedo, f0, roughness, normal, radiance, light.position, scene.camera_position, frag_data.world_position.xyz);
-total_color += get_bdrf(surface_albedo, f0, roughness, normal, light.color, light.position, scene.camera_position, frag_data.world_position.xyz);
    }

    return vec4(total_color, 1.0);
}

Light falls off over distance squared so we just multiply by 1 over that. This is the inverse-power law.

The scene gets a lot darker.

Gamma correction

Some sources I looked at performed gamma correct, which is a process of stretching the rgb space so there is more room for dark colors and better maps the range to human vision. A simple approximation is to raise the color to the power of 1/2.2 .

After some reading it seems we might need to apply this manually. We are drawing to a rgba8unorm texture canvas which is in linear rgb space so there should not be any auto conversion to srgb (gamma corrected rgb) if I understand correctly. In fact, it seems you can’t configure a canvas context with rgba8unorm-srgb format anyway, at least in Chrome. We’ll try it.

In addition we should tone map as well. This just takes the unbounded color which could have components of greater than 1 which we cannot display and reduces them down into the range of 0 to 1. We do this with Reinhard’s operator which is var tone_mapped_color = color / (color + vec3(1.0) This should be done before gamma correction.

@fragment
fn fragment_main(frag_data: VertexOut) -> @location(0) vec4
{   
    var surface_albedo = textureSample(albedo_map, albedo_sampler, frag_data.uv).rgb;
    var roughness_from_map = textureSample(roughness_map, roughness_sampler, frag_data.uv).x;
    var roughness = mix(specular.roughness, roughness_from_map, f32(specular.use_specular_map));
    var f0 = vec3(0.04, 0.04, 0.04);
    var total_color = vec3(0.0);
    var normal = normalize(frag_data.normal);

    for(var i: u32 = 0; i < light_count.count; i++){
        var light = lights[i];
        var light_distance = length(light.position - frag_data.world_position.xyz);
        var attenuation = 1.0 / pow(light_distance, 2.0);
        var radiance = light.color.rgb * attenuation;
        total_color += get_bdrf(surface_albedo, f0, roughness, normal, radiance, light.position, scene.camera_position, frag_data.world_position.xyz);
    }

-   return vec4(total_color, 1.0);
+   var tone_mapped_color = total_color / (total_color + vec3(1.0));
+   return vec4(pow(total_color, vec3(1.0/2.2)), 1.0);
}

We can see more of the dark details although artistically it kinda sucks.

We can play around with the light’s brightness by using a color vector with value that are greater than 1.

Turned up the brightness with color [2.5, 2.5, 2.5, 1]. Not too bad.

Metals

The last things we want to consider are metals. These behave slightly differently in the PBR model because metals have much higher F0 values and they vary by color channel. Let’s try making a gold teapot to test it out. We’ll need to add 2 new properties to the material:

//material.js
export class Material {
  //...stuff
  constructor(options){
    //...stuff
    this.#metalness = options.metalness ?? 0.0;
    this.#baseReflectance = options.baseReflectance ?? [0.04, 0.04, 0.04];
  }
  #metalness;
  #baseReflectance;
  set metalness(val){
    this.#metalness = val;
  }
  get metalness(){
    return this.#metalness;
  }
  set baseReflectance(val){
    this.#baseReflectance = new Float32Array(val);
  }
}

metalness is how much the material is metal (this will most likely be 0 or 1 but we’ll provide a range because it’s easier). baseReflectance is the F0 value in RGB.

physicallybased.info says a good F0 for gold is: 1.059,0.773,0.307.

In the shader we’ll add those props to the Material struct (formerly SpecularMaterial, “specular” was redundant, instances called “specular” were renamed to “material”).

struct Material {
    use_specular_map: u32,
    roughness: f32,
+   metalness: f32,
+   base_reflectance: vec3
}

We use the metalness value in the BDRF:

-fn get_bdrf(surface_albedo: vec3, f0: vec3, roughness: f32, normal: vec3, light_color: vec3, light_pos: vec3, camera_pos: vec3, frag_pos: vec3) -> vec3{
+fn get_bdrf(surface_albedo: vec3, f0: vec3, roughness: f32, metalness: f32, normal: vec3, light_color: vec3, light_pos: vec3, camera_pos: vec3, frag_pos: vec3) -> vec3{
    var to_view = normalize(camera_pos - frag_pos);
    var to_light = normalize(light_pos - frag_pos);
    var half_vector = normalize(to_view + to_light);
    var fresnel = get_fresnel(f0, to_view, half_vector);
    var ks = fresnel;
-   var kd = (vec3(1.0) - ks);
+   var kd = (vec3(1.0) - ks) * (1.0 - metalness);

    var diffuse = get_diffuse(surface_albedo, light_color, normal, to_light);
    var specular = get_specular(fresnel, to_view, to_light, normal, half_vector, roughness);
    return kd * diffuse + specular;
}

Things that are metal don’t diffuse so we essentially turn it off if the metalness is 1. We also vary the f0 value by interpolating between the dielectric value and the passed in value:

@fragment
fn fragment_main(frag_data: VertexOut) -> @location(0) vec4
{   
    var surface_albedo = textureSample(albedo_map, albedo_sampler, frag_data.uv).rgb;
    var roughness_from_map = textureSample(roughness_map, roughness_sampler, frag_data.uv).x;
-   var roughness = max(mix(material.roughness, roughness_from_map, f32(material.use_specular_map)), 0.0);
+   var roughness = max(mix(material.roughness, roughness_from_map, f32(material.use_specular_map)), 0.0);
-   var f0 = vec3(0.04, 0.04, 0.04);
+   var f0 = mix(vec3(0.04, 0.04, 0.04), material.base_reflectance, material.metalness);
    var total_color = vec3(0.0);
    var normal = normalize(frag_data.normal);

    for(var i: u32 = 0; i < light_count.count; i++){
        var light = lights[i];
        var light_distance = length(light.position - frag_data.world_position.xyz);
        var attenuation = 1.0 / pow(light_distance, 2.0);
        var radiance = light.color.rgb * attenuation;
        total_color += get_bdrf(surface_albedo, f0, roughness, normal, radiance, light.position, scene.camera_position, frag_data.world_position.xyz);
    }

    //return vec4(total_color, 1.0);
    var tone_mapped_color = total_color / (total_color + vec3(1.0));
    return vec4(pow(total_color, vec3(1.0/2.2)), 1.0);
}

I also had to clamp the roughness because at 0 the normal distribution function goes to 0 meaning no specular at all. Perfect mirrors don’t behave well. Speaking of not behaving, this is the result:

The gold actually looks pretty good! But the shadowing is very harsh. This is because the metal only has specular so where the light doesn’t hit will get shadowed to the extreme because there is no color transition as we’d get with diffuse lighting.

Note on Ambient Lighting

In reality the scene will have light from more directions the subtle bouncing of light off surfaces like the rug. This can be modeled using expensive and sophisticated ways like path-tracing. I think I’ll leave this topic alone though since it seems to be a bit large.

Directional Lights

Along with point lights we’d like to have directional lights. These are lights that just case one direction but have infinite size. Like the sun from far away. Might as well add them now while we’re in here. We already have the enum for it defined on light we just need a few tweaks.


-fn get_bdrf(surface_albedo: vec3, f0: vec3, roughness: f32, metalness: f32, normal: vec3, light_color: vec3, light_pos: vec3, camera_pos: vec3, frag_pos: vec3) -> vec3{
+fn get_bdrf(surface_albedo: vec3, f0: vec3, roughness: f32, metalness: f32, normal: vec3, light_color: vec3, to_light: vec3, camera_pos: vec3, frag_pos: vec3) -> vec3{
    var to_view = normalize(camera_pos - frag_pos);
-   var to_light = normalize(light_pos - frag_pos);
    var half_vector = normalize(to_view + to_light);
    var fresnel = get_fresnel(f0, to_view, half_vector);
    var ks = fresnel;
    var kd = (vec3(1.0) - ks) * (1.0 - metalness);

    var diffuse = get_diffuse(surface_albedo, light_color, normal, to_light);
    var specular = get_specular(fresnel, to_view, to_light, normal, half_vector, roughness);

    return kd * diffuse + specular;
}

We don’t need light_pos except to calculate the vector to_light. So we’ll calculate that in the main loop.

@fragment
fn fragment_main(frag_data: VertexOut) -> @location(0) vec4
{   
    var surface_albedo = textureSample(albedo_map, albedo_sampler, frag_data.uv).rgb;
    var roughness_from_map = textureSample(roughness_map, roughness_sampler, frag_data.uv).x;
    var roughness = max(mix(material.roughness, roughness_from_map, f32(material.use_specular_map)), 0.0001);
    var f0 = mix(vec3(0.04, 0.04, 0.04), material.base_reflectance, material.metalness);
    var total_color = vec3(0.0);
    var normal = normalize(frag_data.normal);

    for(var i: u32 = 0; i < light_count.count; i++){
        var light = lights[i];
        var light_distance = length(light.position - frag_data.world_position.xyz);
+       var to_light = vec3(0.0);

+       switch light.light_type {
+           case 0: {
+               to_light = normalize(light.position - frag_data.world_position.xyz);
+           }
+           case 1: {
+               to_light = normalize(-light.direction);
+           }
+           default: {}
+       }

        var attenuation = 1.0 / pow(light_distance, 2.0);
        var radiance = light.color.rgb * attenuation;


        total_color += get_bdrf(
            surface_albedo, 
            f0, 
            roughness, 
            material.metalness,
            normal, 
            radiance, 
+           to_light,
            scene.camera_position, 
            frag_data.world_position.xyz
        );
    }

    var tone_mapped_color = total_color / (total_color + vec3(1.0));
    return vec4(pow(total_color, vec3(1.0/2.2)), 1.0);
}

This, as far as I can tell, requires branching mainly because there will be 3 states (I’m just not dealing with spotlights yet). So we switch on the light type. If it’s type 0 (point) we use the previous method of getting the vector from the surface to the point. In the case of a directional light we just need the direction, but to_light is from the surface to the light so it’s actually just the inverse of the light direction. Then we plug in as normal.