Cuda Static Map Rust: When the compiler lies (atomics, tiling and feature completion) 2/x

Compiler bugs, scoped atomics and cooperative tiling, and how we got to feature completion

Where we left off

In the first post, we built the basic pieces: pairs, storage and probing. This time the goal was to go from simple "hashing in place" to a full static map, which means bulk insert, find and contains. We also added cooperative group tiling, so a warp can probe together instead of one lonely thread at a time.

Getting there wasn't only about the code. It also meant getting scoped atomics right so the GPU respects our memory model, figuring out cooperative group tiling so we aren't wasting hardware cycles, and (my personal favorite) a pile of compiler bugs that showed up the second we started passing references and atomics into kernels.

Along the way we added proper hashing and a way to count successful inserts, plus storage that respects group size (and prime capacity for double hashing). I also split the world into two: host-side setup and a small device-side "ref" passed to every kernel. That ref is 16-byte aligned.

The 16-byte alignment bug

The first time we passed that ref into a kernel, it died with IllegalAddress. We got garbage pointers, and the kernel was reading zeros where it expected a struct.

I traced it back to an ABI mismatch in the compiler backend. The host packed our struct with 16-byte alignment, but the compiler generated PTX expecting 8-byte alignment. So the kernel was reading from the wrong offset, and the pointers it pulled out pointed at nothing valid.

The fix was to tell the compiler to treat 16-byte-aligned structs specially, by casting them to 128-bit integer chunks. That had to go upstream, otherwise this post would end right here.

Atomics and scope

On GPUs, atomics are scoped.

The scope tells the hardware who needs to see the operation, whether that's the whole system, the whole device, or just the thread block. So for every compare-and-swap we had to pick the right scope and make sure the generated PTX actually matched what we meant.

Then our atomic loads started misbehaving. And not in a "maybe I wrote a bug" way. Silly me, thinking the compiler would emit the right PTX.

We were using device-scoped atomics, but the macros emitting the load/store instructions were pasting the Rust scope name ("device") into the instruction string instead of the PTX scope ("gpu"). So a relaxed device load became ld.relaxed.device.u32. PTX doesn't know what "device" means in that context, it expects "gpu".

With the wrong scope the hardware makes no guarantees, and in our case that meant the atomics were effectively broken. Once we fixed the macros to emit the PTX scope names, the hardware finally did what it was told.

Cooperative group tiling

So far it was one thread per key, each with its own probe sequence. Fine for a Tuesday, but GPUs are happier when the threads in a warp cooperate, so we added cooperative group tiling. A small tile of threads (2, 4, 8, or a full warp of 32) shares the work of one logical operation.

For an insert, every thread in the tile probes its slice of the table. They sync up and vote: "Did anyone see this key?" (if so, stop, it's a duplicate) or "Who found an empty slot?" We pick a winner to do the actual compare-and-swap, then broadcast the result to the whole tile.

That only works if the probe sequence is tile-aware, meaning the whole tile advances in lockstep. We assign each thread a rank and make the step size equal to the tile width. It's like a search party: everyone checks their own spot, then the whole line steps forward together.

Of course, two more compiler bugs bit us here.

First, the warp vote functions ("did any lane satisfy this?") returned the opposite boolean. The compiler was checking "if predicate is true, return 0, else return 1" and then returning "result != 0". So when everyone agreed, it returned false. Our "did anyone find it?" checks were inverted, and we'd exit early exactly when we found it. We flipped the mapping and sanity returned.

Second, the enum for shuffle direction used Rust's default values (0, 1, 2, 3), but PTX expects specific values (Idx=0, Up=1, Down=2, Xor=3), and our declaration order didn't match what PTX expected. So when we asked to "broadcast from lane X," the compiler emitted "shift down by X."

When wrapping low-level intrinsics, assign enum values explicitly, don't trust declaration order.

The host API

On the host, all of this now sits behind a single API that does the canonical setup:

• Bulk insert resets the counter, launches the kernel with the right thread/tile configuration, and reads the counter back.
• Bulk find/contains passes the device buffers, launches, and copies results.

With that we're at feature parity. Bulk operations and proper hashing are in, and cooperative tiling is there as an option for trading off thread utilization against keys-in-flight.

What's next

We have a working static map. It supports 64-bit keys and values and both linear and double hashing, and it handles collisions like a champ. The compiler bugs are fixed upstream. Atomics use the right scopes now, and the tiling is cooperative.

Next up: benchmarks. We need to see how fast this thing actually is, and where it can get faster.

Until then, happy hashing! 🦀