Hash-Partition an Array on the GPU in C++ and Return Partition Offsets
Company: NVIDIA
Role: Software Engineer
Category: Software Engineering Fundamentals
Difficulty: medium
Interview Round: Technical Screen
Implement hash partitioning on a GPU in C++. You are given an array of `n` keys and a number of partitions `P`; each key belongs to partition `hash(key) % P` for a hash function `hash`. Rearrange the keys so that each partition's keys are contiguous, partition 0 first, and return the offsets: the index at which each partition starts in the output.
This is the classic way GPUs partition data. After the implementation, the interviewer asked many questions about optimizing it for the GPU.
```text
keys = [12, 7, 3, 8, 5, 10, 1] P = 3 hash(key) = key (for this example only)
partition of each key [0, 1, 0, 2, 2, 1, 1]
keys per partition [2, 3, 2]
offsets [0, 2, 5]
one valid output [12, 3, 7, 10, 1, 8, 5]
```
### Constraints and Clarifications
- Write the GPU code as CUDA C++ kernels plus the host code that launches them; library primitives may be used if you can explain what they do.
- Assume the keys are already in GPU memory and an output buffer of the same size is available.
- Not given: the key type, the hash function, the sizes of `n` and `P`, and whether order within a partition matters. Ask, or state your assumptions.
### Clarifying Questions
- What is the key type, and is the hash function given or mine to choose?
- How large are `n` and `P`? Is `P` small enough for one counter per partition to fit in a block's shared memory?
- Must keys keep their original relative order within each partition?
- Should the function return only the offsets, or also the rearranged keys and any values attached to them?
- Is this called once on a large array or many times, so that setup and launch costs matter?
### Part 1 — Partition the array on the GPU
Design the parallel algorithm, write the kernels and the host code, and state the work and the number of passes over global memory.
```hint What a thread cannot know alone
A thread can compute its own key's partition, but not where to write the key. Work out what information about the other keys the write position depends on, and how to compute it in parallel before anyone writes.
```
#### What This Part Should Cover
- A correct multi-pass algorithm and the reason each pass is needed
- What each thread does in each kernel, and how the passes are ordered on the GPU
- The exact definition of the offsets, empty partitions, and whether the result is stable
- Work, memory traffic, and how they depend on `n` and `P`
### Part 2 — Optimize it for the GPU
Explain what limits the speed of your implementation on a GPU and how you would make it faster. Expect several separate optimization questions about it.
```hint Look at the memory addresses
For each pass, look at how many threads update the same memory address at the same time, and whether neighboring threads read and write neighboring addresses.
```
#### What This Part Should Cover
- Contention on shared counters and how to reduce it
- Memory access patterns of the reading and writing passes
- Use of on-chip shared memory, work per thread, and how these choices depend on `P`
- How to measure an improvement and verify that the result is still correct
### What a Strong Answer Covers
- Clarifies the key type, sizes, stability and output format before writing kernels
- A correct parallel algorithm with the offsets defined precisely and every pass justified
- Working CUDA C++ with correct indexing, bounds checks, and no unnecessary host synchronization
- Optimizations that follow from the measured or reasoned bottlenecks, each with its cost
- Testing against a sequential reference, including skewed inputs and empty partitions
### Follow-up Questions
- `P` grows to millions, so one counter per partition no longer fits in shared memory. What changes?
- Almost all keys hash to one partition. What slows down, and how do you fix it?
- Each key carries a 64-byte value that must move with it. How do you move the values efficiently?
- If your counting pass and your placement pass each read every key from global memory, could one read be enough? If not, what can you still save?
Overview: A GPU programming question that asks you to hash-partition an array of keys into P partitions in CUDA C++, making each partition contiguous and returning each partition's starting offset. Follow-up questions test GPU optimization: atomic contention, memory coalescing, shared memory use and the number of passes over global memory.
Read the full NVIDIA Software Engineer interview experience this question came from