Partition Gaussian chunks by a uniform grid before Morton-slicing them

Sources/UntoldEngine/AssetFormat/UntoldGSWriter.swift

@@ -203,7 +203,23 @@ public extension UntoldGSFormat {
         }
 
         let splatsPerChunk = 1 << Int(options.log2ChunkSplats)
-        let chunkCount = (splatCount + splatsPerChunk - 1) / splatsPerChunk
+
+        // Progress: with coarse levels the `chunk` phase is the ordering and the layout and the
+        // chunk loop reports as `coarsen`; without them the loop is the rest of `chunk`, so the
+        // ordering takes the first tenth and the fraction never runs backwards.
+        try progress?.report(.chunk, fraction: 0)
+        let bounds = bounds(of: view)
+        let (order, chunkSizes) = gridChunkPlan(view, boundsMin: bounds.min, boundsMax: bounds.max, splatsPerChunk: splatsPerChunk)
+        let chunkCount = chunkSizes.count
+        let chunkStarts: [Int] = {
+            var starts = [Int](repeating: 0, count: chunkCount)
+            var cursor = 0
+            for (index, size) in chunkSizes.enumerated() {
+                starts[index] = cursor
+                cursor += size
+            }
+            return starts
+        }()
 
         // The coarse levels: automatic above the chunk-count threshold (the template's ratios
         // clamped to the chunk size), or exactly what was asked for.
@@ -222,14 +238,8 @@ public extension UntoldGSFormat {
             return requested
         }()
 
-        // Progress: with coarse levels the `chunk` phase is the ordering and the layout and the
-        // chunk loop reports as `coarsen`; without them the loop is the rest of `chunk`, so the
-        // ordering takes the first tenth and the fraction never runs backwards.
         let orderingShare = coarseOptions == nil ? 0.1 : 1.0
         progress?.setTierHasCoarseLevels(coarseOptions != nil)
-        try progress?.report(.chunk, fraction: 0)
-        let bounds = bounds(of: view)
-        let order = mortonOrder(view, boundsMin: bounds.min, boundsMax: bounds.max)
         try progress?.report(.chunk, fraction: 0.5 * orderingShare)
 
         // The layout is fixed before a chunk is encoded: the tree's node count depends on the
