Split up sorting keys init (done just once) vs distances update (done…

… every frame)
aras-p · Sep 6, 2023 · 8fa8cca · 8fa8cca
1 parent 7e9a9ba
commit 8fa8cca
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/Assets/Scripts/GaussianSplatRenderer.cs b/Assets/Scripts/GaussianSplatRenderer.cs
@@ -198,6 +198,12 @@ public void OnEnable()
         m_GpuSortDistances = new GraphicsBuffer(GraphicsBuffer.Target.Structured, splatCountNextPot, 4);
         m_GpuSortKeys = new GraphicsBuffer(GraphicsBuffer.Target.Structured, splatCountNextPot, 4);
 
+        // init keys buffer to splat indices
+        m_CSSplatUtilities.SetBuffer(0, "_SplatSortKeys", m_GpuSortKeys);
+        m_CSSplatUtilities.SetInt("_SplatCountPOT", m_GpuSortDistances.count);
+        m_CSSplatUtilities.GetKernelThreadGroupSizes(0, out uint gsX, out uint gsY, out uint gsZ);
+        m_CSSplatUtilities.Dispatch(0, (m_GpuSortDistances.count + (int)gsX - 1)/(int)gsX, 1, 1);
+
         m_Material.SetBuffer("_DataBuffer", m_GpuData);
         m_Material.SetBuffer("_OrderBuffer", m_GpuSortKeys);
 
@@ -229,14 +235,14 @@ public void OnDisable()
     void SortPoints(Camera cam)
     {
         // calculate distance to the camera for each splat
-        m_CSSplatUtilities.SetBuffer(0, "_InputPositions", m_GpuPositions);
-        m_CSSplatUtilities.SetBuffer(0, "_SplatSortDistances", m_GpuSortDistances);
-        m_CSSplatUtilities.SetBuffer(0, "_SplatSortKeys", m_GpuSortKeys);
+        m_CSSplatUtilities.SetBuffer(1, "_InputPositions", m_GpuPositions);
+        m_CSSplatUtilities.SetBuffer(1, "_SplatSortDistances", m_GpuSortDistances);
+        m_CSSplatUtilities.SetBuffer(1, "_SplatSortKeys", m_GpuSortKeys);
         m_CSSplatUtilities.SetMatrix("_WorldToCameraMatrix", cam.worldToCameraMatrix);
         m_CSSplatUtilities.SetInt("_SplatCount", m_SplatCount);
         m_CSSplatUtilities.SetInt("_SplatCountPOT", m_GpuSortDistances.count);
-        m_CSSplatUtilities.GetKernelThreadGroupSizes(0, out uint gsX, out uint gsY, out uint gsZ);
-        m_CSSplatUtilities.Dispatch(0, (m_GpuSortDistances.count + (int)gsX - 1)/(int)gsX, 1, 1);
+        m_CSSplatUtilities.GetKernelThreadGroupSizes(1, out uint gsX, out uint gsY, out uint gsZ);
+        m_CSSplatUtilities.Dispatch(1, (m_GpuSortDistances.count + (int)gsX - 1)/(int)gsX, 1, 1);
 
         // sort the splats
         CommandBuffer cmd = new CommandBuffer {name = "GPUSort"};

diff --git a/Assets/Shaders/SplatUtilities.compute b/Assets/Shaders/SplatUtilities.compute
@@ -1,5 +1,6 @@
 #define GROUP_SIZE 1024
 
+#pragma kernel CSSetIndices
 #pragma kernel CSCalcDistances
 
 StructuredBuffer<float3> _InputPositions;
@@ -10,21 +11,33 @@ uint _SplatCount;
 uint _SplatCountPOT;
 
 [numthreads(GROUP_SIZE,1,1)]
-void CSCalcDistances (uint3 id : SV_DispatchThreadID)
+void CSSetIndices (uint3 id : SV_DispatchThreadID)
 {
     uint idx = id.x;
     if (idx >= _SplatCountPOT)
         return;
 
     _SplatSortKeys[idx] = idx;
+}
+
+
+[numthreads(GROUP_SIZE,1,1)]
+void CSCalcDistances (uint3 id : SV_DispatchThreadID)
+{
+    uint idx = id.x;
+    if (idx >= _SplatCountPOT)
+        return;
+
     if (idx >= _SplatCount)
     {
         // fill pad-to-next-POT portion with dummy data that will never get sorted into regular data place
         _SplatSortDistances[idx] = 0;
         return;
     }
 
-    float3 pos = mul(_WorldToCameraMatrix, float4(_InputPositions[idx] * float3(1,1,-1), 1)).xyz;
+    uint origIdx = _SplatSortKeys[idx];
+
+    float3 pos = mul(_WorldToCameraMatrix, float4(_InputPositions[origIdx] * float3(1,1,-1), 1)).xyz;
     // make distance radix sort friendly from http://stereopsis.com/radix.html
     uint fu = asuint(pos.z);
     uint mask = -((int)(fu >> 31)) | 0x80000000;