gi: volume path tracing (for SSS and multiscatter sheen BSDFs)

LICENSE

@@ -2427,6 +2427,7 @@ distribution.
     MDL SDK
     MDL Core, Aux, Support Definitions,
     Omniverse MDL files
+    Advanced Samples for the NVIDIA OptiX 7 Ray Tracing SDK
 ============================================================
 
 
@@ -2462,6 +2463,8 @@ distribution.
   Copyright 2022 NVIDIA Corporation. All rights reserved.
 * For the MDL core definitions, see also this required NOTICE:
   Copyright 2022 NVIDIA Corporation. All rights reserved.
+* For the Advanced Samples for the NVIDIA OptiX 7 Ray Tracing SDK, see also this required NOTICE:
+  Copyright (c) 2013-2024, NVIDIA CORPORATION. All rights reserved.
 
 
 ============================================================

src/gi/shaders/rp_main.chit

@@ -210,20 +210,34 @@ void main()
 
     if (mediumIdx > 0)
     {
+        vec3 weight;
+        float distance = gl_HitTEXT;
+
 #if MEDIUM_STACK_SIZE == 0
         vec3 sigma_a = mdl_volume_absorption_coefficient(shading_state);
+
+        vec3 transmittance = exp(-sigma_a * distance);
+
+        weight = transmittance;
 #else
         Medium m = rayPayload.media[mediumIdx - 1];
-        vec3 sigma_a = m.sigma_t - m.sigma_s;
 
         prevMediumIor = m.ior;
         if (mediumIdx > 1)
         {
             nextMediumIor = rayPayload.media[mediumIdx - 2].ior;
         }
+
+        vec3 transmittance = exp(-m.sigma_t * distance);
+
+        // Unlike the Optix 7 sample, we consider the PDF for the last walk segment
+        // https://github.com/mmp/pbrt-v3/blob/13d871faae88233b327d04cda24022b8bb0093ee/src/media/homogeneous.cpp#L75
+        float pdf = dot(rayPayload.walkSegmentPdf, transmittance);
+
+        weight = transmittance / pdf;
 #endif
 
-        throughput *= exp(-sigma_a * gl_HitTEXT);
+        throughput *= weight;
     }
 
     const vec3 iorCurrent = (isFrontFace || thinWalled) ? prevMediumIor : vec3(BSDF_USE_MATERIAL_IOR);
@@ -407,6 +421,9 @@ void main()
         mediumIdx = (1 - mediumIdx); // toggle between inside and outside
 #endif
 
+        // medium changed -> reset walk
+        rayPayload.bitfield &= ~SHADE_RAY_PAYLOAD_WALK_MASK;
+
         shadeRayPayloadSetMediumIdx(rayPayload, mediumIdx);
     }
 

src/gi/shaders/rp_main.miss

@@ -12,6 +12,28 @@
 
 layout(location = PAYLOAD_INDEX_SHADE) rayPayloadInEXT ShadeRayPayload rayPayload;
 
+#if MEDIUM_STACK_SIZE > 0
+void stepVolume(uint mediumIdx)
+{
+    float distance = gl_RayTmaxEXT;
+
+    Medium m = rayPayload.media[mediumIdx - 1];
+
+    vec3 transmittance = exp(m.sigma_t * -distance);
+
+    vec3 density = m.sigma_t * transmittance;
+
+    float pdf = dot(rayPayload.walkSegmentPdf, density);
+
+    rayPayload.throughput *= (m.sigma_s * transmittance) / pdf;
+
+    rayPayload.ray_origin += gl_WorldRayDirectionEXT * distance;
+
+    rayPayload.bitfield |= SHADE_RAY_PAYLOAD_VOLUME_WALK_MISS_FLAG;
+    shadeRayPayloadIncrementWalk(rayPayload);
+}
+#endif
+
 // Optimized implementation from GLM with only cross products:
 // https://github.com/g-truc/glm/blob/47585fde0c49fa77a2bf2fb1d2ead06999fd4b6e/glm/detail/type_quat.inl#L356-L363
 vec3 quatRotateDir(vec4 q, vec3 dir)
@@ -32,6 +54,17 @@ vec3 sampleDomeLight(uint domeLightIndex, vec3 rayDir)
 
 void main()
 {
+#if MEDIUM_STACK_SIZE > 0
+    uint mediumIdx = shadeRayPayloadGetMediumIdx(rayPayload);
+
+    if (mediumIdx > 0)
+    {
+        stepVolume(mediumIdx);
+
+        return;
+    }
+#endif
+
     rayPayload.bitfield |= SHADE_RAY_PAYLOAD_TERMINATE_FLAG;
 
 #if AOV_ID != AOV_ID_COLOR

src/gi/shaders/rp_main.rgen

@@ -25,6 +25,66 @@ layout(buffer_reference, hitobjectshaderrecordnv) buffer ShaderRecordBuffer
 };
 #endif
 
