Any progress or hint on implementing GPU sorting?

[Yang-Xijie]

I find the following sentence in README in commit 8876629:

MetalSplatter/README.md

Line 23 in 8876629

* Sorting on GPU. Sorting is currently done on the CPU asynchronously at a lower framerate (~10 fps), which increases how often you'll see pops especially when the viewpoint changes quickly

which was removed in commit 2213483:

MetalSplatter/README.md

Line 13 in 2213483

### TODO

However, it seems sorting still happens on the CPU side:

MetalSplatter/MetalSplatter/Sources/SplatRenderer.swift

Line 608 in ec6d0f2

orderAndDepthTempSort.sort { $0.depth > $1.depth }

---

The performance of CPU sort is N*log(N). Here are some performance report (on M2 Pro MacBook Pro, release mode, random numbers):

〇 UInt32
1000000: 0.071 s
10000000: 0.759 s
〇 Float32
1000000 0.084 s
10000000 0.934 s

impeding real-time rendering of 3D gaussians...

---

It seems that radix sort on GPU for Metal is missing. (https://developer.apple.com/forums/thread/105886)

I did find an implementation on Apple Silicon, but it is difficult for me to understand. (https://github.com/ShoYamanishi/AppleNumericalComputing/tree/main/05_radix_sort#54-metal-radix-sort-implementations)

Will there be someone implementing the sorting on GPU to improve the performance? Or is there some library we can directly adopt to accelerate?

[Yang-Xijie]

Not sure if https://github.com/ml-explore/mlx/blob/b05bcfd27f5f1293401b74dce02e38c8fd7ef66a/mlx/backend/metal/kernels/sort.metal could be directly used...

https://github.com/ml-explore/mlx/blob/b05bcfd27f5f1293401b74dce02e38c8fd7ef66a/mlx/backend/metal/kernels/sort.h

// Based on GPU merge sort algorithm at
// https://github.com/NVIDIA/cccl/tree/main/cub/cub

[kemchenj]

@Yang-Xijie I have some trouble using MLX sort/argSort. So I would say don't.

And actually MetalPerformanceShaderGraph provide some ML operators like sort and argSort that we could use here. They have better performance and stability compared to MLX.

Here is a simple implementation based on MPSGraph https://gist.github.com/kemchenj/26e1dad40e5b89de2828bad36c81302f. Feel free to use it if you got interest.

[vade]

https://gist.github.com/kieber-emmons

has a Radix sort in metal - i just stumbled on to this, might give it a whirl

[vade]

FWIW, this might be a path forward - sorry to share 'vibe code' bullshit but - its a helpful starting point to thread stuff together

https://chatgpt.com/share/67f70632-c864-8001-9f5c-df3db25e5a68

[scavenx]

I implemented this a few months ago based on @kemchenj's MPSArgSort. Just pushed it to GitHub, check commit dfab12c. Let me know if it works for you.