brdf.wgsl

webgpu-based path tracer

src/shaders/brdf.wgsl

6.03 KB
// computes the geometry (shadowing/masking) term using smith's method.
fn smith_geometry(normal: vec3<f32>, view_dir: vec3<f32>, light_dir: vec3<f32>, roughness: f32) -> f32 {
    let alpha = roughness * roughness;
    let n_dot_v = max(dot(normal, view_dir), 0.0);
    let n_dot_l = max(dot(normal, light_dir), 0.0);
    let k = (alpha + 1.0) * (alpha + 1.0) / 8.0;
    let geom_v = n_dot_v / (n_dot_v * (1.0 - k) + k);
    let geom_l = n_dot_l / (n_dot_l * (1.0 - k) + k);
    return geom_v * geom_l;
}

// computes the ggx normal distribution function (ndf) d.
fn ggx_distribution(normal: vec3<f32>, half_vec: vec3<f32>, roughness: f32) -> f32 {
    let rgh = max(0.03, roughness);
    let alpha = rgh * rgh;
    let alpha2 = alpha * alpha;
    let n_dot_h = max(dot(normal, half_vec), 0.001);
    let denom = (n_dot_h * n_dot_h) * (alpha2 - 1.0) + 1.0;
    return alpha2 / (PI * denom * denom);
}

// samples a half–vector (h) from the ggx distribution in tangent space.
fn ggx_sample_vndf(view_dir: vec3<f32>, normal: vec3<f32>, roughness: f32, noise: vec2<f32>) -> vec3<f32> {
    // build local frame (t, b, n)
    let up: vec3<f32> = select(vec3<f32>(0.0, 0.0, 1.0), vec3<f32>(1.0, 0.0, 0.0), abs(normal.z) < 0.999);
    let tangent = normalize(cross(up, normal));
    let bitangent = cross(normal, tangent);

    // transform view direction to local space
    let v = normalize(vec3<f32>(dot(view_dir, tangent), dot(view_dir, bitangent), dot(view_dir, normal)));

    // stretch view vector
    let a = roughness;
    let vh = normalize(vec3<f32>(a * v.x, a * v.y, v.z));

    // orthonormal basis
    let lensq = vh.x * vh.x + vh.y * vh.y;
    let t1 = select(vec3<f32>(1.0, 0.0, 0.0), vec3<f32>(- vh.y, vh.x, 0.0) / sqrt(lensq), lensq > 1e-6);
    let t2 = cross(vh, t1);

    // sample point with polar coordinates
    let r = sqrt(noise.x);
    let phi = 2.0 * PI * noise.y;
    let r1 = r * cos(phi);
    let r2 = r * sin(phi);
    let s = 0.5 * (1.0 + vh.z);
    let t2_ = mix(sqrt(1.0 - r1 * r1), r2, s);

    // sampled halfway vector in local space
    let nh = r1 * t1 + t2_ * t2 + sqrt(max(0.0, 1.0 - r1 * r1 - t2_ * t2_)) * vh;

    // unstretch
    let h = normalize(vec3<f32>(a * nh.x, a * nh.y, max(0.0, nh.z)));

    // transform back to world space
    return normalize(h.x * tangent + h.y * bitangent + h.z * normal);
}

fn ggx_specular_sample(view_dir: vec3<f32>, normal: vec3<f32>, seed: vec2<f32>, roughness: f32) -> vec3<f32> {
    let h = ggx_sample_vndf(view_dir, normal, roughness, seed);
    let r = reflect(- view_dir, h);
    return select(normalize(r), vec3<f32>(0.0), dot(r, normal) <= EPSILON);
}

