GPU/wavefront integrator: don't hard-code HenyeyGreenstein as the only phase function.

Now a MultiWorkQueue templated on the phase function type is used for
enqueuing work items for sampling direct/indirect lighting in media.  With
Henyey-Greenstein as the only phase function currently available, this
reduces to the previous behavior, but it makes it easier to add new phase
functions in the future.

Also piped through time in a few places it was missing in the
medium-related work queues.

src/pbrt/media.h

@@ -59,6 +59,8 @@ class HGPhaseFunction {
     PBRT_CPU_GPU
     Float PDF(Vector3f wo, Vector3f wi) const { return p(wo, wi); }
 
+    static const char *Name() { return "Henyey-Greenstein"; }
+
     std::string ToString() const;
 
   private:

src/pbrt/wavefront/integrator.cpp

@@ -285,7 +285,14 @@ WavefrontPathIntegrator::WavefrontPathIntegrator(Allocator alloc, ParsedScene &s
 
     if (haveMedia) {
         mediumSampleQueue = alloc.new_object<MediumSampleQueue>(maxQueueSize, alloc);
-        mediumScatterQueue = alloc.new_object<MediumScatterQueue>(maxQueueSize, alloc);
+
+        // TODO: in the presence of multiple PhaseFunction implementations,
+        // it could be worthwhile to see which are present in the scene and
+        // then initialize havePhase accordingly...
+        pstd::array<bool, PhaseFunction::NumTags()> havePhase;
+        havePhase.fill(true);
+        mediumScatterQueue =
+            alloc.new_object<MediumScatterQueue>(maxQueueSize, alloc, havePhase);
     }
 
     stats = alloc.new_object<Stats>(maxDepth, alloc);

src/pbrt/wavefront/integrator.h

@@ -67,6 +67,8 @@ class WavefrontPathIntegrator {
 
     void TraceShadowRays(int depth);
     void SampleMediumInteraction(int depth);
+    template <typename PhaseFunction>
+    void SampleMediumScattering(int depth);
     void SampleSubsurface(int depth);
 
     void HandleEscapedRays(int depth);

src/pbrt/wavefront/media.cpp

@@ -32,6 +32,16 @@ static inline void rescale(SampledSpectrum &T_hat, SampledSpectrum &lightPathPDF
     }
 }
 
+struct SampleMediumScatteringCallback {
+    int depth;
+    WavefrontPathIntegrator *integrator;
+    template <typename PhaseFunction>
+    void operator()() {
+        integrator->SampleMediumScattering<PhaseFunction>(depth);
+    }
+};
+
+
 // WavefrontPathIntegrator Participating Media Methods
 void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
     RayQueue *nextRayQueue = NextRayQueue(depth);
@@ -114,13 +124,16 @@ void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
                         T_hat *= T_maj * sigma_s;
                         uniPathPDF *= T_maj * sigma_s;
 
-                        // TODO: don't hard code a phase function.
-                        const HGPhaseFunction *phase =
-                            intr.phase.CastOrNullptr<HGPhaseFunction>();
                         // Enqueue medium scattering work.
-                        mediumScatterQueue->Push(MediumScatterWorkItem{
-                            intr.p(), lambda, T_hat, uniPathPDF, *phase, -ray.d,
-                            w.etaScale, ray.medium, w.pixelIndex});
+                        auto enqueue = [=](auto ptr) {
+                            using PhaseFunction =
+                                typename std::remove_reference_t<decltype(*ptr)>;
+                            mediumScatterQueue->Push(MediumScatterWorkItem<PhaseFunction>{
+                                    intr.p(), lambda, T_hat, uniPathPDF, ptr, -ray.d,
+                                    ray.time, w.etaScale, ray.medium, w.pixelIndex});
+                        };
+                        intr.phase.Dispatch(enqueue);
+
                         scattered = true;
 
                         return false;
@@ -183,7 +196,7 @@ void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
             const MixMaterial *mix = material.CastOrNullptr<MixMaterial>();
             while (mix) {
                 SurfaceInteraction intr(w.pi, w.uv, w.wo, w.dpdus, w.dpdvs, w.dndus,
-                                        w.dndvs, 0. /* time */, false /* flip normal */);
+                                        w.dndvs, ray.time, false /* flip normal */);
                 intr.faceIndex = w.faceIndex;
                 MaterialEvalContext ctx(intr);
                 material = mix->ChooseMaterial(BasicTextureEvaluator(), ctx);
@@ -236,15 +249,20 @@ void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
     if (depth == maxDepth)
         return;
 
+    ForEachType(SampleMediumScatteringCallback{depth, this}, PhaseFunction::Types());
+}
+
+template <typename Phase>
+void WavefrontPathIntegrator::SampleMediumScattering(int depth) {
     RayQueue *currentRayQueue = CurrentRayQueue(depth);
+    RayQueue *nextRayQueue = NextRayQueue(depth);
 
-    using PhaseFunction = HGPhaseFunction;
-    std::string desc = std::string("Sample direct/indirect - Henyey Greenstein");
-    ForAllQueued(
-        desc.c_str(), mediumScatterQueue, maxQueueSize,
-        PBRT_CPU_GPU_LAMBDA(MediumScatterWorkItem w) {
+    std::string desc = std::string("Sample direct/indirect - ") + Phase::Name();
+    ForAllQueued(desc.c_str(),
+                 mediumScatterQueue->Get<MediumScatterWorkItem<Phase>>(),
+                 maxQueueSize,
+        PBRT_CPU_GPU_LAMBDA(const MediumScatterWorkItem<Phase> w) {
             RaySamples raySamples = pixelSampleState.samples[w.pixelIndex];
-            Float time = 0;  // TODO: FIXME
             Vector3f wo = w.wo;
 
             // Sample direct lighting at medium scattering event.  First,
@@ -260,7 +278,7 @@ void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
                     ctx, raySamples.direct.u, w.lambda, LightSamplingMode::WithMIS);
                 if (ls && ls->L && ls->pdf > 0) {
                     Vector3f wi = ls->wi;
-                    SampledSpectrum T_hat = w.T_hat * w.phase.p(wo, wi);
+                    SampledSpectrum T_hat = w.T_hat * w.phase->p(wo, wi);
 
                     PBRT_DBG("Phase phase T_hat %f %f %f %f\n", T_hat[0], T_hat[1],
                              T_hat[2], T_hat[3]);
@@ -268,12 +286,12 @@ void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
                     // Compute PDFs for direct lighting MIS calculation.
                     Float lightPDF = ls->pdf * sampledLight->pdf;
                     Float phasePDF =
-                        IsDeltaLight(light.Type()) ? 0.f : w.phase.PDF(wo, wi);
+                        IsDeltaLight(light.Type()) ? 0.f : w.phase->PDF(wo, wi);
                     SampledSpectrum uniPathPDF = w.uniPathPDF * phasePDF;
                     SampledSpectrum lightPathPDF = w.uniPathPDF * lightPDF;
 
                     SampledSpectrum Ld = T_hat * ls->L;
-                    Ray ray(w.p, ls->pLight.p() - w.p, time, w.medium);
+                    Ray ray(w.p, ls->pLight.p() - w.p, w.time, w.medium);
 
                     // Enqueue shadow ray
                     shadowRayQueue->Push(ray, 1 - ShadowEpsilon, w.lambda, Ld, uniPathPDF,
@@ -291,7 +309,7 @@ void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
 
             // Sample indirect lighting.
             pstd::optional<PhaseFunctionSample> phaseSample =
-                w.phase.Sample_p(wo, raySamples.indirect.u);
+                w.phase->Sample_p(wo, raySamples.indirect.u);
             if (!phaseSample || phaseSample->pdf == 0)
                 return;
 
@@ -315,7 +333,7 @@ void WavefrontPathIntegrator::SampleMediumInteraction(int depth) {
                 lightPathPDF *= 1 - q;
             }
 
-            Ray ray(w.p, phaseSample->wi, time, w.medium);
+            Ray ray(w.p, phaseSample->wi, w.time, w.medium);
             bool isSpecularBounce = false;
             bool anyNonSpecularBounces = true;
 

src/pbrt/wavefront/workitems.h

@@ -258,12 +258,14 @@ struct MediumSampleWorkItem {
 };
 
 // MediumScatterWorkItem Definition
+template <typename PhaseFunction>
 struct MediumScatterWorkItem {
     Point3f p;
     SampledWavelengths lambda;
     SampledSpectrum T_hat, uniPathPDF;
-    HGPhaseFunction phase;
+    const PhaseFunction *phase;
     Vector3f wo;
+    Float time;
     Float etaScale;
     Medium medium;
     int pixelIndex;
@@ -469,7 +471,8 @@ class MediumSampleQueue : public WorkQueue<MediumSampleWorkItem> {
 };
 
 // MediumScatterQueue Definition
-using MediumScatterQueue = WorkQueue<MediumScatterWorkItem>;
+using MediumScatterQueue = MultiWorkQueue<
+    typename MapType<MediumScatterWorkItem, typename PhaseFunction::Types>::type>;
 
 // MaterialEvalQueue Definition
 using MaterialEvalQueue = MultiWorkQueue<

src/pbrt/wavefront/workitems.soa

@@ -4,7 +4,7 @@
 // SPDX: Apache-2.0
 
 flat Float;
-flat HGPhaseFunction;
+flat PhaseFunction;
 flat Light;
 flat Material;
 flat Medium;
@@ -137,12 +137,13 @@ soa MediumSampleWorkItem {
     MediumInterface mediumInterface;
 };
 
-soa MediumScatterWorkItem {
+soa MediumScatterWorkItem<PhaseFunction> {
     Point3f p;
     SampledWavelengths lambda;
     SampledSpectrum T_hat, uniPathPDF;
-    HGPhaseFunction phase;
+    const PhaseFunction *phase;
     Vector3f wo;
+    Float time;
     Float etaScale;
     Medium medium;
     int pixelIndex;