Merge branch 'yaobin/prefixSum' into 'main'

Use new prefix sum buffer for NEE Cache to eliminate branch code. See merge request lightspeedrtx/dxvk-remix-nv!800
NVIDIAGameWorks · Apr 16, 2024 · 07f6d27 · 07f6d27
2 parents 40552ea + 9eaa70b
commit 07f6d27
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 18 deletions.
diff --git a/src/dxvk/rtx_render/rtx_accel_manager.cpp b/src/dxvk/rtx_render/rtx_accel_manager.cpp
@@ -602,18 +602,20 @@ namespace dxvk {
 
     // Build prefix sum array
     // Collect primitive count for each surface object
+    // Because we use exclusive prefix sum here, we add one more element to record the scene's total primitive count
     m_reorderedSurfacesPrimitiveIDPrefixSumLastFrame = m_reorderedSurfacesPrimitiveIDPrefixSum;
-    m_reorderedSurfacesPrimitiveIDPrefixSum.resize(m_reorderedSurfaces.size());
+    m_reorderedSurfacesPrimitiveIDPrefixSum.resize(m_reorderedSurfaces.size() + 1);
+    m_reorderedSurfacesPrimitiveIDPrefixSum[0] = 0;
     for (uint32_t i = 0; i < m_reorderedSurfaces.size(); i++) {
       auto surface = m_reorderedSurfaces[i];
       int primitiveCount = 0;
       for (const auto& buildRange: surface->buildRanges) {
         primitiveCount += buildRange.primitiveCount;
       }
-      m_reorderedSurfacesPrimitiveIDPrefixSum[i] = primitiveCount;
+      m_reorderedSurfacesPrimitiveIDPrefixSum[i + 1] = primitiveCount;
     }
 
-    // Calculate prefix sum
+    // Calculate exclusive prefix sum
     uint totalPrimitiveIDOffset = 0;
     for (uint32_t i = 0; i < m_reorderedSurfacesPrimitiveIDPrefixSum.size(); i++) {
       uint primitiveCount = m_reorderedSurfacesPrimitiveIDPrefixSum[i];

diff --git a/src/dxvk/rtx_render/rtx_accel_manager.h b/src/dxvk/rtx_render/rtx_accel_manager.h
@@ -128,8 +128,8 @@ class AccelManager : public CommonDeviceObject {
   void internalBuildTlas(Rc<DxvkContext> ctx);
   std::vector<RtInstance*> m_reorderedSurfaces;
   std::vector<uint32_t> m_reorderedSurfacesFirstIndexOffset;
-  std::vector<uint32_t> m_reorderedSurfacesPrimitiveIDPrefixSum;
-  std::vector<uint32_t> m_reorderedSurfacesPrimitiveIDPrefixSumLastFrame;
+  std::vector<uint32_t> m_reorderedSurfacesPrimitiveIDPrefixSum;              // Exclusive prefix sum for this frame's surface primitive count array
+  std::vector<uint32_t> m_reorderedSurfacesPrimitiveIDPrefixSumLastFrame;     // Exclusive prefix sum for last frame's surface primitive count array
   std::vector<VkAccelerationStructureInstanceKHR> m_mergedInstances[Tlas::Count];
   std::vector<Rc<PooledBlas>> m_blasPool;
 

diff --git a/src/dxvk/shaders/rtx/algorithm/nee_cache_light.slangh b/src/dxvk/shaders/rtx/algorithm/nee_cache_light.slangh
@@ -53,15 +53,7 @@ struct NEECacheUtils
     {
       return NEE_CACHE_INVALID_ID;
     }
-    int prefixSum;
-    if (surfaceIndex == 0)
-    {
-      prefixSum = 0;
-    }
-    else
-    {
-      prefixSum = PrimitiveIDPrefixSumBuffer[surfaceIndex-1];
-    }
+    int prefixSum = PrimitiveIDPrefixSumBuffer[surfaceIndex];
     return primitiveIndex + prefixSum;
   }
 
@@ -78,19 +70,19 @@ struct NEECacheUtils
     int maxSurfaceID = cb.surfaceCount - 1;
     while (minSurfaceID < maxSurfaceID)
     {
-      int midSurfaceID = (minSurfaceID + maxSurfaceID) / 2;
+      int midSurfaceID = (minSurfaceID + maxSurfaceID) / 2 + 1;
       if (PrimitiveIDPrefixSumBuffer[midSurfaceID] <= prefixSumID)
       {
-        minSurfaceID = midSurfaceID + 1;
+        minSurfaceID = midSurfaceID;
       }
       else
       {
-        maxSurfaceID = midSurfaceID;
+        maxSurfaceID = midSurfaceID - 1;
       }
     }
 
     surfaceID = minSurfaceID;
-    int surfacePrefixSum = surfaceID == 0 ? 0 : PrimitiveIDPrefixSumBuffer[surfaceID-1];
+    int surfacePrefixSum = PrimitiveIDPrefixSumBuffer[surfaceID];
     primitiveID = prefixSumID - surfacePrefixSum;
     return true;
   }