Portable GPU Kernels

When writing code for the GPU, the best practice is to first try and use array programming, which can be executed on either the CPU or the GPU. This code is easy to test, and the tests can be run on any machine, without the GPU. If the code works correctly on the CPU, it is highly likely to also work on the GPU - as long as there are no errors. However, there will always be edge cases when your algorithm cannot be implemented in terms of array and broadcast operations, or it is highly inefficient to do so. In these cases, you may have to write the GPU kernel itself.

A GPU kernel is a self-contained program that can be run independently across many cores of a GPU. Traditionally, GPUs could compile shaders, which were able to calculate the colour of a single pixel on the screen. This shader would be run for every pixel to colour a final image. Kernels are just like these old shaders, but perform generic computation, instead of a graphics processing routine. Instead of iterating over each unit of work, we simply describe how to do a single piece of the work, given some identification of what work one should be doing.

Let’s take the example of writing a basic shader which performs vector addition. Luckily, even though CUDA is designed to be a C-like language, we are able to write all of our CUDA code in native Julia (with the help of CUDA.jl for compilation). As a starting point, let’s write the serial version as a basic for loop, describing the work we eventually want to run on the GPU:

function basic_vec_add!(c, a, b)
    for i in eachindex(c, a, b)
        @inbounds c[i] = a[i] + b[i]
    end
    return nothing
end

This code will add the vectors a and b together and store the result in the vector c. Our GPU kernel should make each core perform one of these inner loop additions. KernelAbstractions.jl provides each invocation with a global index, independently of whether the underlying device is NVIDIA or AMD:

@kernel function basic_vec_add_kernel!(c, a, b)
    i = @index(Global, Linear)
    if i <= length(c)
        @inbounds c[i] = a[i] + b[i]
    end
end

@index(Global, Linear) identifies the current unit of work across the complete launch. As with CUDA threads, every invocation runs independently.

We can compile and launch this kernel on whichever supported GPU is active with the following syntax:

a = to_gpu(rand(Float32, 128)); b = to_gpu(rand(Float32, 128)); c = similar(a);
backend = get_backend(c)
basic_vec_add_kernel!(backend, 128)(c, a, b; ndrange=length(c))
gpu_synchronize(c)
@show isapprox(c, a .+ b)

isapprox(c, a .+ b) = true

We can see that our kernel worked as expected. Note that indexing single elements with c[i] is perfectly fine inside a kernel - that code runs on the device itself, so it is not the slow, host-side scalar indexing warned about in the previous chapter. The first argument to basic_vec_add_kernel! selects the compiled back-end, the second is the workgroup size, and ndrange is the number of logical indices. KernelAbstractions maps these concepts onto the native execution model.

• Each GPU has a 3-dimensional grid of blocks.
• Each block contains a 3-dimensional group of threads.
• Each thread is executed independently.

In our case, we restricted the execution to a single dimension $x$, but this just makes the number of elements in the $y$ and $z$ dimensions equal to $1$. One important restriction in the CUDA programming model is making sure that a block has a maximum number of possible threads. This limit is usually $1024$ threads. As we are often dealing with larger arrays, we need to make use of multiple blocks to allow execution across more threads than this limit.

Looking at the figure above, we can visualise how a CUDA launch maps work onto an input array. KernelAbstractions performs the equivalent mapping for us: it computes the global index and simply ignores excess invocations in the bounds guard. We can wrap the kernel in a type-generic function:

function basic_vec_add!(c, a, b)
    n = length(c)
    @assert n == length(a) == length(b)
    n == 0 && return nothing

    backend = get_backend(c)
    basic_vec_add_kernel!(backend, 256)(c, a, b; ndrange=n)
    return nothing
end

The workgroup size of $256$ is a conservative starting point for both CUDA and AMDGPU devices. KernelAbstractions launches enough workgroups to cover ndrange, so we do not calculate a grid size manually. As in a CUDA kernel, the bounds guard is still needed because the logical length may not divide evenly by the workgroup size.

Kernels do not return values and must not allocate memory. This usually means implementing them as in-place operations on pre-allocated memory. If you want to provide a better API for users and abstract away the allocations, wrap the launch:

function basic_vec_add(a, b)
    c = similar(a)
    basic_vec_add!(c, a, b)
    return c
end

Additionally, if the memory is only needed temporarily, you can create it outside the function and release it later. If you are writing the critical loop section of your code, it is often better to write the main algorithm with pre-allocated caches and then create a wrapper (like the one above) to optionally use if you do not want the users of the code to manually manage their cache. This gives the option for re-using a cache in a hot-loop.