# McGridsort: Warping Grids for GPU k-way mergesort _Date: 2026-04-07_ ___ ## Spoiler tl;dr: Throw `k×k` numbers from `k` pre-sorted streams into a square grid, parallel-sort each axis: k parallel sorts of k elements each, 3 axes -- yes, 3 (rows, cols, anti-diagonals). Then, gather candidate `min(k)` elements, `C(k)`, from a provably-smaller, compile-time-known area, sort those and output the `min(k)` elements. Could this be cheaper than sorting `k*k` elements all at once? ## Intro Mergesort. A classic algorithm, probably the first `O(N*logN)` sorting algorithm people learn. Let's make it weird. If you need a refresher, check [[#Typical Mergesort|the typical mergesort idea]]. So let's be concrete: You have a list of size `2^20` which is ~1 million (1M) items. You will definitely sort them, today. This implies that mergesort has to pass through the data ~20 times (well, less because of blocksorts[^1], but don't worry about that). You then think to yourself: why not just... merge more than 2 lists at a time? If you merged `k` lists at a time, then you only need `log_k(N)` passes; e.g. `k=4` implies only 10 passes through the data. Clearly, this implies a more complex algorithm, which maybe doesn't get optimized very well, or has other issues in practice. e.g. The sufficiently smart compiler was not sufficiently smart enough. But you ignore this because you're not looking for "practical", you're looking to optimize for "lol, that's weird". So I will present to you the aftermath of a shower thought 2 years ago: a SIMD-accelerable k-way mergesort, mostly with CUDA in mind. We'll focus on the GPU story for the rest of the post, and unfortunately, I will assume the reader has context around CUDA (or can just ask their local Claude). Note: An older post on [[Faster sorting with SIMD CUDA intrinsics|CUDA bitonic sorting]] lays some groundwork, while another blog post ([[Non-Messing-Up++]]) lays out a supporting theorem with a corresponding proof. Neither are necessary reading for this post, but may be useful in case you want to read more. Pedantic note: This isn't a "full" mergesort. It's a merge engine that a larger mergesort scaffold can utilize, i.e. instead of a normal 2-way merge pass. ## Naive k-way merge So imagine you have `k` sorted streams in VRAM/global memory (GMEM) on your Nvidia GPU. Your threadblock (CTA) selects the min element from the front of those `k` streams. Do this `n` times, and you've merged `n` elements from `k` streams into one sorted stream. Of course, doing this one element at a time is silly, even if we only focus on a single SMSP, which is effectively a "core" (an SM has four of these) comparable to a CPU core. To be clear, this is **not** a CUDA core, which is more like a SIMD lane. For integer operations, 1 SMSP has 16 execution units (32 for FP) -- unless you're on Blackwell, which unified all INT and FP units, so now we can do INT things 32 at a time! Ahem -- but back to the algo. Imagine that you pull a bunch of elements from those `k` streams into shared memory (smem), which is just under 2 orders of magnitude faster than GMEM. Why is smem lowercase? Sorry, I ended up writing things like this in my docs... too bad! So now, we can naively just select `min(k)` elements from the `k` streams, one at a time right? In fact, we can do this "faster" if `k <= 32`. One warp can load all `k` front elements, reduce with warp shuffles (`log_2(32)` = 5 SHFLs) to find the minimum, output and repeat. As it turns out, warp shuffling is ~2x faster than shared memory read-then-write, so we use `__shfl_sync` warp primitives here. Obvious question: Can we do better? Obvious answer: Yeah, because otherwise people would just do this (at least for k=2). But like... **how** do we do better? (No, don't look up GPU-MMS, don't do it!) The \<typical-adjective-here>-reader can read ahead to the [[#But then what?|properties discussion]]. Well, if we can use a fast in-register sorting "primitive" (i.e. not touching smem), maybe we can do something with that? This implies a sorting network, like [bitonic sort](https://en.wikipedia.org/wiki/Bitonic_sorter). With that base idea, and with no further ado, let's put the McGridsort in the bag. ___ ## The Core Idea Instead of solving a `k`-select problem every time, can we use some data structure to do this more quickly? Well, yes. Heapsort is a prototypical example (for `k=2`). Here's mine. Place an `n×n` grid in smem. For concreteness, take `n = 32`: a 1024-element grid occupying 8 KB at 64-bit values. Each of `k` sorted streams (assume min-sorted) feeds `(n^2)/k` elements into the grid. If `n = k`, we can assume each column comes from 1 stream. That's nice so we'll just assume `n = k = 32` in this post. The grid maintains a **triple sorting invariant**: every row is sorted left-to-right, every column is sorted top-to-bottom, and every anti-diagonal is sorted bottom-left to top-right. An anti-diagonal (AD) is a diagonal of the grid which goes from bottom-left to top-right. The simultaneous sorting of the 3 axes is possible by [[Non-Messing-Up++]]. On each extraction round, we pull out the `b` smallest elements (`min(b)`) from the grid and emit them to the output. Then we refill `b` positions from the input streams and restore the invariant. The batch size `b` is a tuning parameter, independent of both `n` and `k`. It turns out to have a sharp sweet spot(s) that we'll get to later. The idea is that the triple-invariant (TI) allows us to more quickly extract the `min(b)` elements, even considering the cost of maintaining the triple-invariant. That's a bold claim but not implausible since each col/row/AD is independent of each other. This means that even a single-core CPU setup can parallelize sorting them with ILP, not even necessarily requiring warp-level parallelism (WLP) on CUDA. Whether or not the claim holds up in reality is a separate story, though. Hey, this is just an idea-post, no claims about real-world numbers, okay? Performance is always harder than it seems! Also, apparently, Claude likes to call this grid a priority queue -- I don't particularly like the framing, but it seems like an approximately-correct interpretation. The global minimum *is* in the top left, though more importantly, we also constrain other almost-minima elements around that area. ## Algorithm Details ### Notation/Variables `N`: the total number of elements to sort. It's big. `n`: the size of one side of our square grid. `k`: the number of streams we're merging. We also use it as our SIMD-parallelism size, chosen to be 32 (or 16). `b`: the number of elements we want to extract from our grid per loop. `min(b)`: the minimum `b` elements, i.e. min(5) is the 5 smallest elements of the grid. `C(b)`: the positions of the grid which can possibly contain the `min(b)` elements. ### Finding min(b) The `min(b)` elements of the `k×k` grid, `G` -- where are they? Without Young Tableau (YT) or triple-invariant (TI), we don't know. Well, that's not quite true: `min(1)` is obviously in the top row (row 0), where the front of the `k` streams are (sorry, in my mind, the streams feed in from below, which contradicts gravity). Well, that's an interesting observation. What if we sorted each row from left to right? Then the `min(1)` element would be forced into the top left corner at `(0,0)`! What about `min(2)`? Besides the global min element at `(0,0)`, the 2nd smallest element will have to be either below it or to the right: `{(1,0),(0,1)}`. Then, as `min(x)` goes from x=1, to 2, to 3, to `b`, the "uncertainty" of where the `min(b)` elements are grows. We'll call this candidate set of possible `min(b)` elements `C(b)`. In fact, `|C(b)|` grows approximately log-linearly, (`O(n*log(n))`), but we won't show that here. Notice that what we've constructed so far is a grid which is YT (satisfies Young Tableau conditions): columns and rows are both sorted, and so we'll specify `C_YT(b)` to indicate `C(b)` in a YT grid. What if we could continue constraining `min(b)`? Being uncreative, we decide, hey, let's sort the anti-diagonals, making the grid TI (triple-invariant holds) (note that YT implies the main diagonals running top-left to bottom-right are already sorted). This implies that `min(2)` can no longer include the position `(0,1)` (to the right of `(0,0)`), or else `(1,0)` would be both smaller than it (by AD sort, which places smaller elements to the bottom-left) and larger than it (since we know `min(1)`, the only smaller element, is in `(0,0)`), a contradiction. Unfortunately, for `min(3)`, the 3rd smallest element is still free to be in 2 possible places, `{(2,0), (0,1)}`, and each successively larger `min(x)` still grows in their uncertainties. Great, we did a lot of extra sorting, but we've only constrained 1 element? Not so fast. It turns out that TI's `|C(b)|` (size of `C(b)`) is ~35% less than `|C_YT(b)|` (slightly different for different values of `b`). Here's what `C(32)` actually looks like: ![[candidates_yt-vs-ti.png|720]] Notice how the number of possible candidates per column drops off sharply. However, if we wanted to avoid sorting 128 elements for a cheaper `sort64`, we still have to drop to `b=29`, instead of `b=32`, which has 69 candidates... nice, but not nice enough for our purposes. Also notice: for any given `n` and `b`, the locations of `C(b)` are compile-time fixed. We can safely just gather the elements at those known positions and immediately process them. ### C(b) and P(r,c) Define `P(r,c)` to be "how many other positions can be smaller than (r,c)": the number of possible positions with elements which are "predecessors" to the element at `(r,c)`. For example, with YT, `(1,1)` has 3 positions on the grid whose items can possibly be smaller than it. Then, to find `C(b)`, we need to find all positions `(r,c)` such that `P(r,c) < b`. e.g. for `b=3`, we just have `P(0,0) = 0, P(1,0) = 1, P(0,1) = 2, P(2,0) = 2` -- 4 positions, 2 of which come from the possible positions of the 3rd-smallest element. TI constrains `C(b)` better than YT because it gives us more information at each position, halving the "incomparable" (unknown) regions: ![[info_yt-vs-ti.png|720]] More technically (for the upper left triangular half of the grid): ``` C(b) = {(r,c) : P(r,c) < b} P_yt(r,c) = (r+1)*(c+1) - 1 P(r,c) = (r+1)*(c+1) + c*(c+1)/2 - 1 ``` The extra term `c*(c+1)/2` grows quadratically in `c`, killing off high columns fast. Setting `r = 0`: `P(0,c) = (c+1) + c*(c+1)/2 - 1 = c*(c+3)/2`. Column `c` drops out entirely when this exceeds `b`. For `b = 32`, columns >= 7 are excluded (`P(0,7) = 35`). The formula has a second case for `r + c >= n`, handling positions where the triangle clips against the grid boundary, but we won't cover it here. Here's what `|C(b)|` looks like in practice (`n = 32`): | `b` | \|C(b)\| | Pads to | Notes | | ------ | -------- | ------- | -------------- | | 24 | 49 | 64 | | | 28 | 60 | 64 | | | **29** | **62** | **64** | **Sweet spot** | | 30 | 66 | 128 | Jumps past 64 | | 32 | 69 | 128 | | | `b` | \|C_yt(b)\| | | --- | ----------- | | 17 | 52 | | 18 | 58 | | **19** | **60** | | 20 | 66 | For a normal YT grid, the 64-candidate requirement drops `b` down to 19, which means the extra AD sort buys us a ~50% larger batch size for extraction. ### Extraction This is the payoff. We need the `min(b)` out of `|C(b)|` candidates. Load all of `C(b)` from their compile-time known grid positions into registers, pad to the next power of two with `+inf`, and run a full bitonic sort (via warp shuffle instructions). At 64 elements this is cheap — 15 compare-and-swap stages per 32-element half, one shared memory round-trip to merge the halves. The lowest `b` lanes hold the winners (ascending sort). In fact, you can pack a sort64 into one warp by using 2 registers per lane -- this improves max throughput via occupancy when you're maxing out threads per CTA, and we can avoid having to go to smem. However, I suspect that in practice, we would prefer to just use more warps despite the smem round trip, which should be hidden by WLP. Do note that each element needs a metadata tag (the id of the stream it came from), so we need 2x move/select instructions per swap during the sorts (still just one comparison though). ### Refill and Restoration After emitting `b` sorted elements, pull replacements from the appropriate streams (the metadata tags say which), place them at the vacated positions, and restore the triple invariant: sort dirty columns, then all rows, then the upper anti-diagonals (Note: we only need upper triangle ADs sorted for `b < 2*n`, left as an exercise to the reader). Each axis sort is an independent bitonic sort -- embarrassingly parallel across warps (but important, they require sync barriers between phases). Not everything needs re-sorting. For `b = 29`: only columns 0–6 are dirty, but all 32 rows must be re-sorted -- column sorting can rearrange any row. And then we just loop until we're done. ### But then what? Of course, regardless of parallelism, this is still a lot of work. It's very unclear this would buy us anything over a simple "just sort 1024 elements and extract 32", or even the naive "select min". **And** we need a barrier after each axis: sort cols -> barrier -> sort rows -> barrier -> sort ADs -> barrier -> extract. Each barrier prevents sorts from pipelining, increasing the latency of the entire algorithm. We can partially get around this, but that kind of optimization won't be in this post. Outside of the practical performance issues, and besides the fact that I haven't shown a rigorous proof of correctness, I think it's a funny algorithm. That's really the point. But let's go over some merits (well, merits in a vacuum). **k-way merge.** With `k = 32` streams, the merge tree has `ceil(log_32(N))` levels instead of `ceil(log_2(N))`. For `N = 10^6` total elements: ~4 passes vs ~20. **In-place operation.** Besides the obvious `O(n^2)` grid per engine (constant in our case), like other merge engines, our "McGridsort" can just stream output back to the GMEM array without an auxiliary GMEM buffer. ## Optimizations? So I've presented the idea pretty much as I had it years ago, but in the making of this post, I had a bit of fun thinking about possible optimizations. The more obvious one is: why a 32x32 grid? Of course, typically the more you batch, the more you win, and CUDA warps obviously map well to processing 32 elements at a time. ### Grid Downsizing But there are some issues with `k=32`. For example, 32x32 is 4 times more elements than 16x16, while `log(16)/log(32) = 4/5` implies only 20% fewer GMEM passes. A 32-element bitonic sort (`sort32`) has 50% more sequential stages (and thus latency) than `sort16` (15 vs 10) -- bitonic sort depth grows with `O(log(n)^2)`. With `k=16`, we only need 8 warps per axis sort: e.g. for 16 rows, each warp (1 of 8) packs in 2 rows (one in lanes 0..15, the other in 16..31) and runs through 2 `sort16`s simultaneously. The cherry on top? `|C(16)| = 30` in a 16x16 grid -- we can just do a `sort32`. Or even better: `|C(32)| = 54 <= 64`! Except... wait, you can't do that. The 17th smallest element of the `k=16` streams might not have been loaded into the grid yet, since it could be that the streams are skewed enough such that the current `min(17)` is actually concentrated all in one of them -- but we've only loaded the `min(16)` elements from each stream into the grid. ### Double Extraction So, assuming we go with the 16x16 picture, can we get back some throughput? Idea: collect and `sort64` the `C(32)` elements anyway, then extract `min(16)` only. Next, refill 16 elements as per usual. But wait! The 17th to 32nd min elements from `C(32)` are already sorted! This means all we have to do for the next 16-element extraction is to `sort16` the refilled elements, then run the last phase of bitonic sort32 with the new elements + survivors (or rather, just run sort32). This ensures us the next `min(16)` elements without having to recover the grid to TI, effectively amortizing the recovery phase. In fact, `C(64)` is within 128 elements, so... yes, we can continue amortizing the recovery beyond just 1 extra extraction, but that requires extra complexity and more expensive sorts. Then, why not just degenerate into sorting the entire grid? Which would then just be a list at that point? Well there's two answers: 1. it may actually be slower (who knows) and 2. not cool enough. But here's what that would look like with 256 elements: `sort256 -> sort32 -> sort64 -> sort32 -> sort128 -> ...`. ### Gridless Merge Path (I won't explain the merge path algorithm here, but basically it allows for parallelizing the merge of 2 sorted lists by having each thread binary-search where an element should end up, in parallel.) So... even for a k-way merge, why would you use McGridsort instead of something like sort + merge-path? You can sort a 16x16=256-element *list* once, trivially extract `min(16)`, refill as normal, then use merge path to parallel-merge the new (sorted) 16 elements with the rest (240, still sorted). Extract again, loop. This would be a 2-way merge path (merge 16 into 240) that happens to implement a k-way merge, and 2-way merge path is known to be fast. Well, unfortunately, I have no idea why one would use McGridsort, or if McGridsort would actually make sense outside of a funny SIMD idea, especially without any empirical numbers. Ironically, perhaps McGridsort could be better on CPU, where I'm not sure it would be easy to get each lane to run merge-path independently (AVX512, appropriately downsized to 8x8). ### What next? Okay! So that's that. My little merge algorithm which I wanted to get off my chest for years. But where's the code? ...Good question! I'll get back to you on that, been playing around with CUDA-izing this stuff (with benchmarks) over the past couple of weekends, as well as thinking about something which could actually make the performance more interesting: incremental TI restoration. Instead of running through 3 sequential sort phases, as it turns out, it's possible to (mostly) merge in new elements with a single (parallel) insert-1 per column, then per AD, save for a couple edge cases (...all my homies hate edge cases). This is **a lot** faster and implies less resource usage (more occupancy) per engine. However, it's pretty gnarly, so I'd rather not promise anything. That being said, my attempts to understand the mechanism were pretty fun, so even if I "fail" in the overall attempt, it might be worth it as a blog post to leave the reader as annoyed about the edge cases as I am :D But even if I don't pick up that thread anytime soon, I hope my blog will soon introduce an even weirder algorithm: data-parallel Hindley-Milner (DPHM), that is to say, CUDA-"accelerated" type inference. ## Appendix ### Typical Mergesort The typical idea: 1. You have a list of numbers. 2. You want to sort them because you're bored, obviously from smallest to largest. 3. You realize that if the list were made of two sorted sublists, you can just pick the smallest at a time from the front of the two sublists. You decide to call this idea "merging two sorted sublists". 4. You decide to punt the problem of having two sorted sublists to future you. 5. Surprise! Future you is now. Damn it. Fine, let's just _recurse_. If the sublists were just 1 number each, we're golden. If not, keep breaking them down until they are. Then start merging all the way back up. 6. Professor-X.aiGPTCode tells you this is `O(N*log(N))`, using `log(N)` levels each passing over the data of `N` elements. 7. You wonder why/where you have to login and the Professor tells you that `log(N)` means the logarithm of `N`, base 2, or `log_2(N)` if you want to be more precise. ### Runtime Analysis and Performance (Without Hard Numbers) (Assuming we have `k^2` parallel processors) Our McGridsort needs 3x phases of `O(log(k)^2)` parallel depth just to get to TI. A `k^2`-element bitonic sort needs just ~4x more parallel depth compared to a *single* phase to fully sort `k^2` items: `log(k^2)^2 = (2*log(k))^2 = 4*log(k)^2`. In fact, 1024-sort is 55 stages whereas 32-sort is 15, and this is before extraction/refill costs. This implies that McGridsort only makes sense if at least one of the following is true: 1. `k`-sort is significantly faster due to hardware constraints. e.g. imagine if `k`-sort were a hardware primitive, or if warp shuffle hardware were a lot faster than smem hardware (not particularly true currently), etc. i.e. hardware-informed constants grew significantly beyond `k` elements. 2. TI maintenance can be made significantly faster/less resource-intensive than 3 phases of compute-intensive sorting, i.e. incremental restoration. ### Why McGridsort? I think I haven't eaten a McGriddle in too many years. [^1]: https://moderngpu.github.io/mergesort.html: "A blocksort kernel sorts random inputs into tile-length sorted lists, communicating with low-latency, high-bandwidth shared memory" -- i.e. if we do this for 1024-element tiles, we skip 10 (GMEM) passes of the overall mergesort.