From 46d269649387dbfddaf76d3ffd6da133c6b69f90 Mon Sep 17 00:00:00 2001
From: Lucas Alber <lucasd.alber@gmail.com>
Date: Tue, 26 Nov 2024 13:43:03 +0100
Subject: [PATCH] merian-shader: Rework ggx

---
 include/merian-shaders/bsdf_diffuse.glsl |   4 +-
 include/merian-shaders/bsdf_ggx.glsl     | 240 ++++++++++++++---------
 include/merian-shaders/common.glsl       |   4 +
 3 files changed, 154 insertions(+), 94 deletions(-)

diff --git a/include/merian-shaders/bsdf_diffuse.glsl b/include/merian-shaders/bsdf_diffuse.glsl
index 0c0eb83..41f99a0 100644
--- a/include/merian-shaders/bsdf_diffuse.glsl
+++ b/include/merian-shaders/bsdf_diffuse.glsl
@@ -9,10 +9,12 @@
 // n must be normalized
 #define bsdf_diffuse_sample(n, random) (make_frame(n) * sample_cos(random))
 
-// solid angle, not error checked against wodotn < 0.
+// solid angle (contains multiplication with wodotn), not error checked against wodotn < 0.
 #define bsdf_diffuse_pdf(wodotn) (INV_PI * wodotn)
 
 // solid angle
 #define bsdf_diffuse_eval(albedo) (INV_PI * albedo)
 
+#define bsdf_diffuse() (INV_PI)
+
 #endif
diff --git a/include/merian-shaders/bsdf_ggx.glsl b/include/merian-shaders/bsdf_ggx.glsl
index dc8b50f..3dd3b4d 100644
--- a/include/merian-shaders/bsdf_ggx.glsl
+++ b/include/merian-shaders/bsdf_ggx.glsl
@@ -6,151 +6,205 @@
 #include "merian-shaders/bsdf_diffuse.glsl"
 #include "merian-shaders/frames.glsl"
 
+// -----------------------------------------------------------------
+
+float bsdf_ggx_roughness_to_alpha(const float roughness) {
+    return roughness * roughness;
+}
+
 // Smiths shadow masking term