+//
+// This path tracing implementation was initially based on the MDL SDK DXR sample:
+// https://github.com/NVIDIA/MDL-SDK/tree/master/examples/mdl_sdk/dxr
+// but has since then received many changes such as flattened ray tracing calls.
+//
+// The volume path tracing part is partly based on the Optix 7 example renderer:
+// https://github.com/NVIDIA/OptiX_Apps/tree/master/apps/MDL_renderer
+//
+
+#if MEDIUM_STACK_SIZE > 0
+float sampleDistance(vec3 albedo, vec3 throughput, vec3 sigma_t, float xi, out vec3 pdf)
+{
+    vec3 weights = throughput * albedo;
+
+    float sum = weights.x + weights.y + weights.z;
+
+    pdf = (sum > 1e-9) ? (weights / sum) : vec3(1.0 / 3.0);
+
+    if (xi < pdf.x)
+    {
+        return sigma_t.x;
+    }
+    else if (xi < (pdf.x + pdf.y))
+    {
+        return sigma_t.y;
+    }
+    else
+    {
+        return sigma_t.z;
+    }
+}
+
+// cmp. d'Eon: "A Hitchhiker's Guide to Multiple Scattering"
+float sampleHenyeyGreensteinCos(float r, float g)
+{
+    if (abs(g) < 1e-3f)
+    {
+        return 1.0 - 2.0 * r; // isotropic medium
+    }
+
+    float s = (1.0 - g * g) / (1.0 - g + 2.0 * g * r);
+
+    return (1.0 + g * g - s * s) / (2.0 * g);
+}
+
+void sampleVolumeScatteringDirection(vec2 xi, float bias, inout vec3 dir)
+{
+    float cos_theta = sampleHenyeyGreensteinCos(xi.x, bias);
+
+    float sin_theta = sqrt(max(0.0, 1.0 - cos_theta * cos_theta));
+
+    float phi = 2.0 * PI * xi.y;
+
+    vec3 t, b;
+    orthonormal_basis(dir, t, b);
+
+    dir = t * sin_theta * cos(phi) + b * sin_theta * sin(phi) + dir * cos_theta;
+}
+#endif
+
 bool russian_roulette(in float random_float, inout vec3 throughput)
 {
     float max_throughput = max(throughput.r, max(throughput.g, throughput.b));
@@ -44,12 +104,15 @@ vec3 evaluate_sample(vec3 ray_origin,
                      vec3 ray_dir,
                      RNG_STATE_TYPE rng_state)
 {
-    rayPayload.throughput = vec3(1.0);
-    rayPayload.bitfield   = 0;
-    rayPayload.radiance   = vec3(0.0);
-    rayPayload.rng_state  = rng_state;
-    rayPayload.ray_origin = ray_origin;
-    rayPayload.ray_dir    = ray_dir;
+    rayPayload.throughput     = vec3(1.0);
+    rayPayload.bitfield       = 0;
+    rayPayload.radiance       = vec3(0.0);
+    rayPayload.rng_state      = rng_state;
+    rayPayload.ray_origin     = ray_origin;
+    rayPayload.ray_dir        = ray_dir;
+#if MEDIUM_STACK_SIZE > 0
+    rayPayload.walkSegmentPdf = vec3(0.0);
+#endif
 
     // Correct clip plane curving by getting the hypothenuse length (our current ray)
     // from the adjacent edge (the forward vector) the angle between both.
@@ -72,6 +135,39 @@ vec3 evaluate_sample(vec3 ray_origin,
              break;
         }
 
+        float tMin = (bounce == 0) ? clipRange.x : 0.0;
+        float tMax = (bounce == 0) ? clipRange.y : FLOAT_MAX;
+
+        // Calculate distance for volume sampling
+#if MEDIUM_STACK_SIZE > 0
+        uint mediumIdx = shadeRayPayloadGetMediumIdx(rayPayload);
+
+        if (mediumIdx > 0)
+        {
+            Medium m = rayPayload.media[mediumIdx - 1];
+
+            rayPayload.walkSegmentPdf = vec3(1.0);
+
+            uint walkLength = shadeRayPayloadGetWalk(rayPayload);
+            bool hasScattering = any(greaterThan(m.sigma_s, vec3(0.0)));
+
+            if (hasScattering && walkLength <= PC.maxVolumeWalkLength)
+            {
+                vec3 albedo = safe_div(m.sigma_s, m.sigma_t);
+
+#ifdef RAND_4D
+                vec2 xi = rng4d_next4f(rayPayload.rng_state).xy;
+#else
+                vec2 xi = rng1d_next2f(rayPayload.rng_state);
+#endif
+
+                float s = sampleDistance(albedo, rayPayload.throughput, m.sigma_t, xi.x, rayPayload.walkSegmentPdf);
+
+                tMax = -log(1.0 - xi.y) / s; // collision free distance (mean free path)
+            }
+        }
+#endif
+
         // Closest hit shading
         rayPayload.neeContrib = vec3(0.0);
 
@@ -87,9 +183,9 @@ vec3 evaluate_sample(vec3 ray_origin,
             2,                     // sbt record stride
             0,                     // miss index
             rayPayload.ray_origin, // ray origin
-            clipRange.x,           // ray min range
+            tMin,                  // ray min range
             rayPayload.ray_dir,    // ray direction
-            clipRange.y,           // ray max range
+            tMax,                  // ray max range
             PAYLOAD_INDEX_SHADE    // payload
         );
 
@@ -113,9 +209,9 @@ vec3 evaluate_sample(vec3 ray_origin,
             2,                     // sbt record stride
             0,                     // miss index
             rayPayload.ray_origin, // ray origin
-            clipRange.x,           // ray min range
+            tMin,                  // ray min range
             rayPayload.ray_dir,    // ray direction
-            clipRange.y,           // ray max range
+            tMax,                  // ray max range
             PAYLOAD_INDEX_SHADE    // payload
         );
 #endif
@@ -184,10 +280,26 @@ vec3 evaluate_sample(vec3 ray_origin,
             }
         }
 
+        // Handle volume miss
+#if MEDIUM_STACK_SIZE > 0
+        if ((rayPayload.bitfield & SHADE_RAY_PAYLOAD_VOLUME_WALK_MISS_FLAG) != 0)
+        {
+#ifdef RAND_4D
+            vec2 xi = rng4d_next4f(rayPayload.rng_state).xy;
+#else
+            vec2 xi = rng1d_next2f(rayPayload.rng_state);
+#endif
+
+            float bias = rayPayload.media[mediumIdx - 1].bias;
+
+            sampleVolumeScatteringDirection(xi, bias, rayPayload.ray_dir);
+
+            rayPayload.bitfield &= ~SHADE_RAY_PAYLOAD_VOLUME_WALK_MISS_FLAG;
+        }
+#endif
+
         // Update bounce count
         rayPayload.bitfield++;
-
-        clipRange = vec2(0.0, FLOAT_MAX); // Only clip primary rays
     }
 
 #if AOV_ID == AOV_ID_DEBUG_BOUNCES

