Merge branch 'dev/jerrans/nee-cache-packing' into 'main'

Raise primitive id bitcount at cost of reducing range bits. Refactor to make changes consistent.

See merge request lightspeedrtx/dxvk-remix-nv!2288

Author: jerrans

src/dxvk/shaders/rtx/algorithm/nee_cache.h

@@ -226,12 +226,12 @@ struct NEECandidate
 
   int getPrimitiveID()
   {
-    return m_data.y & 0xffffff;
+    return m_data.y & PRIMITIVE_INDEX_MAX_VALUE;
   }
 
   [mutating] void setPrimitiveID(int primitiveID)
   {
-    m_data.y = (m_data.y & 0xff000000) | primitiveID;
+    m_data.y = (m_data.y & ~PRIMITIVE_INDEX_MAX_VALUE) | primitiveID;
   }
 
   uint2 getIDData()
@@ -253,12 +253,12 @@ struct NEECandidate
 
   int getRange()
   {
-    return (m_data.y >> 24) & 0xff;
+    return (m_data.y >> PRIMITIVE_INDEX_BIT_COUNT) & NEE_RANGE_MAX;
   }
 
   [mutating] void setRange(uint range)
   {
-    m_data.y = (m_data.y & 0xffffff) | (range << 24);
+    m_data.y = (m_data.y & PRIMITIVE_INDEX_MAX_VALUE) | (range << PRIMITIVE_INDEX_BIT_COUNT);
   }
 
   static NEECandidate create(uint surfaceID, uint primitiveID, uint range)
