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.

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.

Launching kernel on the backend

To launch the kernel on a backend-supported backend isa(backend, KA.GPU) (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.