-float smith_g1(const float roughness, const float wodotn) {
-    const float alpha_sq = pow(roughness, 4);
-    return 2.0 * wodotn / (wodotn + sqrt(alpha_sq + (1.0 - alpha_sq) * MERIAN_SQUARE(wodotn)));
+float bsdf_ggx_smith_g1(const float minuswidotn, const float alpha_squared) {
+    return 2.0 * minuswidotn / (minuswidotn + sqrt(alpha_squared + (1.0 - alpha_squared) * MERIAN_SQUARE(minuswidotn)));
+}
+
+// Smiths shadow masking term
+float bsdf_ggx_smith_g1_over_minuswidotn(const float minuswidotn, const float alpha_squared) {
+    return 2.0 / (minuswidotn + sqrt(alpha_squared + (1.0 - alpha_squared) * MERIAN_SQUARE(minuswidotn)));
 }
 
 // minuswi = V
-float smith_g_over_minuswidotn(const float roughness, const float minuswidotn, const float wodotn) {
-    const float alpha = MERIAN_SQUARE(roughness);
-    const float g1 = minuswidotn * sqrt(MERIAN_SQUARE(alpha) + (1.0 - MERIAN_SQUARE(alpha)) * MERIAN_SQUARE(wodotn));
-    const float g2 = wodotn * sqrt(MERIAN_SQUARE(alpha) + (1.0 - MERIAN_SQUARE(alpha)) * MERIAN_SQUARE(minuswidotn));
+float bsdf_ggx_smith_g2_over_minuswidotn(const float wodotn, const float minuswidotn, const float alpha_squared) {
+    const float g1 = minuswidotn * sqrt(alpha_squared + (1.0 - alpha_squared) * MERIAN_SQUARE(wodotn));
+    const float g2 = wodotn * sqrt(alpha_squared + (1.0 - alpha_squared) * MERIAN_SQUARE(minuswidotn));
     return 2.0 * wodotn / (g1 + g2);
 }
 
-// minuswi = V
-vec3 bsdf_ggx_times_wodotn(const vec3 minuswi, const vec3 wo, const vec3 n, const float roughness, const vec3 F0) {
-    const vec3 H = normalize(wo + minuswi);
+float bsdf_ggx_D(const float ndoth_squared, const float alpha_squared) {
+    return alpha_squared / (M_PI * merian_square(ndoth_squared * alpha_squared + (1 - ndoth_squared)));
+}
 
-    const float wodotn = clamp(dot(n, wo), 0, 1);
-    const float minuswi_dot_h = clamp(dot(minuswi, H), 0, 1);
-    const float ndotv = clamp(dot(n, minuswi), 0, 1);
-    const float ndoth = clamp(dot(n, H), 0, 1);
+// -----------------------------------------------------------------
 
-    if (wodotn > 0) {
-        const float G = smith_g_over_minuswidotn(roughness, ndotv, wodotn);
-        const float alpha = MERIAN_SQUARE(roughness);
-        const float D = MERIAN_SQUARE(alpha) / (M_PI * MERIAN_SQUARE(MERIAN_SQUARE(ndoth) * MERIAN_SQUARE(alpha) + (1 - MERIAN_SQUARE(ndoth))));
+float bsdf_ggx_times_wodotn(const float wodotn, const float minuswidotn, const float ndoth, const float alpha) {
+    const float G2 = bsdf_ggx_smith_g2_over_minuswidotn(wodotn, minuswidotn, alpha * alpha);
+    const float D = bsdf_ggx_D(ndoth * ndoth, alpha * alpha);
 
-        const vec3 F = fresnel_schlick(minuswi_dot_h, F0);
+    return (D * G2) / 4;
+}
 
-        return F * (D * G / 4);
-    }
-    return vec3(0);
+
+float bsdf_ggx_times_wodotn(const vec3 wi, const vec3 wo, const vec3 n, const float alpha) {
+    const vec3 h = normalize(wo - wi);
+
+    const float wodotn = dot(wo, n);
+    const float minuswidotn = dot(-wi, n);
+    const float ndoth = dot(n, h);
+
+    return bsdf_ggx_times_wodotn(wodotn, minuswidotn, ndoth, alpha);
 }
 
-float bsdf_ggx_times_wodotn(const vec3 minuswi, const vec3 wo, const vec3 n, const float roughness, const float F0) {
-    const vec3 H = normalize(wo + minuswi);
+// with fresnel
+float bsdf_ggx_times_wodotn(const vec3 wi, const vec3 wo, const vec3 n, const float alpha, const float F0) {
+    const vec3 h = normalize(wo - wi);
 
-    const float wodotn = clamp(dot(n, wo), 0, 1);
-    const float minuswi_dot_h = clamp(dot(minuswi, H), 0, 1);
-    const float ndotv = clamp(dot(n, minuswi), 0, 1);
-    const float ndoth = clamp(dot(n, H), 0, 1);
+    const float wodotn = dot(wo, n);
+    const float minuswidotn = dot(-wi, n);
+    const float ndoth = dot(n, h);
 
-    if (wodotn > 0) {
-        const float G = smith_g_over_minuswidotn(roughness, ndotv, wodotn);
-        const float alpha = MERIAN_SQUARE(roughness);
-        const float D = MERIAN_SQUARE(alpha) / (M_PI * MERIAN_SQUARE(MERIAN_SQUARE(ndoth) * MERIAN_SQUARE(alpha) + (1 - MERIAN_SQUARE(ndoth))));
+    const float bsdf = bsdf_ggx_times_wodotn(wodotn, minuswidotn, ndoth, alpha);
+    const float minuswidoth = dot(-wi, h);
 
-        const float F = fresnel_schlick(minuswi_dot_h, F0);
+    const float F = fresnel_schlick(minuswidoth, F0);
 
-        return F * (D * G / 4);
-    }
-    return 0;
+    return F * bsdf;
 }
 
-float bsdf_ggx_diffuse_mix_times_wodotn(const vec3 minuswi, const vec3 wo, const vec3 n, const float roughness, const float F0) {
-    const vec3 H = normalize(wo + minuswi);
+// with fresnel
+vec3 bsdf_ggx_times_wodotn(const vec3 wi, const vec3 wo, const vec3 n, const float alpha, const vec3 F0) {
+    const vec3 h = normalize(wo - wi);
+
+    const float wodotn = dot(wo, n);
+    const float minuswidotn = dot(-wi, n);
+    const float ndoth = dot(n, h);
+
+    const float bsdf = bsdf_ggx_times_wodotn(wodotn, minuswidotn, ndoth, alpha);
+    const float minuswidoth = dot(-wi, h);
+
+    const vec3 F = fresnel_schlick(minuswidoth, F0);
+
+    return F * bsdf;
+}
+
+
+// -----------------------------------------------------------------
+
+float bsdf_ggx_diffuse_mix_times_wodotn(const vec3 wi, const vec3 wo, const vec3 n, const float alpha, const float F0) {
+    const vec3 h = normalize(wo - wi);
 
     const float wodotn = clamp(dot(n, wo), 0, 1);
     if (wodotn <= 0) {
         return 0;
     }
 
-    const float minuswi_dot_h = clamp(dot(minuswi, H), 0, 1);
-    const float ndotv = clamp(dot(n, minuswi), 0, 1);
-    const float ndoth = clamp(dot(n, H), 0, 1);
+    const float minuswidotn = dot(-wi, n);
+    const float ndoth = dot(n, h);
+    const float minuswidoth = dot(-wi, h);
 
-    const float G = smith_g_over_minuswidotn(roughness, ndotv, wodotn);
-    const float alpha = MERIAN_SQUARE(roughness);
-    const float D = MERIAN_SQUARE(alpha) / (M_PI * MERIAN_SQUARE(MERIAN_SQUARE(ndoth) * MERIAN_SQUARE(alpha) + (1 - MERIAN_SQUARE(ndoth))));
+    const float F = fresnel_schlick(minuswidoth, F0);
 
-    const float F = fresnel_schlick(minuswi_dot_h, F0);
-
-    return mix(INV_PI * wodotn, (D * G / 4), F);
+    return mix(bsdf_diffuse() * wodotn,
+               bsdf_ggx_times_wodotn(wodotn, minuswidotn, ndoth, alpha),
+               F);
     
 }
 
-// returns a half vector in tangent space
-vec3 bsdf_ggx_sample_H(const vec2 random, const float roughness) {
-    const float phi = TWO_PI * random.x;
-    const float cosTheta = sqrt((1 - random.y) / (1 + (pow(roughness, 4) - 1) * random.y));
-    const float sinTheta = sqrt(1 - cosTheta * cosTheta);
-
-    return vec3(sinTheta * cos(phi), 
-                sinTheta * sin(phi),
-                cosTheta);
+/**
+ * Sampling Visible GGX Normals with Spherical Caps
+ * Jonathan Dupuy, Anis Benyoub
+ * High Performance Graphics 2023
+ * 
+ * https://gist.github.com/jdupuy/4c6e782b62c92b9cb3d13fbb0a5bd7a0
+ */
+vec3 bsdf_ggx_VNDF_H_sample(const vec3 minuswi, const vec2 alpha, const vec2 random) {
+    // warp to the hemisphere configuration
+    vec3 wiStd = normalize(vec3(minuswi.xy * alpha, minuswi.z));
+    // sample a spherical cap in (-minuswi.z, 1]
+    float phi = (2.0f * random.x - 1.0f) * M_PI;
+    float z = fma((1.0f - random.y), (1.0f + wiStd.z), -wiStd.z);
+    float sinTheta = sqrt(clamp(1.0f - z * z, 0.0f, 1.0f));
+    float x = sinTheta * cos(phi);
+    float y = sinTheta * sin(phi);
+    vec3 c = vec3(x, y, z);
+    // compute halfway direction as standard normal
+    vec3 wmStd = c + wiStd;
+    // warp back to the ellipsoid configuration
+    vec3 wm = normalize(vec3(wmStd.xy * alpha, wmStd.z));
+    // return final normal
+    return wm;
 }
 
-// https://github.com/NVIDIAGameWorks/donut
-// MIT Licensed
-vec3 bsdf_ggx_VNDF_sample_H(const vec2 random, const float roughness, const vec3 Ve) {
-    const float alpha = MERIAN_SQUARE(roughness);
-    const vec3 Vh = normalize(vec3(alpha * Ve.x, alpha * Ve.y, Ve.z));
-
-    const float lensq = MERIAN_SQUARE(Vh.x) + MERIAN_SQUARE(Vh.y);
-    const vec3 T1 = lensq > 0.0 ? vec3(-Vh.y, Vh.x, 0.0) * (1 / sqrt(lensq)) : vec3(1.0, 0.0, 0.0);
-    const vec3 T2 = cross(Vh, T1);
-
-    const float r = sqrt(random.x);// sqrt(random.x * ndf_trim);
-    const float phi = 2.0 * M_PI * random.y;
-    const float t1 = r * cos(phi);
-    float t2 = r * sin(phi);
-    const float s = 0.5 * (1.0 + Vh.z);
-    t2 = (1.0 - s) * sqrt(1.0 - MERIAN_SQUARE(t1)) + s * t2;
-
-    const vec3 Nh = t1 * T1 + t2 * T2 + sqrt(max(0.0, 1.0 - MERIAN_SQUARE(t1) - MERIAN_SQUARE(t2))) * Vh;
+/**
+ * Sampling Visible GGX Normals with Spherical Caps
+ * Jonathan Dupuy, Anis Benyoub
+ * High Performance Graphics 2023
+ * 
+ * https://gist.github.com/jdupuy/4c6e782b62c92b9cb3d13fbb0a5bd7a0
+ */
+vec3 bsdf_ggx_VNDF_H_sample(const vec3 minuswi, const vec3 n, const float alpha, const vec2 random) {
+    // decompose the vector in parallel and perpendicular components
+    const vec3 wi_z = n * dot(minuswi, n);
+    const vec3 wi_xy = minuswi - wi_z;
+    // warp to the hemisphere configuration
+    const vec3 wiStd = normalize(wi_z - alpha * wi_xy);
+    // sample a spherical cap in (-wiStd.z, 1]
+    const float wiStd_z = dot(wiStd, n);
+    const float phi = (2.0f * random.x - 1.0f) * M_PI;
+    const float z = (1.0f - random.y) * (1.0f + wiStd_z) - wiStd_z;
+    const float sinTheta = sqrt(clamp(1.0f - z * z, 0.0f, 1.0f));
+    const float x = sinTheta * cos(phi);
+    const float y = sinTheta * sin(phi);
+    const vec3 cStd = vec3(x, y, z);
+    // reflect sample to align with normal
+    const vec3 up = vec3(0, 0, 1);
+    const vec3 wr = n + up;
+    const vec3 c = dot(wr, cStd) * wr / wr.z - cStd;
+    // compute halfway direction as standard normal
+    const vec3 wmStd = c + wiStd;
+    const vec3 wmStd_z = n * dot(n, wmStd);
+    const vec3 wmStd_xy = wmStd_z - wmStd;
+    // warp back to the ellipsoid configuration
+    const vec3 wm = normalize(wmStd_z + alpha * wmStd_xy);
+    // return final normal
+    return wm;
+}
 
-    return vec3(alpha * Nh.x,
-                alpha * Nh.y,
-                max(0.0, Nh.z)
-        );
+vec3 bsdf_ggx_VNDF_sample(const vec3 wi, const vec3 n, const float alpha, const vec2 random) {
+    const vec3 h = bsdf_ggx_VNDF_H_sample(-wi, n, alpha, random);
+    return reflect(wi, h);
 }
 
-float bsdf_ggx_VNDF_pdf(const float roughness, const vec3 N, const vec3 minus_wi, const vec3 wo) {
-    const vec3 H = normalize(wo + minus_wi);
-    const float ndoth = clamp(dot(N, H), 0, 1);
-    const float minuswi_dot_h = clamp(dot(minus_wi, H), 0, 1);
+float bsdf_ggx_VNDF_pdf(const vec3 wi, const vec3 wo, const vec3 n, const float alpha) {
+    const vec3 h = normalize(wo - wi);
+    const float minuswidotn = dot(-wi, n);
+    const float ndoth = clamp(dot(n, h), 0, 1);
 
-    const float alpha = MERIAN_SQUARE(roughness);
-    const float D = MERIAN_SQUARE(alpha) / (M_PI * MERIAN_SQUARE(MERIAN_SQUARE(ndoth) * MERIAN_SQUARE(alpha) + (1 - MERIAN_SQUARE(ndoth))));
-    return (minuswi_dot_h > 0.0) ? D / (4.0 * minuswi_dot_h) : 0.0;
+    const float G1 = bsdf_ggx_smith_g1_over_minuswidotn(minuswidotn, alpha * alpha);
+    const float D = bsdf_ggx_D(ndoth * ndoth, alpha * alpha);
+
+    return (D * G1) / 4;
 }
 
-float bsdf_ggx_diffuse_mix_pdf(const vec3 minuswi, const vec3 wo, const vec3 n, const float roughness) {
+float bsdf_ggx_diffuse_mix_pdf(const vec3 wi, const vec3 wo, const vec3 n, const float alpha) {
     const float wodotn = dot(wo, n);
+
     if (wodotn <= 0) {
         return 0;
     }
 
     const float diffuse_pdf = bsdf_diffuse_pdf(wodotn);
-    const float ggx_vndf_pdf = bsdf_ggx_VNDF_pdf(roughness, n, minuswi, wo);
+    const float ggx_vndf_pdf = bsdf_ggx_VNDF_pdf(wi, wo, n, alpha);
 
-    const float fresnel = fresnel_schlick(dot(minuswi, n), 0.02);
+    const float fresnel = fresnel_schlick(dot(-wi, n), 0.02);
     return mix(diffuse_pdf, ggx_vndf_pdf, fresnel);
 }
 
-vec3 bsdf_ggx_diffuse_mix_sample(const vec3 minuswi, const vec3 n, const float roughness, const vec3 random) {
-    const float fresnel = fresnel_schlick(dot(minuswi, n), 0.02);
+vec3 bsdf_ggx_diffuse_mix_sample(const vec3 wi, const vec3 n, const float alpha, const vec3 random) {
+    const float fresnel = fresnel_schlick(dot(-wi, n), 0.02);
 
     if (fresnel < random.x) {
         return bsdf_diffuse_sample(n, random.yz);
     } else {
-        const mat3x3 frame = make_frame(n);
-        const vec3 H = normalize(bsdf_ggx_VNDF_sample_H(random.yz, roughness, world_to_frame(frame, minuswi)));
-        return reflect(-minuswi, frame_to_world(frame, H));
+        return bsdf_ggx_VNDF_sample(wi, n, alpha, random.yz);
     }
 }
 
diff --git a/include/merian-shaders/common.glsl b/include/merian-shaders/common.glsl
index 652ce43..df83939 100644
--- a/include/merian-shaders/common.glsl
+++ b/include/merian-shaders/common.glsl
@@ -11,6 +11,10 @@
 
 #define MERIAN_SQUARE(x) ((x) * (x))
 
+float merian_square(const float x) {
+    return x * x;
+}
+
 #define MERIAN_WORKGROUP_INDEX (gl_WorkGroupID.x + gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.z * gl_NumWorkGroups.x * gl_NumWorkGroups.y)
 #define MERIAN_GLOBAL_INVOCATION_INDEX (MERIAN_WORKGROUP_INDEX * gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z + gl_LocalInvocationIndex)