Cuda Static Map Rust: Building a GPU Hash Map from Scratch 1/x

From Rust type system alignment hacks to GPU vectorisation, laying the groundwork for GPU hashing

Blog image

The goal

We're building a statically-sized hash map that runs entirely on CUDA-enabled GPUs, written in Rust. It's roughly Rust's HashMap, but with a few constraints:

This is the first in a series of posts documenting the journey. This one covers what the goal is, surviving the dependency nightmare, a build setup that compiles for both CPU and GPU, and implementing the core building blocks.

The dependency nightmare

The first question in any Rust GPU project is "how do I even compile Rust code for the GPU?" And unfortunately it's not as simple as adding #[target = "nvptx64-nvidia-cuda"] to your Cargo.toml.

The Rust CUDA ecosystem is still in its early days. The old cuda_builder crate on crates.io was last updated 4 years ago. Then there's the Rust-CUDA project. For some reason, I decided to not use Docker, and duck my system dependancies instead.

I started by trying to build rustc_codegen_nvvm (the codegen backend that compiles Rust to NVVM IR, which then becomes PTX) from source. This meant installing:

I wasted time wrestling with version mismatches and mysterious linker errors before deciding I don't want to fight this battle. Life is too short to debug why libnvvm.so can't find libcudart.so.12 when you're just trying to write a hash map.

I just chickened out (good to have fear sometimes I must say), and just used a docker image instead.

It took me a bit to get the docker environment correctly, but was much better than bricking my system :)

Build setup

With the dependency nightmare out of the way, on to what we find fun, actually building things. The catch is that the code has to be compiled twice, once for the CPU (host code) and once for the GPU (device code).

Host code runs on your CPU. It's in charge of:

Device code runs on the GPU and handles things like:

In Rust you split the two with conditional compilation, #[cfg(target_arch = "nvptx64")]. Anything marked with it only compiles when targeting the GPU, so the same code can work in both contexts:

#[cfg(target_arch = "nvptx64")]
// This code only compiles for GPU
use cuda_std::thread;

#[cfg(not(target_arch = "nvptx64"))]
// This code only compiles for CPU
use cust::memory::DeviceBuffer;

For the build itself I went with xtask, a pattern for build automation in Rust projects. Instead of remembering whether you need SOME_OBSCURE_FLAG cargo build --target nvptx64-nvidia-cuda or SOME_OBSCURE_FLAG2 cargo build, you just run:

Under the hood our xtask deals with the annoying bits:

The project layout is clean and idiomatic:

Splitting it this way means we write the code once and it works in both contexts. We need that here, because the map has to be usable from the CPU (for setup) and from the GPU (for operations).

The building blocks

Next, the actual data structures. A hash map needs three fundamental pieces:

  1. Pair, a key-value pair that can be stored in GPU memory
  2. Storage, a buffer that holds all the pairs
  3. Probing, a strategy for finding empty slots when collisions occur

When Pair alignment becomes a type system puzzle

Pair looks simple enough, it's just a key and a value. But GPU memory has alignment requirements, and the pairs need to be aligned correctly for optimal performance. In C++ this is straightforward:

template<typename First, typename Second>
constexpr size_t get_alignment() {
    return std::min(16, std::bit_ceil(sizeof(First) + sizeof(Second)));
}

struct alignas(get_alignment<First, Second>()) pair {
    First first;
    Second second;
};

Rust can't do this directly. #[repr(align())] requires a literal value, not a computed one. You can't write #[repr(align(compute_alignment()))], the compiler just rejects it.

After receiving help from some lovely folk at the Rust Forum (RIP Stack Overflow for niche Rust questions), I figured we can't use computed values in attributes, but we can use types. So the trick is to map each alignment value to a type, and then let that type enforce the alignment.

#![feature(generic_const_exprs)]
#![allow(incomplete_features)]

/// Computes the alignment for a pair type.
pub const fn alignment<First, Second>() -> usize {
    let x = (size_of::<First>() + size_of::<Second>()).next_power_of_two();
    // Const-compatible min: if x > 16, return 16, else return x
    if x > 16 { 16 } else { x }
}

/// Trait for mapping alignment values to aligned types
pub trait AlignedTo<const ALIGN: usize> {
    type Aligned: Clone + Copy + core::fmt::Debug + cust_core::DeviceCopy;
}

/// Type alias to get the aligned type for a given alignment value
pub type Aligned<const ALIGN: usize> = <() as AlignedTo<ALIGN>>::Aligned;

/// Macro to generate `AlignedTo` implementations for specific alignment values
macro_rules! aligned_to {
    ( $($align:literal),* $(,)? ) => {
        $(
            const _: () = {
                #[repr(align($align))]
                #[derive(Clone, Copy, Debug)]
                pub struct Aligned;
                
                impl AlignedTo<$align> for () {
                    type Aligned = Aligned;
                }
                
                unsafe impl cust_core::DeviceCopy for Aligned {}
            };
        )*
    }
}

