Kernels

The KernelAbstractions.jl package provides a macro-based dialect that hides the intricacies of vendor-specific GPU programming. It allows one to write hardware-agnostic kernels that can be instantiated and launched for different device backends without modifying the high-level code nor sacrificing performance.

In the following, we show how to write and launch kernels on various backends. We also explain the concept of a Launcher in Chmy.jl, that complements the default kernel launching, allowing us to hide the latency between the bulk of the computations and boundary conditions or MPI communications.

Writing Kernels

This section highlights some important features of KernelAbstractions.jl that are essential for understanding the high-level abstraction of the kernel concept that is used throughout our package. As it barely serves for illustrative purposes, for more specific examples, please refer to their documentation.

using KernelAbstractions

# Define a kernel that performs element-wise operations on A
@kernel function mul2!(A)
    # use @index macro to obtain the global Cartesian index of the current work item.
    I = @index(Global, Cartesian)
    A[I] *= 2
end

The kernel mul2! being defined using the @kernel macro, we can launch it on the desired backend to perform the element-wise operations on host.

# Define array and work group size
A       = ones(1024, 1024)
backend = get_backend(A) # CPU

# Launch kernel and explicitly synchronize
kernel = mul2!(backend)
kernel(A, ndrange=size(A))
KernelAbstractions.synchronize(backend)

# Result assertion
@assert(all(A .== 2.0) == true)

To launch the kernel on GPU devices, one could simply define A as CuArray, ROCArray or oneArray as detailed in the section "launching kernel on the backend". More fine-grained memory access is available using the @index macro as described here.

Thread Indexing

Thread indexing is essential for memory usage on GPU devices; however, it can quickly become cumbersome to figure out the thread index, especially when working with multi-dimensional grids of multi-dimensional blocks of threads. The performance of kernels can also depend significantly on access patterns.

In the example above, we saw the usage of I = @index(Global, Cartesian), which retrieves the global index of threads for the two-dimensional array A. Such powerful macros are provided by KernelAbstractions.jl for conveniently retrieving the desired index of threads.

The following table is non-exhaustive and provides a reference of commonly used terminology. Here, KernelAbstractions.@index is used for index retrieval, and KernelAbstractions.@groupsize is used for obtaining the dimensions of blocks of threads.

KernelAbstractions	CPU	CUDA	AMDGPU
@index(Local, Linear)	mod(i, g)	threadIdx().x	workitemIdx().x
@index(Local, Cartesian)[2]		threadIdx().y	workitemIdx().y
@index(Local, Cartesian)[3]		threadIdx().z	workitemIdx().z
@index(Group, Linear)	i ÷ g	blockIdx().x	workgroupIdx().x
@index(Group, Cartesian)[2]		blockIdx().y	workgroupIdx().y
@groupsize()[3]		blockDim().z	workgroupDim().z
@index(Global, Linear)	i	global index computation needed	global index computation needed
@index(Global, Cartesian)[2]		global index computation needed	global index computation needed
@index(Global, NTuple)		global index computation needed	global index computation needed

The @index(Global, NTuple) returns a NTuple object, allowing more fine-grained memory control over the allocated arrays.

@kernel function memcpy!(a, b)
    i, j = @index(Global, NTuple)
    @inbounds a[i, j] = b[i, j]
end

A tuple can be splatted with ... Julia operator when used to avoid explicitly using i, j indices.

@kernel function splatting_memcpy!(a, b)
    I = @index(Global, NTuple)
    @inbounds a[I...] = b[I...]
end

Kernel Launcher

In Chmy.jl, the KernelLaunch module is designed to provide handy utilities for performing different grid operations on selected data entries of Fields that are involved at each kernel launch, in which the grid geometry underneath is also taken into account.

Followingly, we define a kernel launcher associated with an UniformGrid object, supporting CUDA backend.

# Define backend and geometry
arch = Arch(CUDABackend())
grid = UniformGrid(arch; origin=(-1, -1), extent=(2, 2), dims=(126, 126))

# Define launcher
launch = Launcher(arch, grid)

We also have two kernel functions compute_q! and update_C! defined, which shall update the fields q and C using grid operators (see section Grid Operators) ∂x, ∂y, divg that are anchored on some grid g accordingly.

@kernel inbounds = true function compute_q!(q, C, χ, g::StructuredGrid, O)
    I = @index(Global, Cartesian)
    I = I + O
    q.x[I] = -χ * ∂x(C, g, I)
    q.y[I] = -χ * ∂y(C, g, I)
end

@kernel inbounds = true function update_C!(C, q, Δt, g::StructuredGrid, O)
    I = @index(Global, Cartesian)
    I = I + O
    C[I] -= Δt * divg(q, g, I)
end

To spawn the kernel, we invoke the launcher using the launch function to perform the field update at each physical timestep, and specify desired boundary conditions for involved fields in the kernel.

# Define physics, numerics, geometry ...
for it in 1:nt
    # without boundary conditions
    launch(arch, grid, compute_q! => (q, C, χ, grid))

    # with Neumann boundary conditions and MPI exchange
    launch(arch, grid, update_C! => (C, q, Δt, grid); bc=batch(grid, C => Neumann(); exchange=C))
end

Overlap and implicit synchronization

With outer_width, the launcher computes the inner points on the task-local stream, and the boundary regions, their boundary conditions and the MPI exchange on the streams of worker tasks. CUDA.jl (≥ 6.3) and AMDGPU.jl (≥ 2.11) order every use of an array on a stream after its last use on another stream, which would serialize these. The launcher therefore waits for the work pending on the task-local stream and suspends this implicit synchronization for the arrays passed to the kernel while it runs. Work on these arrays issued from other tasks must be synchronized by the caller before the launch. On AMDGPU.jl < 2.11 the implicit synchronization cannot be suspended: the results are correct, but computation and communication do not overlap.

Overlap on AMD GPUs

HIP maps the streams of a process to hardware queues in creation order, up to GPU_MAX_HW_QUEUES (default 4) per priority. With ROCm 7.2 on MI300A and MI250X, a kernel on the queue four places after a busy one (e.g. the 5th queue while the first dispatches a large kernel) starts late. With the defaults, one worker of the overlapped launcher gets such a queue, which makes short launches (around a millisecond) up to 15% slower than with the plain launcher. Setting GPU_MAX_HW_QUEUES=2 in the environment of the run avoids this. Binding each MPI rank to the CPU cores of its GPU's NUMA domain (e.g. srun --cpu-bind=map_ldom:...) shortens the launch latency. On MI250X, the boundary kernels slow down the inner kernel that runs concurrently, so the plain launcher (without outer_width) is usually faster. On MI300A, compare both launchers for the problem at hand.