rtx: environment-map MIS (balance heuristic)

devices/rtx/device/gpu/evalShading.h

@@ -119,7 +119,7 @@ VISRTX_DEVICE NextRay materialNextRay(
 {
   if (shadingState.callableBaseIndex
       == ~DeviceObjectIndex(0)) // No next ray by default
-    return NextRay{vec3(0.0f), vec3(0.0f)};
+    return NextRay{vec3(0.0f), vec3(0.0f), 0.0f};
 
   return optixDirectCall<NextRay>(shadingState.callableBaseIndex
           + int(SurfaceShaderEntryPoints::EvaluateNextRay),
@@ -128,6 +128,24 @@ VISRTX_DEVICE NextRay materialNextRay(
       &rs);
 }
 
+// Solid-angle pdf the material's BSDF sampler would assign to direction `wi`
+// given outgoing direction `wo` (both world space), for environment MIS. 0 for
+// no material and for materials whose sampler cannot produce `wi` (e.g. Matte,
+// which has no continuation ray).
+VISRTX_DEVICE float materialEvalPdf(const MaterialShadingState &shadingState,
+    const vec3 &wo,
+    const vec3 &wi)
+{
+  if (shadingState.callableBaseIndex == ~DeviceObjectIndex(0))
+    return 0.0f;
+
+  return optixDirectCall<float>(shadingState.callableBaseIndex
+          + int(SurfaceShaderEntryPoints::EvaluatePdf),
+      &shadingState.data,
+      &wo,
+      &wi);
+}
+
 VISRTX_DEVICE vec3 materialShadeSurface(
     const MaterialShadingState &shadingState,
     const SurfaceHit &hit,

devices/rtx/device/gpu/gpu_util.h

@@ -403,14 +403,41 @@ VISRTX_DEVICE bool getBackgroundLight(
       // For orthonormal matrices, inverse = transpose
       const mat3 xfmInv = glm::transpose(mat3(hdriLight.xfm));
       const vec3 localRayDir = xfmInv * rayDir;
-      outRadiance += sampleHDRI(light, localRayDir);
+      // sampleHDRI applies hdri.scale; tint by light.color to match the NEE
+      // radiance in sampleHDRILight (raw * hdri.scale * color), so env MIS
+      // deposits identical radiance on the NEE and BSDF-escape sides.
+      outRadiance += sampleHDRI(light, localRayDir) * light.color;
       hasVisibleHDRI = true;
     }
   }
 
   return hasVisibleHDRI;
 }
 
+// Solid-angle sampling pdf of the visible HDRI environment(s) at `rayDir`, used
+// as the light-sampling density on the escape side of environment MIS. It must
+// match the NEE importance pdf in sampleHDRILight exactly: raw-texel luminance
+// (NO scale/color) times pdfWeight, with the same instance/HDRI transform chain
+// as getBackgroundLight. Summed over visible HDRIs — exact for the single-HDRI
+// case, a mixture-pdf approximation when several are visible.
+VISRTX_DEVICE float envPdf(const FrameGPUData &fd, const vec3 &rayDir)
+{
+  float pdf = 0.0f;
+  for (size_t i = 0; i < fd.world.numHdriLightInstances; i++) {
+    const auto &hdriLight = fd.world.hdriLightInstances[i];
+    const auto &light = fd.registry.lights[hdriLight.lightIndex];
+    if (!light.hdri.visible)
+      continue;
+    const vec3 localRayDir = glm::transpose(mat3(hdriLight.xfm)) * rayDir;
+    const vec3 d = light.hdri.xfm * localRayDir;
+    const vec2 thetaPhi = sphericalCoordsFromDirection(d);
+    const vec2 uv = vec2(thetaPhi.y / kTwoPi, thetaPhi.x / kPi);
+    pdf += dot(sampleHDRI(light, uv), vec3(0.2126f, 0.7152f, 0.0722f))
+        * light.hdri.pdfWeight;
+  }
+  return pdf;
+}
+
 VISRTX_DEVICE uint32_t computeGeometryPrimId(const SurfaceHit &hit)
 {
   if (!hit.foundHit)

devices/rtx/device/gpu/sbt.h

@@ -47,6 +47,7 @@ enum class SurfaceShaderEntryPoints
   EvaluateTransmission,
   EvaluateNormal,
   Shade,
+  EvaluatePdf,
   Count
 };
 

devices/rtx/device/gpu/shadingState.h

@@ -61,6 +61,12 @@ struct NextRay
 {
   vec3 direction;
   vec3 contributionWeight;
+  // Solid-angle pdf of `direction`, used for balance-heuristic environment MIS.
+  // +inf marks a lobe whose env contribution the escape estimator owns outright
+  // (primary ray, transmission); 0 marks a dead ray. Must equal the value the
+  // material's EvaluatePdf callable returns for the same direction on the
+  // reflection side, so NEE-side and escape-side MIS weights partition to 1.
+  float pdf{INFINITY};
   uint32_t flags{NEXT_RAY_NONE};
 };
 

devices/rtx/device/material/shaders/MDLShader_ptx.cu

@@ -220,10 +220,23 @@ NextRay __direct_callable__nextRay(
       ? NEXT_RAY_CONTINUES_THROUGH_SURFACE
       : NEXT_RAY_NONE;
 
+  // Env-MIS solid-angle pdf for the sampled direction, matching
+  // __direct_callable__evaluatePdf so the balance heuristic partitions to 1.
+  // A specular (delta) lobe can't be evaluated by NEE, and a through-surface
+  // continuation is past the NEE hemisphere gate — both report +inf so the BSDF
+  // escape owns the environment (w_bsdf = 1). Glossy/diffuse reflections report
+  // the finite sampling pdf and are MIS-combined with NEE.
+  const bool isSpecular =
+      (sample_data.event_type & mi::neuraylib::BSDF_EVENT_SPECULAR) != 0;
+  const float pdf =
+      (isSpecular || (flags & NEXT_RAY_CONTINUES_THROUGH_SURFACE))
+      ? INFINITY
+      : sample_data.pdf;
   return NextRay{direction,
       vec3(sample_data.bsdf_over_pdf.x,
           sample_data.bsdf_over_pdf.y,
           sample_data.bsdf_over_pdf.z),
+      pdf,
       flags};
 }
 
@@ -275,3 +288,31 @@ vec3 __direct_callable__evaluateNormal(const MDLShadingState *shadingState)
 {
   return make_vec3(shadingState->state.normal);
 }
+
+// Env-MIS BSDF density at `wi` given outgoing `wo` (both world space): the
+// balance-heuristic light-side weight. mdlBsdf_evaluate fills `eval_data.pdf`
+// with the solid-angle sampling pdf, matching NextRay.pdf in nextRay (MDL's
+// evaluate-pdf and sample-pdf are the same density). A pure specular lobe
+// evaluates to pdf 0 (NEE can't reach a delta) — consistent with the escape
+// owning it via +inf. Mirrors shadeSurface's ior/k1/k2 setup exactly.
+VISRTX_CALLABLE float __direct_callable__evaluatePdf(
+    const MDLShadingState *shadingState, const vec3 *wo, const vec3 *wi)
+{
+  BsdfEvaluateData eval_data = {};
+  if (shadingState->isFrontFace) {
+    eval_data.ior1 = make_float3(1.0f, 1.0f, 1.0f);
+    eval_data.ior2.x = MI_NEURAYLIB_BSDF_USE_MATERIAL_IOR;
+  } else {
+    eval_data.ior1.x = MI_NEURAYLIB_BSDF_USE_MATERIAL_IOR;
+    eval_data.ior2 = make_float3(1.0f, 1.0f, 1.0f);
+  }
+  eval_data.k1 = make_float3(normalize(*wo));
+  eval_data.k2 = make_float3(normalize(*wi));
+
+  mdlBsdf_evaluate(&eval_data,
+      &shadingState->state,
+      &shadingState->resData,
+      shadingState->argBlock);
+
+  return eval_data.pdf;
+}

devices/rtx/device/material/shaders/MatteShader_ptx.cu

@@ -59,7 +59,7 @@ VISRTX_CALLABLE void __direct_callable__init(MatteShadingState *shadingState,
 VISRTX_CALLABLE NextRay __direct_callable__nextRay(
     const MatteShadingState *, const Ray *, RandState *)
 {
-  return NextRay{vec3(0.0f), vec3(0.0f)};
+  return NextRay{vec3(0.0f), vec3(0.0f), 0.0f};
 }
 
 VISRTX_CALLABLE
@@ -105,3 +105,12 @@ VISRTX_CALLABLE vec3 __direct_callable__shadeSurface(
   return shadingState->baseColor * kInvPi * NdotL * lightSample->radiance
       / lightSample->pdf;
 }
+
+// Matte has no continuation ray (nextRay returns a dead ray), so its BSDF can
+// never produce the environment direction: report pdf 0 so env MIS leaves NEE
+// owning the environment (w_nee = 1), matching the pre-MIS behavior.
+VISRTX_CALLABLE float __direct_callable__evaluatePdf(
+    const MatteShadingState *, const vec3 *, const vec3 *)
+{
+  return 0.0f;
+}