// Generate implementations for all possible alignment values (1, 2, 4, 8, 16, 32)
// The alignment function returns min(16, next_power_of_two(size)), so these cover all cases
aligned_to!(1, 2, 4, 8, 16, 32);

How the magic works, in four parts:

  1. generic_const_exprs. This unstable feature lets us use generic type parameters in constant expressions. Without it, there's no way to compute alignment::<First, Second>() at compile time.

  2. The trait system. AlignedTo<const ALIGN: usize> maps alignment values to types. For each alignment value (1, 2, 4, 8, 16, 32), the macro generates a zero-sized struct with that alignment using #[repr(align(N))].

  3. Fully Qualified Syntax (FQS). The type alias Aligned<const ALIGN: usize> uses FQS to look up the right aligned type, <() as AlignedTo<ALIGN>>::Aligned. You can read it as "treat () as implementing AlignedTo<ALIGN>, and give me its Aligned type."

  4. The zero-sized marker. Pair gets a zero-sized _marker field with the computed alignment:

pub struct Pair<First, Second>
where
    (): AlignedTo<{ alignment::<First, Second>() }>,
{
    _marker: Aligned<{ alignment::<First, Second>() }>,
    pub first: First,
    pub second: Second,
}

Even though _marker takes up no space, Rust requires the struct's alignment to be at least the maximum of all field alignments. So if _marker has alignment 8, the entire Pair must be aligned to at least 8 bytes. It's a hack, but a beautiful one: the type system ends up enforcing a runtime property (memory alignment) at compile time. It's a perfect example of Rust's type system being flexible enough to work around language limitations, even if it takes some creative thinking.

Initializing Storage

Storage manages the device memory buffer that holds all the key-value pairs. When you allocate GPU memory it's uninitialized, so it has to be filled with sentinel values (special key-value pairs that mean "this slot is empty").

For small maps you can just build a vector on the CPU and copy it over to the GPU. For large maps (millions of slots) that's slow, and it's better to launch a GPU kernel that initializes all slots in parallel.

Our Storage supports both:

// Option 1: Host-to-device copy (simple, works for small maps)
storage.initialize(sentinel, None)?;

// Option 2: Kernel-based initialization (fast for large maps)
let module = Module::from_ptx(PTX, &[])?;
storage.initialize(sentinel, Some(&module))?;

The kernel version launches thousands of threads, and each one initializes a chunk of memory. Copying 1GB of data from CPU to GPU is slow. Having 10,000 GPU threads each initialize 100KB in parallel is fast.

Probing

When you insert a key, you hash it and use the hash to pick a slot. If that slot is already taken, you need a probing scheme, i.e. a strategy for finding the next available slot.

We implemented two classic ones.

Linear probing is the simplest. If slot h is taken, try h+1, then h+2, and so on, wrapping around when you hit the end. It's efficient when occupancy is low, but it can suffer from clustering (long chains of occupied slots).

Double hashing uses two hash functions. The first picks the starting position and the second picks the step size, so the probe sequence goes h1(key), h1(key) + h2(key), h1(key) + 2*h2(key), etc. That cuts down on clustering, which makes it the better pick once occupancy gets high.

Both implement a ProbingScheme trait that creates a ProbingIterator, an iterator that yields the sequence of slot indices to check. The iterator deals with wrap-around itself, so the calling code just checks each slot until it finds an empty one (or figures out the key doesn't exist).

Parallel linear probing with cooperative groups

Everything above is scalar, one thread checking one slot at a time. But this is a GPU, so the work can be split across threads.

Take a cooperative group of cg_size threads working together. Instead of one thread checking one slot, several threads check several slots in parallel. That's what bucket_size is for, each thread handles bucket_size consecutive slots. So the total "stride" one cooperative group covers is stride = bucket_size * cg_size.

Blog image

The hash table is conceptually split into "groups" of size stride. Hashing a key tells you which group it belongs to, group_index = hash(key) % (capacity / stride), and the probe sequence starts at the beginning of that group, init_base = group_index * stride.

Within a group, threads get ranks from 0 to cg_size - 1, and each one starts at its own offset, init = init_base + thread_rank * bucket_size. So if cg_size = 4 and bucket_size = 2, thread 0 checks slots 0-1, thread 1 checks slots 2-3, thread 2 checks slots 4-5, and thread 3 checks slots 6-7, all within the same group.

On each iteration, all threads in the group advance together by stride. This means they jump to the next group and check their assigned slots there. It's like a synchronized dance: all threads move in lockstep, but each thread is responsible for its own bucket of slots.

Next up

Where things stand after this post:

Hopefully soon, we'll put it all together and implement the actual StaticMap type (the public API that users will interact with). That means bulk operations (inserting thousands of keys at once), device-side operations (querying the map from GPU kernels), and the sentinel value system that makes it all work.

Until then, happy hashing! 🦀