@@ -245,8 +255,7 @@ public extension UntoldGSFormat {
         var cursor = payloadOffset
         let payloadOffsets: [Int] = (0 ..< chunkCount).map { chunk in
             let offset = cursor
-            let count = min(splatsPerChunk, splatCount - chunk * splatsPerChunk)
-            cursor += alignedToPage(count * (coreRecordSize + shCount))
+            cursor += alignedToPage(chunkSizes[chunk] * (coreRecordSize + shCount))
             return offset
         }
         var fileSize = cursor
@@ -268,8 +277,8 @@ public extension UntoldGSFormat {
         let results = ChunkResults(count: chunkCount)
         let work: @Sendable (Int) -> Void = { chunk in
             do {
-                let start = chunk * splatsPerChunk
-                let end = min(start + splatsPerChunk, splatCount)
+                let start = chunkStarts[chunk]
+                let end = start + chunkSizes[chunk]
                 let mortonOrdered = order[start ..< end].map { Int($0) }
                 var ordered = mortonOrdered
                 if options.sortByImportanceWithinChunk {
@@ -541,6 +550,56 @@ public extension UntoldGSFormat {
         return sortedByKeyThenIndex(keys)
     }
 
+    /// The tier's Morton order, cut into chunks that never cross a uniform grid cell — instead
+    /// of a chunk being "the next `splatsPerChunk` splats in scene-wide Morton order" (which can
+    /// span the whole scene when density is uneven), it's "the next `splatsPerChunk` splats in
+    /// Morton order *and* in the same grid cell". A 63-bit Morton key interleaves 21 bits per
+    /// axis with the coarsest bits on top, so a uniform grid at `2^bitsPerAxis` cells per axis is
+    /// exactly the key's top `3 * bitsPerAxis` bits — the scene-wide order is already grouped by
+    /// cell, ascending, with no separate cell sort. `bitsPerAxis` is chosen so the grid has about
+    /// as many cells as the tier would have made chunks the old way, so a uniformly dense region
+    /// still yields chunks close to `splatsPerChunk`; a sparse or empty cell just yields a
+    /// smaller trailing chunk instead of merging into whatever cell comes next in Morton order.
+    /// Returns the permutation (identical to `mortonOrder(_:boundsMin:boundsMax:)`) and each
+    /// resulting chunk's own splat count, since a cell's last chunk — or a whole sparse cell —
+    /// can be smaller than `splatsPerChunk`.
+    internal static func gridChunkPlan(
+        _ view: UntoldGSStoreView, boundsMin: SIMD3<Float>, boundsMax: SIMD3<Float>, splatsPerChunk: Int
+    ) -> (order: [UInt32], chunkSizes: [Int]) {
+        let count = view.count
+        guard count > 0 else { return ([], []) }
+        let order = mortonOrder(view, boundsMin: boundsMin, boundsMax: boundsMax)
+        guard splatsPerChunk > 0 else { return (order, [count]) }
+
+        let targetChunks = max(1, (count + splatsPerChunk - 1) / splatsPerChunk)
+        let bitsPerAxis = min(21, max(0, Int((log2(Double(targetChunks)) / 3).rounded(.down))))
+        let cellShift = UInt64(3 * (21 - bitsPerAxis))
+
+        var chunkSizes: [Int] = []
+        chunkSizes.reserveCapacity(targetChunks + targetChunks / 4)
+        view.withUnsafePointers { pointers in
+            var currentCell: UInt64 = .max
+            var currentSize = 0
+            for splatIndex in order {
+                let position = pointers.position(pointers.storeIndex(Int(splatIndex)))
+                let cell = UntoldGSPacking.mortonKey(position, boundsMin: boundsMin, boundsMax: boundsMax) >> cellShift
+                if currentSize > 0, cell == currentCell, currentSize < splatsPerChunk {
+                    currentSize += 1
+                } else {
+                    if currentSize > 0 {
+                        chunkSizes.append(currentSize)
+                    }
+                    currentCell = cell
+                    currentSize = 1
+                }
+            }
+            if currentSize > 0 {
+                chunkSizes.append(currentSize)
+            }
+        }
+        return (order, chunkSizes)
+    }
+
     /// Indices `0 ..< keys.count` ordered by `(key, index)`: a least-significant-digit radix sort
     /// in 16-bit digits, stable, so equal keys keep ascending indices.
     private static func sortedByKeyThenIndex(_ keys: UnsafeSharedBuffer<UInt64>) -> [UInt32] {

Sources/UntoldEngine/RuntimeAssets/GaussianChunkLoader.swift

@@ -522,6 +522,9 @@ enum GaussianChunkLoader {
         let pagesPerChunk = GaussianPagingPolicy.pagesPerChunk(splatsPerChunk: header.splatsPerChunk)
         let slotBytes = ranksPerPage * (UntoldGSFormat.coreRecordSize + shBytesPerSplat)
         let assetBytes = GaussianPagingPolicy.assetBytes(splatCount: splatCount, shBytesPerSplat: shBytesPerSplat)
+        let assetSlotCount = index.chunks.reduce(into: 0) { count, chunk in
+            count += (Int(chunk.splatCount) + ranksPerPage - 1) / ranksPerPage
+        }
         // Sized and claimed in one step under the registry's lock: two loads running at once
         // (tiers of one progressive entity, streamed entities) each see the other's claim, so
         // the pools together stay within the residency budget. The claim becomes the pager's
@@ -531,6 +534,7 @@ enum GaussianChunkLoader {
         let reservation = registry.reserve { allocatedBytes in
             slotCount = GaussianPagingPolicy.poolSlotCount(
                 assetBytes: assetBytes,
+                assetSlotCount: assetSlotCount,
                 slotBytes: slotBytes,
                 residencyBudgetBytes: residencyBudgetBytes,
                 allocatedBytes: allocatedBytes

Sources/UntoldEngine/RuntimeAssets/GaussianPagingPolicy.swift

@@ -261,10 +261,11 @@ public enum GaussianPagingPolicy {
 
     /// The pool's slot count: what the residency budget leaves after the pools already
     /// allocated, clamped between the minimum and the smaller of the asset and the platform cap.
-    public static func poolSlotCount(assetBytes: Int, slotBytes: Int, residencyBudgetBytes: Int, allocatedBytes: Int, poolMaxBytes: Int = pagePoolMaxBytes, minPoolSlots: Int = minPoolSlots) -> Int {
+    public static func poolSlotCount(assetBytes: Int, assetSlotCount: Int? = nil, slotBytes: Int, residencyBudgetBytes: Int, allocatedBytes: Int, poolMaxBytes: Int = pagePoolMaxBytes, minPoolSlots: Int = minPoolSlots) -> Int {
         guard slotBytes > 0 else { return minPoolSlots }
-        // The asset in whole slots: its last tier is short but takes a slot.
-        let assetSlots = (assetBytes + slotBytes - 1) / slotBytes
+        // Every chunk tier occupies its own slot. `assetBytes / slotBytes` is only exact when
+        // chunks are packed continuously; grid partitioning can leave several short chunk tails.
+        let assetSlots = assetSlotCount ?? (assetBytes + slotBytes - 1) / slotBytes
         let poolCap = min(assetSlots * slotBytes, poolMaxBytes)
         let remaining = residencyBudgetBytes - allocatedBytes
         let poolBytes = min(max(remaining, minPoolSlots * slotBytes), max(poolCap, minPoolSlots * slotBytes))

Tests/UntoldEngineRenderTests/GaussianChunkCullTest.swift

@@ -38,15 +38,16 @@ final class GaussianChunkCullTest: BaseRenderSetup {
     private var legacyTwin: GaussianLegacyTwin?
     private var indexResolver: GaussianSplatIndexResolver?
 
-    /// The 200-splat fixture baked with 16 splats per chunk: 13 chunks.
-    private let expectedChunkCount = 13
+    /// The 200-splat fixture baked with 16 splats per chunk: 13 chunks by count alone, 16 once
+    /// the grid partition (gridChunkPlan) also cuts a chunk at each of its 8 cells' boundary.
+    private let expectedChunkCount = 16
 
     /// Visible chunks at each of `cameras`, so the frustum boundary is known to be exercised.
-    private let expectedVisibleChunkCounts = [13, 8, 12]
+    private let expectedVisibleChunkCounts = [16, 6, 15]
 
     /// Three views of the fixture (x, y in ±1.4, z in −0.05…0.55, 16-splat Morton chunks): the
-    /// whole asset from afar (13 of 13 chunks), a close view of the +x/+y corner that leaves
-    /// five chunks outside the frustum (8 of 13), a close view of the −x side (12 of 13).
+    /// whole asset from afar (16 of 16 chunks), a close view of the +x/+y corner that leaves
+    /// several chunks outside the frustum (6 of 16), a close view of the −x side (15 of 16).
     private let cameras: [(eye: simd_float3, target: simd_float3)] = [
         (simd_float3(0, 3, 7), .zero),
         (simd_float3(1.0, 1.0, 0.6), simd_float3(1.0, 1.0, 0)),

Tests/UntoldEngineRenderTests/GaussianChunkLevelTest.swift

@@ -301,6 +301,24 @@ final class GaussianChunkLevelTest: BaseRenderSetup {
         return areas
     }
 
+    /// Each kept chunk's own density (its real splat count over its area) — not a fixed
+    /// splatsPerChunk, since the grid partition (gridChunkPlan) leaves many chunks short.
+    private func mirrorDensities(_ fixture: Fixture) throws -> [Int: Float] {
+        let areas = try mirrorAreas(fixture)
+        return areas.reduce(into: [:]) { densities, entry in
+            densities[entry.key] = Float(fixture.index.chunks[entry.key].splatCount) / entry.value
+        }
+    }
+
+    /// `chunkLevels`-style helpers pass a chunk's own splat count, not a fixed one, to
+    /// `GaussianChunkCullMath.level` — the grid partition leaves many chunks short of
+    /// `splatsPerChunk`, so the level rule and this mirror must agree on each chunk's real count.
+    private func levelsByChunk(_ fixture: Fixture, areas: [Int: Float], floor: Float, previous: Int, shifts: (Int, Int)) -> [Int: Int] {
+        Dictionary(uniqueKeysWithValues: areas.map { chunk, area in
+            (chunk, GaussianChunkCullMath.level(densityCap: .infinity, densityFloor: floor, splatCount: fixture.index.chunks[chunk].splatCount, screenArea: area, previous: previous, available: 7, tierShifts: shifts))
+        })
+    }
+
     /// A camera above the slab's centre on the line (0, h, 0.6 h), raised or lowered until the
     /// median visible chunk covers about `pixels` pixels of the 1920 × 1080 viewport.
     @discardableResult
@@ -593,7 +611,7 @@ final class GaussianChunkLevelTest: BaseRenderSetup {
         // floor, where the synthetic slab's tiny splats have already dropped below a pixel.
         let view = try placeCameraWithMedianChunkPixels(300, fixture)
         XCTAssertEqual(view.medianPixels, 300, accuracy: 30)
-        let medianDensity = try mirrorAreas(fixture).values.map { 1024 / $0 }.sorted()[fixture.chunkCount / 2]
+        let medianDensity = try mirrorDensities(fixture).values.sorted()[fixture.chunkCount / 2]
         let medianTier = GaussianChunkCullMath.densityTier(density: medianDensity)
         XCTAssertGreaterThanOrEqual(GaussianChunkCullMath.densityTier(density: densityFloor) - medianTier, -4, "sanity — fine at the default floor")
 
@@ -809,7 +827,7 @@ final class GaussianChunkLevelTest: BaseRenderSetup {
         XCTAssertGreaterThan(areas.count, 200)
         // The median chunk's density tier; floors that put it at Δ = −1 (fine), −(s1 + 1) (just
         // inside level 1), then one and two tiers back up.
-        let medianDensity = areas.values.map { 1024 / $0 }.sorted()[areas.count / 2]
+        let medianDensity = try mirrorDensities(fixture).values.sorted()[areas.count / 2]
         let medianTier = GaussianChunkCullMath.densityTier(density: medianDensity)
         let fineTier = medianTier - 1
         let coarseTier = medianTier - shifts.0 - 1
@@ -820,9 +838,7 @@ final class GaussianChunkLevelTest: BaseRenderSetup {
         }
         func chunkLevels(at tier: Int, previous: Int) -> [Int: Int] {
             let floor = GaussianChunkCullMath.densityTierFloor(tier) * 1.19
-            return areas.mapValues { area in
-                GaussianChunkCullMath.level(densityCap: .infinity, densityFloor: floor, splatCount: 1024, screenArea: area, previous: previous, available: 7, tierShifts: shifts)
-            }
+            return levelsByChunk(fixture, areas: areas, floor: floor, previous: previous, shifts: shifts)
         }
         // The chunks the rule keeps fine at the fine floor and sends to level 1 at the coarse floor.
         let switching = Set(chunkLevels(at: fineTier, previous: 0).filter { $0.value == 0 }.keys).intersection(chunkLevels(at: coarseTier, previous: 0).filter { $0.value == 1 }.keys)
@@ -1013,7 +1029,7 @@ final class GaussianChunkLevelTest: BaseRenderSetup {
             let fadeFrames = GaussianPagingPolicy.fadeFrames
             try placeCameraWithMedianChunkPixels(600, fixture)
             let areas = try mirrorAreas(fixture)
-            let medianDensity = areas.values.map { 1024 / $0 }.sorted()[areas.count / 2]
+            let medianDensity = try mirrorDensities(fixture).values.sorted()[areas.count / 2]
             let medianTier = GaussianChunkCullMath.densityTier(density: medianDensity)
             let fineTier = medianTier - 1
             let level1Tier = medianTier - shifts.0 - 1
@@ -1024,9 +1040,7 @@ final class GaussianChunkLevelTest: BaseRenderSetup {
             }
             func chunkLevels(at tier: Int) -> [Int: Int] {
                 let floor = GaussianChunkCullMath.densityTierFloor(tier) * 1.19
-                return areas.mapValues { area in
-                    GaussianChunkCullMath.level(densityCap: .infinity, densityFloor: floor, splatCount: 1024, screenArea: area, previous: 0, available: 7, tierShifts: shifts)
-                }
+                return levelsByChunk(fixture, areas: areas, floor: floor, previous: 0, shifts: shifts)
             }
             // The chunks fine at the fine floor, level 1 at the level-1 floor, level 2 at the level-2 floor.
             let switching = Set(chunkLevels(at: fineTier).filter { $0.value == 0 }.keys)
@@ -1321,7 +1335,7 @@ final class GaussianChunkLevelTest: BaseRenderSetup {
         GaussianRuntimeLimits.workingSetSplatsOverride = nil
         GaussianSharedWorkingSet.shared.resetBudgetHysteresis()
         placeGaussianTestCamera(eye: simd_float3(0, 8, 14), target: simd_float3(0, 0, -3))
-        let densities = try mirrorAreas(fixture).values.map { 1024 / $0 }.sorted()
+        let densities = try mirrorDensities(fixture).values.sorted()
         XCTAssertEqual(densities.count, fixture.chunkCount, "every chunk in view")
         let spread = GaussianChunkCullMath.densityTier(density: densities[densities.count - 1]) - GaussianChunkCullMath.densityTier(density: densities[0])
         XCTAssertGreaterThanOrEqual(spread, 4, "sanity — the densities spread over four tiers or more")