@@ -385,16 +385,19 @@ struct NEECell
 
   static bool isLightTask(uint2 value)
   {
-    return (value.x & (1 << 24)) != 0;
+    return (value.x & (1 << NEE_ISLIGHT_BIT)) != 0;
   }
 
   void insertSlotTask(uint task, vec3 radiance, bool isLightTask) {
+    if (!isLightTask && task > PRIMITIVE_INDEX_MAX_VALUE) {
+      return;
+    }
     float accumulateValue = calcBt709Luminance(radiance);
     float randomOffset = (reversebits(asuint(accumulateValue)) >> 22) / 1024.0f;
     uint index = getSlotBinHash(task + cb.frameIdx);
     int taskAddress = getHashTaskAddress(index);
     uint sortValueI = firstbithigh(uint(min(accumulateValue, 50) / 0.001));
-    task |= (sortValueI << 25) | (isLightTask ? (1 << 24) : 0);
+    task |= (sortValueI << NEE_SORT_SHIFT) | (isLightTask ? (1 << NEE_ISLIGHT_BIT) : 0);
 
     // Clamp min/max value before accumulation to improve stability.
     // The min value is required because floating point atomics is not supported on all platforms,

src/dxvk/shaders/rtx/algorithm/nee_cache_data.h

@@ -20,6 +20,7 @@
 * DEALINGS IN THE SOFTWARE.
 */
 #pragma once
+#include "rtx/pass/instance_definitions.h"
 
 #define NEE_CACHE_PROBE_RESOLUTION 32
 #define NEE_CACHE_TOTAL_PROBE (NEE_CACHE_PROBE_RESOLUTION * NEE_CACHE_PROBE_RESOLUTION * NEE_CACHE_PROBE_RESOLUTION)

src/dxvk/shaders/rtx/algorithm/nee_cache_light.slangh

@@ -23,7 +23,8 @@
 
 #include "rtx/concept/light/light_helper.slangh"
 #include "rtx/concept/light/light.slangh"
-#define NEE_CACHE_INVALID_ID 0xffffff
+#include "rtx/pass/instance_definitions.h"
+#define NEE_CACHE_INVALID_ID PRIMITIVE_INDEX_MAX_VALUE
 
 struct NEECacheUtils
 {
@@ -241,7 +242,7 @@ struct NEECacheUtils
     {
       trianglePdf /= float(range);
       float triangleExplorationProbability = cb.neeCacheArgs.triangleExplorationProbability;
-      if (uvw.z < triangleExplorationProbability && range < 255)
+      if (uvw.z < triangleExplorationProbability && range < NEE_RANGE_MAX)
       {
         // Sample a triangle outside the given range with a low probability to search potential emissive triangles.
         // The code here is slightly biased because we use the probability for triangles inside the range,

src/dxvk/shaders/rtx/pass/instance_definitions.h

@@ -138,12 +138,21 @@
 
 // Engine-wide index limits.
 // SurfaceIndex: 21 bits (fits in the 24-bit instanceCustomIndex alongside 2-bit material type + 1-bit view-model flag).
-// PrimitiveIndex: 26 bits (max ~67M triangles per scene).
+// PrimitiveIndex: 27 bits (max 134M triangles per scene). Other packing constants are derived from this value.
 #define SURFACE_INDEX_BIT_COUNT       21
 #define SURFACE_INDEX_MAX_VALUE       ((1 << SURFACE_INDEX_BIT_COUNT) - 1)
-#define PRIMITIVE_INDEX_BIT_COUNT     26
+#define PRIMITIVE_INDEX_BIT_COUNT     27
 #define PRIMITIVE_INDEX_MAX_VALUE     ((1 << PRIMITIVE_INDEX_BIT_COUNT) - 1)
 
+// NEE packing constants
+// The sort key packed above this shift is limited to 32 - PRIMITIVE_INDEX_BIT_COUNT - 1 (for NEE_ISLIGHT_BIT)
+// Current max value in nee_cache.h is: firstbithigh(50 / 0.001); which requires 4 bits so with a PRIMITIVE_INDEX_BIT_COUNT
+// of 27 there is no capacity to increase the PRIMITIVE_INDEX_BIT_COUNT, only to reduce it (or change the structure)
+#define NEE_ISLIGHT_BIT               PRIMITIVE_INDEX_BIT_COUNT
+#define NEE_SORT_SHIFT                (NEE_ISLIGHT_BIT + 1)
+#define NEE_RANGE_BITS                (32 - PRIMITIVE_INDEX_BIT_COUNT)
+#define NEE_RANGE_MAX                 ((1 << NEE_RANGE_BITS) - 1)
+
 // Custom Index encoding (24-bit VkAccelerationStructureInstanceKHR.instanceCustomIndex)
 //   Bits  0..20 : surface index  (CUSTOM_INDEX_SURFACE_MASK)
 //   Bits 21..22 : material type  (CUSTOM_INDEX_MATERIAL_TYPE_MASK)

src/dxvk/shaders/rtx/pass/nee_cache/update_nee_cache.comp.slang

@@ -43,8 +43,6 @@
 #define UPDATE_NEE_CACHE 1
 #include "rtx/algorithm/nee_cache.h"
 
-#define INVALID_ID 0xffffff
-
 #define CANDIDATE_LENGTH NEE_CACHE_ELEMENTS * 2
 #define CANDIDATE_GROUP_SIZE 16
 #define BATCH_SIZE 8
@@ -226,8 +224,8 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
     int taskID = localIndex.x;
     uint2 value = cell.getSlotTaskValue(taskID);
     int surfaceID, primitiveID;
-    uint oldRange = (value.x & 0xff000000);
-    value.x = convertToThisFramePrefixSumID(value.x & 0xffffff, surfaceID, primitiveID);
+    uint oldRange = (value.x & ~PRIMITIVE_INDEX_MAX_VALUE);
+    value.x = convertToThisFramePrefixSumID(value.x & PRIMITIVE_INDEX_MAX_VALUE, surfaceID, primitiveID);
     value.x |= oldRange;
     int delta = max(value.y * cb.neeCacheArgs.learningRate, 1);
     value.y = clamp(int(value.y) - delta, 0, 1 << 25);
@@ -264,7 +262,7 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
     {
       continue;
     }
-    value.x &= 0xffffff;
+    value.x &= PRIMITIVE_INDEX_MAX_VALUE;
     int surfaceID, primitiveID;
     value.x = convertToThisFramePrefixSumID(value.x, surfaceID, primitiveID);
     value.y = min(maxDeltaValue, value.y);
@@ -288,8 +286,8 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
       bool found = false;
       for (int j = 0; j < existingCount; ++j)
       {
-        int firstPrimitiveID = (s_candidateList[j][localIndex.y].x & 0xffffff);
-        int range = (s_candidateList[j][localIndex.y].x >> 24);
+        int firstPrimitiveID = (s_candidateList[j][localIndex.y].x & PRIMITIVE_INDEX_MAX_VALUE);
+        int range = (s_candidateList[j][localIndex.y].x >> PRIMITIVE_INDEX_BIT_COUNT);
         if (value.x >= firstPrimitiveID && value.x < firstPrimitiveID + range)
         {
           s_candidateList[j][localIndex.y].y += value.y;
@@ -300,7 +298,7 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
 
       if (!found)
       {
-        value.x = ((1 << 24) | value.x);
+        value.x = ((1 << PRIMITIVE_INDEX_BIT_COUNT) | value.x);
         insertTask(value, localIndex.y);
       }
     }
@@ -314,9 +312,9 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
   for (int i = localIndex.x; i < NEE_CACHE_HASH_TASK_COUNT && cb.neeCacheArgs.triangleExplorationProbability > 0; i += CANDIDATE_GROUP_SIZE)
   {
     uint2 data = s_candidateList[i][localIndex.y];
-    uint range = (data.x >> 24);
-    uint primitiveID = (data.x & 0xffffff);
-    uint2 newData = uint2(data.y, (primitiveID << 8) | range);
+    uint range = (data.x >> PRIMITIVE_INDEX_BIT_COUNT);
+    uint primitiveID = (data.x & PRIMITIVE_INDEX_MAX_VALUE);
+    uint2 newData = uint2(data.y, (primitiveID << NEE_RANGE_BITS) | range);
     s_candidateList[i][localIndex.y] = newData;
   }
 
@@ -335,8 +333,8 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
     for (int i = 0; i < count; i++)
     {
       uint2 data = s_candidateList[i][localIndex.y];
-      uint newRange = (data.y & 0xff);
-      uint newPrefixSumID = (data.y >> 8);
+      uint newRange = (data.y & NEE_RANGE_MAX);
+      uint newPrefixSumID = (data.y >> NEE_RANGE_BITS);
       uint newValue = data.x;
       totalValue += newValue;
     }
@@ -350,24 +348,24 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
     for (int i = 0; i < count; i++)
     {
       uint2 data = s_candidateList[i][localIndex.y];
-      uint newRange = (data.y & 0xff);
-      uint newPrefixSumID = (data.y >> 8);
+      uint newRange = (data.y & NEE_RANGE_MAX);
+      uint newPrefixSumID = (data.y >> NEE_RANGE_BITS);
       uint newValue = data.x;
 
       int newSurfaceID, newPrimitiveID;
       NEECacheUtils.convertPrefixSumIDToID(newPrefixSumID, PrimitiveIDPrefixSum, newSurfaceID, newPrimitiveID);
       int primitiveCount = NEECacheUtils.getSurfacePrimitiveCount(newSurfaceID, PrimitiveIDPrefixSum);
       int gab = min(cb.neeCacheArgs.triangleExplorationMaxRange, cb.neeCacheArgs.triangleExplorationRangeRatio * primitiveCount) * 
         cb.neeCacheArgs.triangleExplorationAcceptRangeRatio + 1;
-      if (newSurfaceID == oldSurfaceID && oldPrimitiveID < newPrimitiveID + newRange + gab && oldPrimitiveID + oldRange - newPrimitiveID < 255 &&
+      if (newSurfaceID == oldSurfaceID && oldPrimitiveID < newPrimitiveID + newRange + gab && oldPrimitiveID + oldRange - newPrimitiveID < NEE_RANGE_MAX &&
         oldValue + newValue < maxMergeValue)
       {
         // Merge
         oldRange = oldPrimitiveID + oldRange - newPrimitiveID;
         oldPrefixSumID = newPrefixSumID;
         oldPrimitiveID = newPrimitiveID;
         oldValue += newValue;
-        s_candidateList[oldIndex][localIndex.y] = uint2((oldRange << 24) | oldPrefixSumID, oldValue);
+        s_candidateList[oldIndex][localIndex.y] = uint2((oldRange << PRIMITIVE_INDEX_BIT_COUNT) | oldPrefixSumID, oldValue);
         newValue = 0;
         newRange = 0;
       }
@@ -382,7 +380,7 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
         oldValue = newValue;
       }
 
-      s_candidateList[i][localIndex.y] = uint2((newRange << 24) | newPrefixSumID, newValue);
+      s_candidateList[i][localIndex.y] = uint2((newRange << PRIMITIVE_INDEX_BIT_COUNT) | newPrefixSumID, newValue);
     }
   }
 
@@ -400,7 +398,7 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
   if (cb.neeCacheArgs.triangleExplorationProbability == 0.0)
   {
     uint2 task = s_candidateList[localIndex.x][localIndex.y];
-    task.x = (task.x & 0xffffff) | (1 << 24);
+    task.x = (task.x & PRIMITIVE_INDEX_MAX_VALUE) | (1 << PRIMITIVE_INDEX_BIT_COUNT);
 
     s_candidateList[localIndex.x][localIndex.y].x = task.x;
     GroupMemoryBarrierWithGroupSync();
@@ -421,8 +419,8 @@ void updateTriangleTask(NEECell cell, uint2 localIndex)
     int candidateID = localIndex.x;
     uint2 task = s_candidateList[candidateID][localIndex.y];
     int surfaceID, primitiveID;
-    NEECacheUtils.convertPrefixSumIDToID(task.x & 0xffffff, PrimitiveIDPrefixSum, surfaceID, primitiveID);
-    NEECandidate candidate = NEECandidate.create(surfaceID, primitiveID, max(1, task.x >> 24));
+    NEECacheUtils.convertPrefixSumIDToID(task.x & PRIMITIVE_INDEX_MAX_VALUE, PrimitiveIDPrefixSum, surfaceID, primitiveID);
+    NEECandidate candidate = NEECandidate.create(surfaceID, primitiveID, max(1, task.x >> PRIMITIVE_INDEX_BIT_COUNT));
 
     candidate.setSampleProbability(s_candidateLight[candidateID][localIndex.y]);
     cell.setCandidate(candidateID, candidate);
@@ -539,7 +537,7 @@ void updateLightTask(NEECell cell, uint2 localIndex)
     {
       continue;
     }
-    value.x &= 0xffffff;
+    value.x &= PRIMITIVE_INDEX_MAX_VALUE;
     value.x = convertLightIndex(value.x);
     value.y = min(maxDeltaValue, value.y);