// cosine-weighted hemisphere sample in local space, then converted to world space.
fn cosine_hemisphere_sample(normal: vec3<f32>, noise: vec2<f32>) -> vec3<f32> {
    // var current_seed = seed;
    let r1 = noise.x;
    let r2 = noise.y;

    // let r1 = uniform_float(state);
    // let r2 = uniform_float(state);

    let phi = 2.0 * PI * r2;
    let cos_theta = sqrt(1.0 - r1);
    let sin_theta = sqrt(r1);

    // sample in local coordinates (with z as the normal)
    let local_sample = vec3<f32>(sin_theta * cos(phi), sin_theta * sin(phi), cos_theta);

    // build tangent space and transform sample to world space.
    // let up = select(vec3<f32>(0.0, 0.0, 1.0), vec3<f32>(1.0, 0.0, 0.0), abs(normal.z) > 0.999);

    let tangent = normalize(cross(select(vec3<f32>(0.0, 1.0, 0.0), vec3<f32>(1.0, 0.0, 0.0), abs(normal.y) > 0.99), normal));
    let bitangent = cross(normal, tangent);

    // let tangent = normalize(cross(up, normal));
    // let bitangent = cross(normal, tangent);
    let world_dir = normalize(local_sample.x * tangent + local_sample.y * bitangent + local_sample.z * normal);
    return world_dir;
}

fn eval_f0(metallic: f32, albedo: vec3<f32>) -> vec3<f32> {
    let dielectric_f0 = vec3<f32>(0.04);
    return mix(dielectric_f0, albedo, metallic);
}

fn fresnel_schlick_roughness(cos_theta: f32, f0: vec3<f32>, roughness: f32) -> vec3<f32> {
    let one_minus_cos = 1.0 - cos_theta;
    let factor = pow(one_minus_cos, 5.0);
    let fresnel = f0 + (max(vec3f(1.0 - roughness), f0) - f0) * factor;
    return fresnel;
}

fn disney_diffuse(albedo: vec3<f32>, roughness: f32, n_dot_l: f32, n_dot_v: f32, l_dot_h: f32) -> vec3<f32> {
    let fd90 = 0.5 + 2.0 * l_dot_h * l_dot_h * roughness;
    let light_scatter = 1.0 + (fd90 - 1.0) * pow(1.0 - n_dot_l, 5.0);
    let view_scatter = 1.0 + (fd90 - 1.0) * pow(1.0 - n_dot_v, 5.0);
    return albedo * light_scatter * view_scatter * (1.0 / PI);
}

fn eval_brdf(normal: vec3<f32>, view_dir: vec3<f32>, light_dir: vec3<f32>, material: Material) -> vec3<f32> {
    let n = normal;
    let v = view_dir;
    let l = light_dir;
    let h = normalize(v + l);

    let ndot_l = max(dot(n, l), EPSILON);
    let ndot_v = max(dot(n, v), EPSILON);
    let ndot_h = max(dot(n, h), EPSILON);
    let vdot_h = max(dot(v, h), EPSILON);
    let ldot_h = max(dot(l, h), EPSILON);

    let f0 = mix(vec3f(0.04), material.albedo.rgb, material.metallic);
    let f = fresnel_schlick_roughness(vdot_h, f0, material.roughness);

    // frostbite specular
    let d = ggx_distribution(n, h, material.roughness);
    let g = smith_geometry(n, v, l, material.roughness);
    let spec = (d * g * f) / (4.0 * ndot_v * ndot_l + 0.001);

    let diffuse = (1.0 - material.metallic) * disney_diffuse(material.albedo.rgb, material.roughness, ndot_l, ndot_v, ldot_h);
    return diffuse + spec;
}

fn ggx_pdf(view_dir: vec3<f32>, normal: vec3<f32>, h: vec3<f32>, roughness: f32) -> f32 {
    let d = ggx_distribution(normal, h, roughness);
    let n_dot_h = max(dot(normal, h), EPSILON);
    let v_dot_h = max(dot(view_dir, h), EPSILON);
    // change–of–variables: pdf(r) = d(h) * n_dot_h / (4 * v_dot_h)
    return (d * n_dot_h) / (4.0 * v_dot_h);
}

fn cosine_pdf(normal: vec3<f32>, dir: vec3<f32>) -> f32 {
    // cosine weighted density: d = cos(theta) / PI.
    let cos_theta = max(dot(normal, dir), EPSILON);
    if cos_theta <= 0.0 {
        return 0.0;
        // no contribution when direction is opposite or perpendicular to the normal.
    }
    return cos_theta / PI;
}