make sort32 fast

[39ali]

try to improve the performance of sort32, on avg it's 30-40% faster .

things that changed :

• pass 2 no longer re-reads keys[] , scratch stores a packed u64 of (inverted_key << 32 | original_index). pass 2 reads the high 16 bits directly from scratch with kv >> 48 making it a sequential scan

• histogram and scatter are now branchless to help llvm vectorize the loop

• manually unrolled histogram and both scatter passes to 8-wide

[mrxz]

Is there a specific reason this is marked unsafe? It compiles just fine without.

[39ali]

i had many experiments with simd, which didn't make it marginally faster so i removed it for simplicity sake but forgot to remove the unsafe, will clean it up

[mrxz]

Awesome work, gave it a try and can confirm that it improves sorting performance. In my limited testing I saw ~20% reduction in sorting time (~25% faster).

manually unrolled histogram and both scatter passes to 8-wide

Without this change the performance gain seems to be roughly the same, or at least I didn't observe any significant difference. The majority of the benefit seems to come from making it branchless.

[39ali]

@mrxz i squeezed a bit more performance ~<=1ms by removing more branches from hot loops, and what you noticed seems about right, it will differ from one wasm engine to another, and arch to another(specially cache sizes and arch) so it's hard to give a solid number but it'll still be a pump in performance

[dmarcos]

can you remove the changes in the dist directory?

[asundqui]

@39ali great work! This looks like a cool win indeed, thanks for the work. Could remove the build in the dist/ folder from your branch, then we can merge it in?

[mrxz]

Could the macros used for unrolling the loops also be used for the body of remainder loops? Both should be identical, so if we could avoid the duplication we avoid the risk of it ever getting out of sync.

[dmarcos]

@39ali any chance for you to implement @mrxz suggestions? Thanks so much for the contribution

[39ali]

I'll implement the changes

[39ali]

@mrxz @dmarcos @asundqui done !

[asundqui]

@39ali really great work here! I've done some benchmarking on your method, and I'm actually getting 2x - 4x speedups in sorting from this. I'm truly shocked that this was possible! This will have a great impact on Spark's sorting performance. On a 10M splat scene on my M3 it goes from 250 ms to 60 ms or so. It's possible that the speedup is not as great on other environments, such as @mrxz was reporting 25% speedups on his system.

I went through and carefully separated the optimizations and measured them:

• I found that approx 65% of the gain could be explained by storing (key, index) as a packed u64 in scratch, which turned the next loop from a random gather into a sequential read. So much better cache performance as a result.

• The next 20% came from doing unsafe { unchecked... } array accesses where the compiler couldn't be sure that it would always be in bounds, so it has to check it every iteration. I went through the logic and it looks like it should always be in bounds.

• The final 15% or so came from unrolling the loops. I had thought the unrolling couldn't possibly help because of branch prediction + cpu instruction reordering but it all helps!

It did seem like there was one error though: the second branchless loop seems problematic... I think writing to the array and only advancing the pointer if it's "valid" could overwrite things. So I removed it. I don't think it does very much for the performance anyway.

Finally I reverted some unnecessary changes to make it closer to @mrxz 's original formulation. I think we should merge this in @dmarcos , @mrxz ! WDYT? This should really help with #225 .

Interestingly, because the sorting is so much faster, it sort of exposes the next bottleneck more: uploading the ordering frequently to the GPU can cause stuttering sometimes when the counts get large. Now this happens more often!

[mrxz]

On a 10M splat scene on my M3 it goes from 250 ms to 60 ms or so. It's possible that the speedup is not as great on other environments, such as @mrxz was reporting 25% speedups on his system.

There is bound to be some variability between setups, but that's quite the discrepancy between the measurements. My guess would be that besides the gains of being branchless and more cache friendly it somehow avoids a slow-path on the M3? Regardless, it's a net positive, even the comparatively modest speed-up I've measured is a very welcome improvement.

It did seem like there was one error though: the second branchless loop seems problematic... I think writing to the array and only advancing the pointer if it's "valid" could overwrite things. So I removed it. I don't think it does very much for the performance anyway.

The overwriting was intentional AFAICT. Since the 'invalid' entries aren't tallied you don't want to count them when scattering either. The corollary is that by only advancing after writing a valid entry, you guarantee that the final write to the array at a given position is a valid entry. The only way this goes wrong is once a bucket has been fully scattered, as then it'll point into the start of the next bucket. The inverted Infinity value, however, has 0xFFFF as its low bytes, always placing it in the final bucket which can safely overflow up to max_splats.

That said, I do think the explicit check is more readable and some quick testing doesn't show a meaningful performance difference on my end between the two.

I think we should merge this in @dmarcos , @mrxz ! WDYT? This should really help with #225 .

I don't see any reason not to merge this.

Regarding #225 it will most definitely reduce the time a stale ordering is visible on screen. Whether or not this will be enough is the question. Similar to the pesky FOUC in web-design, no matter how short it will remain an "issue". Not sure what mkkellogg did differently, but it at the very least it degenerated more gracefully. It does have logic to detect large camera changes and queues partial sorts, though even before a sort update arrives, the rendered frame somehow looks less bad (subjectively).

[dmarcos]

Thanks everyone. Good stuff!