src/gi/shaders/rp_main_payload.glsl

@@ -1,5 +1,6 @@
 #include "common.glsl"
 
+#define SHADE_RAY_PAYLOAD_VOLUME_WALK_MISS_FLAG 0x40000000u
 #define SHADE_RAY_PAYLOAD_MEDIUM_IDX_MASK 0x0f000000u
 #define SHADE_RAY_PAYLOAD_MEDIUM_IDX_OFFSET 24
 #define SHADE_RAY_PAYLOAD_WALK_MASK 0x00fff000u
@@ -20,7 +21,8 @@ struct ShadeRayPayload
     /* inout */ vec3 throughput;
 
     /*               1000 0000 0000 0000 0000 0000 0000 0000 terminate
-     *               0111 0000 0000 0000 0000 0000 0000 0000 unused
+     *               0100 0000 0000 0000 0000 0000 0000 0000 volume walk miss
+     *               0011 0000 0000 0000 0000 0000 0000 0000 unused
      *               0000 1111 0000 0000 0000 0000 0000 0000 medium index [0, 256)
      *               0000 0000 1111 1111 1111 0000 0000 0000 walk length [0, 4096)
      *               0000 0000 0000 0000 0000 1111 1111 1111 bounces [0, 4096) */
@@ -32,6 +34,7 @@ struct ShadeRayPayload
 
 #if MEDIUM_STACK_SIZE > 0
     /* inout */ Medium media[MEDIUM_STACK_SIZE];
+    /* inout */ vec3 walkSegmentPdf;
 #endif
 
     /* out */   vec3 ray_origin;
@@ -50,6 +53,23 @@ struct ShadowRayPayload
     /* out */   bool shadowed;
 };
 
+#if MEDIUM_STACK_SIZE > 0
+void shadeRayPayloadIncrementWalk(inout ShadeRayPayload payload)
+{
+    uint bitfield = payload.bitfield;
+    bitfield &= SHADE_RAY_PAYLOAD_WALK_MASK;
+    bitfield = min(bitfield + 1, SHADE_RAY_PAYLOAD_WALK_MASK);
+
+    payload.bitfield &= ~SHADE_RAY_PAYLOAD_WALK_MASK;
+    payload.bitfield |= bitfield;
+}
+
+uint shadeRayPayloadGetWalk(in ShadeRayPayload payload)
+{
+    return (payload.bitfield & SHADE_RAY_PAYLOAD_WALK_MASK) >> SHADE_RAY_PAYLOAD_WALK_OFFSET;
+}
+#endif
+
 uint shadeRayPayloadGetMediumIdx(in ShadeRayPayload payload)
 {
     uint bitfield = payload.bitfield;