GPU kernels

A kernel fn runs once per thread on the GPU. On macOS native code launches it on Metal; elsewhere, or with SSPUR_GPU=0, it runs on the CPU. Every tier gives the same results and the same traps, so kernels are tested like any other code.

A first kernel

kernel fn saxpy(a: F32, x: &[F32], y: &mut [F32]) @grid(y.len, 256)
  pre x.len == y.len
= y[gid] := a * x[gid] + y[gid]

• @grid(y.len, 256) launches y.len threads in groups of 256. gid is the thread’s index.
• Parameters are scalars or slices. &[F32] is read-only and &mut [F32] writable; writing a read-only slice is E_KERNEL_WRITE.
• pre is checked once, before the launch.
• The checker proves that no two threads write the same element. Here every thread writes y[gid]; an index that isn’t gid plus a uniform offset is E_KERNEL_RACE.

F32 arithmetic rounds like IEEE single precision in every tier, so results are the same on the GPU and the CPU.

Shared memory and barriers

A reduction sums 256 values per group in shared memory:

kernel fn block_sums(x: &[F32], out: &mut [F32]) @grid(x.len, 256)
= do
  sh = shared[F32](256)
  sh[lid] := x[gid]
  barrier()
  for k in 0..8
    s = 128.shr(k)
    if lid < s then sh[lid] := sh[lid] + sh[lid + s]
    barrier()
  if lid == 0 then out[group_id] := sh[0]

shared[F32](256) is a zeroed array per group, and lid is the thread’s index within its group. barrier() waits for the whole group, so it may only sit under conditions that are the same for every thread in the group (E_KERNEL_BARRIER otherwise). Between barriers, each thread writes its own sh[lid] and reads sh[lid + s] only under if lid < s, which the race check accepts because the ranges can’t overlap.

Launching from the host

fn sum_on_gpu(n: Int) -> F64 ! dev
= do
  x = dev_f32((0..n).map(i => (i % 10).to_f64))
  parts = dev_f32((0..n / 256).map(i => 0.0))
  block_sums(&x, &mut parts)
  parts.to_list.sum

fn scaled(a: F64) -> List[F64] ! dev
= do
  var ys = [1.0, 2.0, 3.0]
  saxpy(a, &[10.0, 20.0, 30.0], &mut ys)
  ys

Calling a kernel launches it and performs the dev effect. &mut ys copies a list in and the result back. dev_f32(xs) copies a list to device memory once, as a DevBuf that kernels update in place, and parts.to_list copies it back. Launches on DevBufs are queued and finish at the next to_list or gpu_sync().

$ sspur run saxpy.ssp
saxpy: [6.0, 12.0, 18.0]
sum: 294900.0
$ sspur run --interp saxpy.ssp
saxpy: [6.0, 12.0, 18.0]
sum: 294900.0
$ sspur test saxpy.ssp
1 passed, 0 failed

Set SSPUR_GPU_TRACE=1 to see each launch, the device and the thread count.

Emitting kernel source

sspur gpu prints the kernels as Metal or OpenCL C, or compiles SPIR-V and PTX with a clang that has those backends:

$ sspur gpu saxpy.ssp --emit metal
...
kernel void k_saxpy(constant float& p0 [[buffer(0)]], device const float* p1 [[buffer(1)]], constant long& l1 [[buffer(2)]], device float* p2 [[buffer(3)]], constant long& l2 [[buffer(4)]], constant long& ss_n [[buffer(5)]], device atomic_uint* ss_flag [[buffer(6)]], device atomic_uint* ss_mask [[buffer(7)]], uint ss_t [[thread_position_in_grid]]) {
  long gid = (long)ss_t;
  if (gid >= ss_n) return;
...

Every array access is bounds-checked and every integer operation overflow-checked on the device too. When a thread would trap, or meets a subnormal F32, the launch reruns on the CPU to give exactly the interpreter’s result. Atomics, 2D grids (@grid2) and helper functions are covered in the language reference and in ADR 0020 and ADR 0023.