Writing kernels

These kernel language constructs are intended to be used inside @kernel functions. They are not valid in ordinary Julia code.

Constant arguments

Kernel functions allow input arguments to be marked with the @Const macro. It informs the compiler that the memory accessed through that argument will not be written to as part of the kernel, and that it does not alias any other memory in the kernel. If you are used to CUDA C, this is similar to const restrict.

using KernelAbstractions

@kernel function saxpy!(a, @Const(X), Y)
    I = @index(Global)
    @inbounds Y[I] = a * X[I] + Y[I]
end

a = 2.0
X = collect(1.0:8.0)
Y = fill(1.0, 8)
backend = CPU()
saxpy!(backend, 8, size(Y))(a, X, Y)
synchronize(backend)
Y

# output

8-element Vector{Float64}:
  3.0
  5.0
  7.0
  9.0
 11.0
 13.0
 15.0
 17.0

Indexing

The @index macro returns the index of the current work item. Choose a granularity and an optional kind:

using KernelAbstractions

@kernel function fill_diagonal!(A, val)
    I = @index(Global, Cartesian)
    if I[1] == I[2]
        @inbounds A[I] = val
    end
end

A = collect(reshape(1.0:16.0, 4, 4))
backend = CPU()
fill_diagonal!(backend, 4, size(A))(A, 42)
synchronize(backend)
A

# output

4×4 Matrix{Float64}:
 42.0   5.0   9.0  13.0
  2.0  42.0  10.0  14.0
  3.0   7.0  42.0  15.0
  4.0   8.0  12.0  42.0

using KernelAbstractions

@kernel function linear_example!(A)
    I = @index(Global, Linear)   # 1, 2, 3, ...
    g = @index(Group, Linear)    # workgroup id
    l = @index(Local, Linear)    # work item within workgroup
    @inbounds A[I] = g + l
end

A = collect(1.0:16.0)
backend = CPU()
linear_example!(backend, 4, size(A))(A)
synchronize(backend)
A

# output

16-element Vector{Float64}:
 2.0
 3.0
 4.0
 5.0
 3.0
 4.0
 5.0
 6.0
 4.0
 5.0
 6.0
 7.0
 5.0
 6.0
 7.0
 8.0

Inside a kernel, @groupsize and @ndrange query the launch configuration:

using KernelAbstractions

@kernel function scale!(A, factor)
    N = @uniform prod(@groupsize())
    I = @index(Global, Linear)
    lmem = @localmem Float32 (N,)
    i = @index(Local, Linear)
    lmem[i] = factor
    @synchronize()
    @inbounds A[I] = A[I] * lmem[i]
end

A = collect(1.0:16.0)
backend = CPU()
scale!(backend, 8, size(A))(A, 2)
synchronize(backend)
A

# output

16-element Vector{Float64}:
  2.0
  4.0
  6.0
  8.0
 10.0
 12.0
 14.0
 16.0
 18.0
 20.0
 22.0
 24.0
 26.0
 28.0
 30.0
 32.0

Local memory, synchronization, and private memory

@localmem declares storage shared by all work items in a workgroup. Only static local memory is supported at the moment: the allocation size must be known at compile time (for example @localmem Int (32,) or @localmem Int (N,) where N = prod(@groupsize()) and the workgroup size is fixed when the kernel is constructed). Reads and writes must be separated by @synchronize if they are performed by different work items:

using KernelAbstractions

@kernel function reverse_block!(A)
    I = @index(Global, Linear)
    i = @index(Local, Linear)
    N = @uniform prod(@groupsize())
    buf = @localmem Int (N,)
    buf[i] = i
    @synchronize()
    @inbounds A[I] = buf[N - i + 1]
end

A = collect(1.0:16.0)
backend = CPU()
reverse_block!(backend, 8, size(A))(A)
synchronize(backend)
A

# output

16-element Vector{Float64}:
 8.0
 7.0
 6.0
 5.0
 4.0
 3.0
 2.0
 1.0
 8.0
 7.0
 6.0
 5.0
 4.0
 3.0
 2.0
 1.0

@private declares per-work-item storage that survives across @synchronize statements, and @uniform evaluates an expression outside the work-item scope so it can be reused across @synchronize statements. For scratch storage that does not need to survive across @synchronize, an MArray can be used instead.

Launching kernels

Construct a kernel by calling the kernel function on a backend and optional static sizes, then launch it with ndrange:

# dynamic sizes — supply ndrange (and optionally workgroupsize) at launch
kernel = my_kernel(backend)
kernel(A, ndrange=size(A))

# static workgroup size
kernel = my_kernel(backend, 256)
kernel(A, ndrange=size(A))

# static workgroup size and ndrange — fewer runtime checks, specialized per size
kernel = my_kernel(backend, 32, size(A))
kernel(A)

Index ranges

Each entry of ndrange is either an extent (indices 1:n) or a range of indices, so a kernel can iterate over a region whose indices do not start at 1. ndrange can also be given as a single range or as a CartesianIndices:

# static ndrange over the indices -2:N+3 along x and 0:M+1 along y
kernel = my_kernel(backend, (16, 16), (-2:N+3, 0:M+1))
kernel(A)

# dynamic ndrange
kernel = my_kernel(backend, (16, 16))
kernel(A, ndrange=(-2:N+3, 0:M+1))
kernel(A, ndrange=CartesianIndices(A))   # e.g. for an OffsetArray

Inside the kernel @index(Global, Cartesian) and @index(Global, NTuple) return the shifted indices, @index(Global, Linear) counts the indices of the region from 1 in column-major order, and @ndrange() returns the extents.

Obtain the backend from an array with get_backend and always call synchronize before reading results on the host. See the Quickstart for a full walkthrough and the Examples section of the manual for larger patterns.

Loops without a kernel

A kernel whose body is a loop over the indices of an array needs none of the kernel language beyond the index itself. foreach_index launches such a loop directly, with one work item per index of eachindex(y, x):

function scale!(y, x)
    foreach_index(y, x) do i
        @inbounds y[i] = 2 * x[i] + 1
    end
    return y
end

The body is an ordinary Julia function, and it receives the index a for i in eachindex(y, x) loop would: a linear index if the arrays have IndexLinear style, a CartesianIndex otherwise. Passing every array the body indexes makes sure the index is valid for each, and that they all live on the same backend. The launch is asynchronous like any other, and bounds checks are not elided.

To loop over other indices than those of an array, give the backend and the index space, a range or a CartesianIndices like an ndrange. The body then receives these indices as they are, e.g. those of the interior of a matrix:

function smooth!(B, A)
    interior = CartesianIndices((2:size(A, 1)-1, 2:size(A, 2)-1))
    foreach_index(get_backend(A), interior) do I
        i, j = Tuple(I)
        @inbounds B[I] = (A[i-1, j] + A[i+1, j] + A[i, j-1] + A[i, j+1]) / 4
    end
    return B
end

Because the body becomes a kernel, every value it captures must have a known type — which is why the examples wrap the loop in a function, and why the body must not assign to a captured variable. Write the kernel out with @kernel when it needs more than an index: workgroup-level indices, local memory, or synchronization.