repository·main·Indexed 19 days ago
https://github.com/juliagpu/kernelabstractions.jlA Julia package providing a unified interface for writing high-performance, heterogeneous kernels that target multiple execution backends, including NVIDIA CUDA (via CUDA.jl), AMD ROCm (via AMDGPU.jl), Intel oneAPI (via oneAPI.jl), Apple Metal (via Metal.jl), and CPU. It offers a minimal abstraction layer with kernel language macros like @kernel and @index, memory management tools, and integration with Atomix.jl for thread-safe atomic operations.
KernelAbstractions.jl (KA) is a performance-portable library that allows you to write GPU-like kernels targeting different execution backends. It emulates GPU semantics to provide a consistent programming model across both CPU and GPU hardware.
Key Concepts:
@synchronize might be ignored on the CPU, while others like @localmem might be swapped for CPU-equivalent structures (e.g., MVector). This may result in the CPU performing extra work to maintain the programming model, though it remains fast.When implementing a backend for KernelAbstractions.jl, the KernelAbstractions.synchronize function must be cooperative.
This means the implementation cannot block inside an external library (such as a C or Fortran library). Instead, it must implement a cooperative wait that yields the current task, returning the scheduling slice to the Julia runtime. This cooperative behavior is critical for allowing the overlapping of communication and computation, particularly when using MPI.
In KernelAbstractions (specifically since v0.9.34), the @synchronize construct requires convergent execution.
The Issue:
Most GPU implementations execute workgroups on static blocks. If a kernel's ndrange is smaller than the static block size (e.g., ndrange=(32, 30) on a (32, 32) block), KA previously inserted dynamic bounds-checks. This caused implicit divergent execution when calling @synchronize, which can lead to mis-compilations on backends like OpenCL.
The Solution:
KA now lowers kernels to ensure @synchronize is called in a way that respects convergence. If you need to avoid the implicit bounds-checking behavior that might cause divergence, you can use the unsafe_indices=true flag on your kernel, but you must manually derive your global indices using @index(Group) and @index(Local) instead of relying on @index(Global).
# Example of a kernel using unsafe_indices to avoid implicit bounds-check divergence
@kernel unsafe_indices=true function localmem(A)
N = @uniform prod(@groupsize())
gI = @index(Group, Linear)
i = @index(Local, Linear)
lmem = @localmem Int (N,)
lmem[i] = i
@synchronize
I = (gI - 1) * N + i
if i <= N && I <= length(A)
A[I] = lmem[N - i + 1]
end
endKernelAbstractions.jl supports several hardware backends through specific driver packages. Depending on your hardware, you can target:
CUDA.jl)AMDGPU.jl)oneAPI.jl)Metal.jl)To use a specific backend, you must ensure the corresponding driver package is installed and available in your Julia environment.
You can enqueue multiple kernels on the same backend before calling synchronize. To overlap kernel execution with other asynchronous host work, use Julia's task-based parallelism via Threads.@spawn.
On GPU backends, synchronize(backend) is cooperative: it yields to the Julia scheduler rather than blocking the entire thread, allowing other Julia tasks to make progress while the kernel runs on the hardware.
function exchange_and_compute!(backend, A, B)
recv = Threads.@spawn begin
mul2_kernel(backend, 64)(A, ndrange=length(A))
synchronize(backend) # cooperative on GPU backends
end
send = Threads.@spawn begin
mul2_kernel(backend, 64)(B, ndrange=length(B))
synchronize(backend)
end
wait(recv)
wait(send)
endUse the @localmem macro to declare storage shared by all work items in a workgroup.
Note: Currently, only static local memory is supported. The allocation size must be known at compile time (e.g., @localmem Int (32,) or @localmem Int (N,) where N = prod(@groupsize()) and the workgroup size is fixed during kernel construction).
If different work items perform reads and writes to the same local memory, you must separate these operations with the @synchronize macro to avoid race conditions.
@private: Declares per-work-item storage that survives across @synchronize statements.@uniform: Evaluates an expression outside the work-item scope so it can be reused across @synchronize statements.@synchronize, use an MArray instead.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)When using KernelAbstractions instead of direct backend libraries like CUDA.jl or AMDGPU.jl, note the following semantic differences:
ndrange.nothing.KernelAbstractions.jl uses terminology inspired by CUDA to describe parallel execution models:
To profile the performance of your kernels (such as measuring memory throughput and SM utilization), you can use the nv-nsight-cu-cli tool. This is particularly useful for analyzing the 'Speed of Light' metrics on NVIDIA GPUs.
Run the following command to profile the performance example provided in the repository:
nv-nsight-cu-cli --nvtx --profile-from-start=off --section=SpeedOfLight --section=julia --project=examples examples/performance.jlTo run kernels on a GPU, you must use arrays provided by the specific backend (e.g., CuArray for CUDA, ROCArray for AMD, oneArray for oneAPI, or MtlArray for Metal).
You can use get_backend(A) to automatically retrieve the correct backend from a device array. The kernel is then constructed with that backend and launched using the ndrange keyword.
using CUDA: CuArray
A = CuArray(ones(1024, 1024))
backend = get_backend(A)
# Launching on the backend
mul2_kernel(backend, 64)(A, ndrange=size(A))
synchronize(backend)
@assert all(A .== 2)Kernels are constructed by calling the kernel function on a backend and then launched with an ndrange. There are three primary ways to configure the launch:
ndrange (and optionally workgroupsize) at launch time.ndrange at launch.Always call synchronize(backend) before reading results on the host. You can obtain the backend from an array using get_backend(array).
# 1. Dynamic sizes
kernel = my_kernel(backend)
kernel(A, ndrange=size(A))
# 2. Static workgroup size
kernel = my_kernel(backend, 256)
kernel(A, ndrange=size(A))
# 3. Static workgroup size and ndrange
kernel = my_kernel(backend, 32, size(A))
kernel(A)