From e599e455710e15c733bcabdb13f209f1b371c99d Mon Sep 17 00:00:00 2001
From: adunn <adunn@nvidia.com>
Date: Wed, 25 Sep 2024 14:42:16 -0700
Subject: [PATCH] Introduce more safety checks around POM which were relaxed
 during a recent refactoring change.  Suspect these checks are required with
 current code to avoid a crash.  The fix is to check whether POM
 direct/indirect lighting is enabled before taking the POM path - otherwise we
 can end up with bad surface table index while performing POM.

---
 src/dxvk/shaders/rtx/algorithm/integrator_direct.slangh  | 9 ++++++---
 .../shaders/rtx/algorithm/integrator_indirect.slangh     | 6 +++++-
 src/dxvk/shaders/rtx/algorithm/visibility.slangh         | 2 +-
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/src/dxvk/shaders/rtx/algorithm/integrator_direct.slangh b/src/dxvk/shaders/rtx/algorithm/integrator_direct.slangh
index e109d4b4..a38f2309 100644
--- a/src/dxvk/shaders/rtx/algorithm/integrator_direct.slangh
+++ b/src/dxvk/shaders/rtx/algorithm/integrator_direct.slangh
@@ -60,15 +60,18 @@ void evalNEEPrimary(
   uint8_t rayMask = OBJECT_MASK_OPAQUE | (geometryFlags.objectMask & OBJECT_MASK_ALL_DYNAMIC);
   if (cb.enableDirectTranslucentShadows) rayMask |= OBJECT_MASK_TRANSLUCENT;
 
-  const bool enablePOM = cb.pomEnableDirectLighting && opaqueSurfaceMaterialInteractionHasHeightTexture(opaqueSurfaceMaterialInteraction);
-  const uint16_t primarySurfaceIndex = enablePOM ? uint16_t(SharedSurfaceIndex[pixelCoordinate]) : BINDING_INDEX_INVALID;
+  // We can encounter a POM surface during resolving (as it's a shared function) but if direct lighting for POM
+  //  is disabled, we must also handle that user choice gracefully.
+  const bool pomOpaqueSurfaceEncountered = geometryFlags.pomOpaqueSurfaceEncountered && cb.pomEnableDirectLighting;
+  
+  const uint16_t primarySurfaceIndex = pomOpaqueSurfaceEncountered ? uint16_t(SharedSurfaceIndex[pixelCoordinate]) : BINDING_INDEX_INVALID;
 
   const bool isSubsurface = isSubsurfaceMaterial(opaqueSurfaceMaterialInteraction);
 
   const VisibilityResult visibility = traceVisibilityRay<visibilityModeEnablePom | visibilityModeAccurateHitDistance | visibilityModeEnableTranslucentMaterials | visibilityModeEnableSubsurfaceMaterials>(
     minimalSurfaceInteraction,
     lightSample.position, rayMask,
-    sampledTransportPortalIndex, geometryFlags.portalSpace, geometryFlags.pomOpaqueSurfaceEncountered,
+    sampledTransportPortalIndex, geometryFlags.portalSpace, pomOpaqueSurfaceEncountered,
     minimalRayInteraction.coneRadius, viewRay.spreadAngle, 1.0, geometryFlags.isViewModel, false,
     primarySurfaceIndex, SharedTextureCoord[pixelCoordinate],
     isSubsurface, opaqueSurfaceMaterialInteraction.shadingNormal);
diff --git a/src/dxvk/shaders/rtx/algorithm/integrator_indirect.slangh b/src/dxvk/shaders/rtx/algorithm/integrator_indirect.slangh
index 46d95209..aaf99544 100644
--- a/src/dxvk/shaders/rtx/algorithm/integrator_indirect.slangh
+++ b/src/dxvk/shaders/rtx/algorithm/integrator_indirect.slangh
@@ -279,6 +279,10 @@ void integratePathVertex(
     pomOpaqueSurfaceEncountered,
     isRaytracedRenderTarget);
 
+  // We can encounter a POM surface during resolving (as it's a shared function).  But if POM is disabled in indirect cases
+  //  we must also handle that case and ensure nothing tries to use POM elsewhere.
+  pomOpaqueSurfaceEncountered = pomOpaqueSurfaceEncountered && cb.pomEnableIndirectLighting;
+  
   // Output ReSTIR GI geometry information
 
   if (cb.enableReSTIRGI && !pathState.restirGiHasHitGeometry)
@@ -627,7 +631,7 @@ void integratePathVertex(
       {
         const bool isViewModel = pathState.rayMask & OBJECT_MASK_ALL_VIEWMODEL;
         
-        uint16_t surfaceIndex = (cb.pomEnableIndirectLighting && pomOpaqueSurfaceEncountered) ? uint(rayInteraction.surfaceIndex) : BINDING_INDEX_INVALID;
+        uint16_t surfaceIndex = pomOpaqueSurfaceEncountered ? uint(rayInteraction.surfaceIndex) : BINDING_INDEX_INVALID;
     
         // Todo: Ray Portal transport sampling in the future.
         evalNEESecondary(
diff --git a/src/dxvk/shaders/rtx/algorithm/visibility.slangh b/src/dxvk/shaders/rtx/algorithm/visibility.slangh
index 8ba34e29..822dbda1 100644
--- a/src/dxvk/shaders/rtx/algorithm/visibility.slangh
+++ b/src/dxvk/shaders/rtx/algorithm/visibility.slangh
@@ -456,7 +456,7 @@ VisibilityResult traceVisibilityRay<let VisibilityMode : uint>(
   {
     const MemoryPolymorphicSurfaceMaterial memoryPolymorphicSurfaceMaterial = surfaceMaterials[primarySurfaceIndex];
     OpaqueSurfaceMaterial opaqueSurfaceMaterial = opaqueSurfaceMaterialCreate(primarySurfaceIndex, memoryPolymorphicSurfaceMaterial);
-
+    
     pomAttenuation = opaqueSurfaceMaterialInteractionCalcHeightThroughput(
       minimalSurfaceInteraction, visibilityRay.direction, opaqueSurfaceMaterial.heightTextureIndex,
       opaqueSurfaceMaterial.samplerIndex, texCoords, opaqueSurfaceMaterial.displaceIn