Quickstart

Terminology

Because CUDA is the most popular GPU programming environment, we can use it as a reference for defining terminology in KA. A workgroup is called a block in NVIDIA CUDA and designates a group of threads acting in parallel, preferably in lockstep. For the GPU, the workgroup size is typically around 256, while for the CPU, it is usually a multiple of the natural vector-width. An ndrange is called a grid in NVIDIA CUDA and designates the total number of work items. If using a workgroup of size 1 (non-parallel execution), the ndrange is the number of items to iterate over in a loop.

Like other GPU programming models, KernelAbstractions assumes that the work is roughly balanced across the work-items. A launch without a workgroup size gets one chosen under that assumption: on the CPU, each workgroup runs on one thread and is vectorized across its work-items, so the default splits a launch into about one workgroup per thread, as long as the workgroups stay large enough to vectorize. If the cost of a work-item varies a lot, e.g., with a triangular loop bound, pass a smaller workgroup size, so that threads that finish early can pick up more workgroups.

Writing your first kernel

Kernel functions are marked with the @kernel. Inside the @kernel macro you can use the kernel language. As an example, the mul2 kernel below will multiply each element of the array A by 2. It uses the @index macro to obtain the global linear index of the current work item.

using KernelAbstractions

@kernel function mul2_kernel(A)
  I = @index(Global)
  A[I] = 2 * A[I]
end

mul2_kernel (generic function with 4 methods)

Launching kernel on the host

You can construct a kernel for a specific backend by calling the kernel with mul2_kernel(CPU(), 16). The first argument is a backend of type KA.Backend, the second argument being the workgroup size. This returns a generated kernel executable that is then executed with the input argument A and the additional argument being a static ndrange.

dev = CPU()
A = ones(1024, 1024)
mul2_kernel(dev, 64)(A, ndrange=size(A))
synchronize(dev)
@assert all(A .== 2.0)

All kernels are launched asynchronously. The synchronize blocks the host until the kernel has completed on the backend.

Static workgroup size and ndrange

When the workgroup size and ndrange are known ahead of time, pass them to the kernel constructor to enable additional compile-time optimizations and avoid supplying them at every launch:

# workgroup size 32, ndrange size(A) — fixed for this kernel object
kernel = mul2_kernel(dev, 32, size(A))
kernel(A)  # ndrange inferred from construction
synchronize(dev)
@assert all(A .== 4)

See also Memcopy with static NDRange.

Opting out of the bounds check

The ndrange is padded to a multiple of the workgroup size, so a kernel checks for every work-item whether it is inside the ndrange before running the body, and the body only runs for the work-items that are. A kernel can opt out of the check with unsafe_indices = true. Its body then also runs for the padded work-items, so it has to test the index itself before touching memory:

@kernel unsafe_indices = true function mul2_unsafe_kernel(A)
    g = @index(Group, Linear)
    l = @index(Local, Linear)
    I = (g - 1) * prod(@groupsize()) + l
    if I <= length(A)
        A[I] = 2 * A[I]
    end
end

mul2_unsafe_kernel(dev, 64)(A, ndrange = length(A))
synchronize(dev)
@assert all(A .== 8)

Don't use @index(Global) in any form in such a kernel. Compute the global index from @index(Group), @index(Local) and @groupsize instead, as above, and compare it with the size of the range.

Launching kernel on the backend

To launch the kernel on a backend (e.g., CUDABackend(), ROCBackend(), oneAPIBackend(), MetalBackend()), we generate the kernel for this backend.

First, we initialize the array using the Array constructor of the chosen backend with

using CUDA: CuArray
A = CuArray(ones(1024, 1024))

using AMDGPU: ROCArray
A = ROCArray(ones(1024, 1024))

using oneAPI: oneArray
A = oneArray(ones(1024, 1024))

using Metal: MtlArray
A = MtlArray(ones(Float32, 1024, 1024))

The kernel generation and execution are then

backend = get_backend(A)
mul2_kernel(backend, 64)(A, ndrange=size(A))
synchronize(backend)
@assert all(A .== 2)

Synchronization

The code around KA may heavily rely on GPUArrays, for example, to initialize variables.

function mymul(A)
    A .= 1.0
    backend = get_backend(A)
    ev = mul2_kernel(backend, 64)(A, ndrange=size(A))
    synchronize(backend)
    @assert all(A .== 2.0)
end

mymul(A)

function mymul(A, B)
    A .= 1.0
    B .= 3.0
    backend = get_backend(A)
    @assert get_backend(B) == backend
    mul2_kernel(backend, 64)(A, ndrange=size(A))
    mul2_kernel(backend, 64)(B, ndrange=size(B))
    synchronize(backend)
    @assert all(A .+ B .== 8.0)
end

mymul(A, ones(size(A)))

Using task programming to launch kernels in parallel

As shown in the Synchronization section above, multiple kernels can be enqueued on the same backend before a single synchronize call. The same pattern extends to Julia's task-based parallelism: launch kernels from tasks when you want to overlap kernel execution with other asynchronous host work, or with each other.

Backends may give each Julia task its own queue, so kernels launched from different tasks can run concurrently, but are not ordered with respect to each other. Use KernelAbstractions.@spawn in place of Threads.@spawn to launch kernels from a task. It behaves like Threads.@spawn, and additionally guarantee that kernels run after everything the spawning task had already queued, and that once wait(task) or fetch(task) returns, the kernel results are ready to use:

function exchange_and_compute!(backend, A, B)
    recv = KernelAbstractions.@spawn backend begin
        mul2_kernel(backend, 64)(A, ndrange=length(A))
    end
    send = KernelAbstractions.@spawn backend begin
        mul2_kernel(backend, 64)(B, ndrange=length(B))
    end
    wait(recv)
    wait(send)
end

Waiting on a backend, whether with synchronize or at the end of a spawned task, yields to the Julia scheduler, so other tasks keep making progress while a kernel runs.

Which device a task runs on

A task started with plain Threads.@spawn runs on an implementation defined device for chosen backend. KernelAbstractions.@spawn selects the device explicitly: by default the one active in the spawning task, or the one named by device, a 1-based index into 1:ndevices(backend):

function compute_on_both!(backend, A, B)
    here = KernelAbstractions.@spawn backend begin
        mul2_kernel(backend, 64)(A, ndrange=length(A))
    end
    there = KernelAbstractions.@spawn backend device=2 begin
        mul2_kernel(backend, 64)(B, ndrange=length(B))
    end
    wait(here)
    wait(there)
end

The ordering guarantee holds across the switch: the second task's work on device 2 is still ordered after what the spawning task had queued on its own device.

Prefer device= over calling device! inside the body. device! is not a synchronization point — work queued after it is ordered neither against the spawning task nor against what the body queued before the switch. If you do switch by hand, order it with record_event and wait_event, which apply to the device that is active when each is called:

event = KernelAbstractions.record_event(backend)   # captures work on the current device
KernelAbstractions.device!(backend, 2)
KernelAbstractions.wait_event(backend, event)      # device 2 waits for that work

A plain synchronize before the device! works too, at the cost of blocking the task until the first device is idle.

A full MPI example that overlaps communication with device copies is in examples/mpi.jl.