Rust SIMD on the GPU

AtVectorWare, we are building the firstGPU-native software company. Today, we are excited to
announce that we can successfully use Rust’s portable SIMD
(core::simd) on the GPU. This
milestone marks a significant step towards our vision of enabling developers to write
complex, high-performance applications that leverage the full power of GPU hardware
using familiar Rust abstractions.

When webrought Rust threads to the GPU, we mapped eachstd::threadto a GPUwarp. This let us run many concurrent threads on the GPU but
did not use the parallellaneswithin each thread/warp.

On the CPU, the abstraction for parallelism within a thread isSIMD.A single instruction
operates on several data elements packed into a vector unit: where scalar code adds
two numbers, a SIMD add takes two vectors of, say, eightf32values and produces eight
sums at once. This data parallelism isinsidea single thread, below the level where the
operating system schedules anything.

Historically, writing SIMD in Rust meant reaching for the architecture-specific vendor
intrinsics incore::arch, such as_mm256_add_pson
x86-64 orvaddq_f32on
Arm. These intrinsics are specific to a single instruction set, so a program
that runs on more than one architecture needs a separate implementation for each.

Rust’sportable SIMDinstead adds a layer
of abstraction above these
intrinsics. It provides a single generic typeSimdthat represents a
vector ofNelements of typeT. A program writes its arithmetic, comparisons,
reductions, and lane shuffles once againstSimdand the compiler lowers them to whatever
vector instructions the target CPU has.

At VectorWare, we realized the GPU is just one more piece of vector hardware for
portable SIMD to target.As a bonus, portable SIMD lives incorerather thanstdand it does not even need thestdsupport we brought to the
GPU.

GPUs execute in a model NVIDIA callsSIMT, or Single
Instruction, Multiple Thread. A warp issues one instruction, and each of its 32 lanes runs
that instruction on its own data. One instruction operating on many data elements isexactlywhat SIMD means, and the per-lane addressing that SIMT adds does not change
that.A warp is a wide vector unit and a portable SIMD vector maps onto that unit directly.

For example, aSimdgives onei16element to each of the warp’s 32 lanes, and adding two such vectors compiles to a single warp
instruction in which every lane adds its element at once.

This new mapping completes the parallelism hierarchy from our earlier work. On the CPU, a
thread contains SIMD lanes, and on the GPUourstd::threadis a
warpwhose hardware lanes play the same role. In both cases,core::simddrives those lanes.

As with our earlier posts, this is hard to show visually because the code is ordinary
Rust. The samecore::simdtypes that lower to x86-64 SIMD on a laptop lower to warp
operations on the GPU, with no change to the source.

Here we define a small portable SIMD routine and call it frommain. It exercises
the core features of the model: elementwise arithmetic, a comparison that produces a
lane mask, aselectdriven by that mask, and a horizontal reduction across lanes.

The entry point is a normalfn mainwith no GPU-specific annotations. Our toolchain
compiles it to a GPU kernel, and the result is printed from the device using ourstdsupport.