Rust SIMD on the GPU

At VectorWare, we are building the first GPU-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.

Parallelism below the thread

When we brought Rust threads to the GPU, we mapped each std::thread to a GPU warp. This let us run many concurrent threads on the GPU but did not use the parallel lanes within each thread/warp.

On the CPU, the abstraction for parallelism within a thread is SIMD. 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, eight f32 values and produces eight sums at once. This data parallelism is inside a single thread, below the level where the operating system schedules anything.

CPU threadSIMD op012N⋯SIMD lanesCPU thread

Rust's portable SIMD

Historically, writing SIMD in Rust meant reaching for the architecture-specific vendor intrinsics in core::arch, such as _mm256_add_ps on x86-64 or vaddq_f32 on 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's portable SIMD instead adds a layer of abstraction above these intrinsics. It provides a single generic type Simd that represents a vector of N elements of type T. A program writes its arithmetic, comparisons, reductions, and lane shuffles once against Simd and 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 in core rather than std and it does not even need the std support we brought to the GPU.

SIMT is SIMD

GPUs execute in a model NVIDIA calls SIMT, 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 is exactly what 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.

CPU thread012N⋯SIMD lanes≈GPU warp012N⋯warp lanes

For example, a Simd gives one i16 element 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.

CPUlet a: Simd = [1, 1, 1, ..., 1];let b: Simd = [2, 2, 2, ..., 2];let c = a + b;compiles tovpaddw %zmm2, %zmm1, %zmm0a0+b0lane 0a1+b1lane 1a2+b2lane 2a31+b31lane 31⋯println!("{c:?}"); GPUlet a: Simd = [1, 1, 1, ..., 1];let b: Simd = [2, 2, 2, ..., 2];let c = a + b;compiles toadd.s16 %rs3, %rs1, %rs2;a0+b0lane 0a1+b1lane 1a2+b2lane 2a31+b31lane 31⋯println!("{c:?}");

This new mapping completes the parallelism hierarchy from our earlier work. On the CPU, a thread contains SIMD lanes, and on the GPU our std::thread is a warp whose hardware lanes play the same role. In both cases, core::simd drives those lanes.

CPU⋯thread 0012N⋯thread 1012N⋯thread N012N⋯SIMD lanes≈GPU⋯warp 0012N⋯warp 1012N⋯warp N012N⋯warp lanes

A world first: core::simd on the GPU

As with our earlier posts, this is hard to show visually because the code is ordinary Rust. The same core::simd types 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 from main. It exercises the core features of the model: elementwise arithmetic, a comparison that produces a lane mask, a select driven by that mask, and a horizontal reduction across lanes.

#![feature(portable_simd)]

use core::simd::cmp::SimdPartialOrd; use core::simd::num::SimdFloat; use core::simd::{Select, Simd};

// Portable SIMD. This exact function also compiles and runs on the CPU, // where it lowers to x86-64, Arm, or scalar code depending on the target. fn relu_dot(a: Simdf32, 32>, b: Simdf32, 32>) -> f32 { // Elementwise multiply: 32 products computed at once. let products = a * b;

// Per-lane comparison produces a mask, one boolean per lane. let positive = products.simd_gt(Simd::splat(0.0));

// Keep the positive products, replace the rest with zero. let clamped = positive.select(products, Simd::splat(0.0));

// Horizontal add across all lanes down to a single scalar. clamped.reduce_sum()

fn main() { // Two 32-wide vectors, built with ordinary Rust. let a = Simd::f32, 32>::splat(2.0); let b = Simd::f32, 32>::from_array(std::array::from_fn(|i| i as f32 - 16.0));

// Elementwise ops, a comparison mask, a select, and a reduction: // all ordinary portable SIMD, all running on the GPU. let result = relu_dot(a, b);

// Printed from the GPU using our std support. println!("relu_dot = {result}"); The entry...

Rust SIMD on the GPU

Related Articles

(no title)

Scientists reverse brain aging, with a nasal spray

NHS apologises and admits Palantir have access to identifiable patient data

AI has torched the market for junior programmers

Al Vigier: Canada's AI strategy shouldn't include secret Palantir bills