KernelInterface

KernelInterface

The KernelInterface (or KI) module defines the API interface for backends to define various lower-level device and host-side functionality. The KI interface is used to define the higher-level device-side functionality in KernelAbstractions.

Both provide APIs for host and device-side functionality, but KI focuses on lower-level functionality that is shared amongst backends, while KernelAbstractions provides higher-level functionality such as writing kernels that work on arrays with an arbitrary number of dimensions, or convenience functions like allocating arrays on a backend.

Backend

Abstract supertype for all KernelInterface backends. Backends subtype it directly.

A backend value identifies a backend and its configuration (e.g. compiler options). The device and the queue that operations use are task-local: each task selects its active device with device!. Host-side queries answer for the calling task's active device, and work is queued on the calling task's queue of that device. Use get_backend to obtain the backend of an array and allocate to create storage on a backend.

Example

backend = get_backend(A)
kernel = my_kernel(backend, 256)
kernel(A, ndrange=length(A))
synchronize(backend)

get_backend(A::AbstractArray)::Backend

Get a Backend instance suitable for array A.

get_local_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}

The 1-based index of the work-item within its work-group, as integers of type T.

T is a fixed-width integer type of at most 64 bits (e.g. Int32 or UInt64); the result is the exact value modulo T, as if computed with x % T.

@device_override get_local_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T}

get_group_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}

The 1-based index of the work-group within the launch, as integers of type T.

See get_local_id for the supported types T.

@device_override get_group_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T}

get_local_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}

The number of work-items in a work-group, as integers of type T.

See get_local_id for the supported types T.

@device_override get_local_size(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T}

get_num_groups([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}

The number of work-groups in the launch, as integers of type T.

See get_local_id for the supported types T.

@device_override get_num_groups(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T}

get_global_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}

The 1-based index of the work-item within the launch, as integers of type T: (get_group_id(T) - 1) * get_local_size(T) + get_local_id(T) per dimension.

See get_local_id for the supported types T.

@device_override get_global_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T}

get_global_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T}

The number of work-items in the launch, as integers of type T: get_local_size(T) * get_num_groups(T) per dimension. For an ndrange launch, this is the ndrange padded to whole work-groups.

See get_local_id for the supported types T.

@device_override get_global_size(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T}

get_sub_group_size([::Type{T}=Int])::T

The number of work-items in the sub-group, at most the sub-group width (get_max_sub_group_size). Which sub-groups have fewer work-items than the width is unspecified: when the work-group size isn't a multiple of the width, there can be more than one, e.g. one per row of a multi-dimensional work-group.

See get_local_id for the supported types T.

@device_override get_sub_group_size(::Type{T})::T where {T}

get_max_sub_group_size([::Type{T}=Int])::T

The sub-group width, sub_group_size(backend) on the host.

See get_local_id for the supported types T.

@device_override get_max_sub_group_size(::Type{T})::T where {T}

get_num_sub_groups([::Type{T}=Int])::T

The number of sub-groups in the work-group. It is at least cld(prod(get_local_size()), get_max_sub_group_size()), but can be larger, since more than one sub-group can be partial. Size storage for a value per sub-group for up to one sub-group per work-item.

See get_local_id for the supported types T.

@device_override get_num_sub_groups(::Type{T})::T where {T}

get_sub_group_id([::Type{T}=Int])::T

The 1-based index of the sub-group within the work-group, between 1 and get_num_sub_groups. How it relates to get_local_id is unspecified.

See get_local_id for the supported types T.

@device_override get_sub_group_id(::Type{T})::T where {T}

get_sub_group_local_id([::Type{T}=Int])::T

The 1-based index of the work-item within its sub-group (its lane), between 1 and get_sub_group_size. It doesn't depend on which work-items of the sub-group are active, e.g. in a divergent branch.

See get_local_id for the supported types T.

@device_override get_sub_group_local_id(::Type{T})::T where {T}

barrier()

Wait until all work-items of the work-group have reached the barrier. Afterwards, the writes to global and local memory that each work-item made before the barrier are visible to all work-items of the work-group.

This does not order memory between work-groups.

All work-items of a work-group have to reach the same barrier() (not in a divergent branch).

@device_override barrier()

sub_group_barrier()

Like barrier, for the work-items of a sub-group: wait until all work-items of the sub-group have reached the barrier, and make their writes to global and local memory before it visible to the sub-group.

All work-items of a sub-group have to reach the same sub_group_barrier().

@device_override sub_group_barrier()

localmemory(::Type{T}, dims)

Declare an array of element type T and size dims in memory that is local to a work-group. dims has to be known at compile time.

Every call site of localmemory in a kernel has its own memory, shared by all work-items of a work-group. It is uninitialized, and lives until the work-group finishes. Executing the same call site again, e.g. in a loop, returns the same memory. A function containing a call that is itself called from several places may get the same memory at each of them, or different memory, depending on whether it is inlined: don't rely on either. Use barrier to make writes visible to the other work-items.

@device_override localmemory(::Type{T}, ::Val{Dims}) where {T, Dims}

shfl_down(val::T, offset::Integer)::T

Return val of the work-item offset lanes further in the sub-group, i.e. with get_sub_group_local_id equal to get_sub_group_local_id() + offset. When there is no such work-item, the result is an unspecified value (of type T).

All work-items of the sub-group have to execute shfl_down together (not in a divergent branch), with the same offset.

shfl_down exchanges values, not memory: it is not a memory fence.

@device_override shfl_down(val::T, offset::Integer) where T

_print(args...)

Print args from a kernel; the backend hook behind KernelAbstractions.@print.

@device_override _print(args...)

The generic fallback prints on the host, which keeps CPU backends working. Val arguments are unwrapped, since KernelAbstractions.@print uses them to pass literal strings through to backends that require compile-time format strings.

allocate(::Backend, Type, dims...; unified=false)::AbstractArray

Allocate an uninitialized array on the active device of the backend. unified=true allocates unified memory, accessible from the host and the device without explicit copies, if the backend supports it and throws otherwise. Use supports_unified to determine whether it is supported by a backend.

zeros(::Backend, Type, dims...; unified=false)::AbstractArray

Allocate an array with allocate and fill it with zeros.

This is generic: backends implement allocate (and fill! for their array type).

ones(::Backend, Type, dims...; unified=false)::AbstractArray

Allocate an array with allocate and fill it with ones.

This is generic: backends implement allocate (and fill! for their array type).

copyto!(::Backend, dest::AbstractArray, src::AbstractArray)::typeof(dest)

Copy the elements of src to dest, ordered with respect to the other work on the calling task's queue: after work queued before the copy, and before work queued after it. Returns dest.

Either array can be a host array or an array of backend. dest and src must have the same length, otherwise an ArgumentError is thrown. Backends only have to support dense (contiguous) arrays with the same element type.

The copy may be asynchronous with respect to the host, but doesn't have to be: it can also block until it has completed. For a simple, synchronous copy, use Base.copyto!.

arr = zeros(64)
GC.@preserve arr begin
    copyto!(backend, arr, ...)
    # other operations
    synchronize(backend)
end

pagelock!(::Backend, dest::AbstractArray)::Union{Nothing, Missing}

Pagelock (pin) a host memory buffer for a backend device. This may be necessary for copyto! to perform asynchronously with respect to the host.

This function returns nothing, or missing if not implemented.

unsafe_free!(x::AbstractArray)

Release the memory of an array for reuse by future allocations, reducing pressure on the allocator. The array may not be used afterwards.

This is a hint: releasing the memory is allowed to do nothing.

synchronize(::Backend)

Block the calling task until all work it has queued on the active device of backend has completed.

record_event(backend::Backend)

Capture the work the calling task has queued on backend's active device before this call, and return a handle for wait_event. Recording need not wait for that work to complete. The handle is only meant for wait_event.

wait_event(backend::Backend, event)

Order the work the calling task subsequently queues on backend's active device after the work captured by event, which was returned by record_event. Returning does not mean that the captured work has completed.

The wait applies to the queue of the device that is active when wait_event is called; switching devices adds no ordering. To order work across a device switch, select the device first and wait afterwards:

event = record_event(backend)   # captures work on the current device
device!(backend, 2)
wait_event(backend, event)      # orders this task's work on device 2 after it

priority!(::Backend, prio::Symbol)::Nothing

Set the priority for the backend stream/queue. This is an optional feature that backends may or may not implement. If a backend shall support priorities it must accept :high, :normal, :low. Where :normal is the default.

device(backend::Backend)::Int

Return the 1-based index of the currently active device for backend.

device(backend::Backend, A::AbstractArray)::Int

Return the 1-based index of the device that owns the memory of A, independently of the currently active device.

ndevices(backend::Backend)::Int

Return the number of devices available to backend.

device!(backend::Backend, id::Int)::Nothing

Select the active device for backend. id is a 1-based device index; an id outside 1:ndevices(backend) throws an ArgumentError.

device! is not a synchronization point: work queued before the switch is not ordered with respect to work queued after it. To order across a switch, either synchronize beforehand, or bracket the switch with record_event and wait_event.

Example

device!(CUDABackend(), 2)  # use the second CUDA device

functional(::Backend)::Union{Bool, Missing}

Queries if the provided backend is functional. This may mean different things for different backends, but generally should mean that the necessary drivers and a compute device are available.

This function should return a Bool or missing if not implemented.

versioninfo(io::IO=stdout, backend::Backend)::Nothing

Print information about backend to io. It is up to the backends to determine what is relevant.

supports_unified(::Backend)::Bool

Whether allocate supports unified=true on the active device: memory that can be accessed from both the host and the device without explicit copies.

supports_atomics(::Backend)::Bool

Whether kernels on the active device support Atomix.jl's atomic operations: at least add and compare-and-swap on 32-bit integers and floats in global memory.

supports_float64(::Backend)::Bool

Whether kernels on the active device support Float64 values.

supports_subgroups(::Backend)::Bool

Whether kernels on the active device support sub-groups: the sub-group queries (get_sub_group_size etc.), sub_group_barrier, and a fixed sub-group width sub_group_size. See the manual for what KernelInterface guarantees about how work-groups are divided into sub-groups; a backend that can't ensure that reports false.

Which types shfl_down supports is queried separately with supports_shuffle.

supports_shuffle(::Backend, ::Type{T})::Bool

Whether kernels on the active device support shfl_down for values of type T.

max_work_group_size(backend)::Int
max_work_group_size(kernel::Kernel)::Int

The largest number of work-items a work-group can have: on the active device of backend, or for launches of the compiled kernel (which may be lower, e.g. because of the kernel's register use). Launching a larger work-group is an error.

The work-group size that performs best is often smaller; see launch_configuration.

max_work_group_size(backend::NewBackend)::Int
max_work_group_size(kernel::Kernel{<:NewBackend})::Int

launch_configuration(kernel::Kernel; nitems=nothing, max_work_group_size=typemax(Int))::@NamedTuple{workgroupsize::Int}

The recommended number of work-items per work-group for launching kernel over nitems work-items in total (nothing if unknown), at most max_work_group_size. This is what an ndrange launch without a workgroupsize uses, passing the number of work-items in the ndrange (saturated at typemax(Int)). nitems and max_work_group_size are positive.

Unlike max_work_group_size, this is advice: backends may base it on occupancy or on the size of the launch, e.g. to prefer more work-groups over larger ones.

launch_configuration(kernel::Kernel{<:NewBackend}; nitems::Union{Int, Nothing}=nothing,
                     max_work_group_size::Int=typemax(Int))::@NamedTuple{workgroupsize::Int}

max_work_group_dims(backend)::NTuple{3, Int}

The maximum number of work-items along each dimension of a work-group, for the active device of backend. max_work_group_size bounds their product.

max_work_group_dims(backend::NewBackend)::NTuple{3, Int}

max_num_groups(backend)::NTuple{3, Int}

The maximum number of work-groups along each dimension of a launch, for the active device of backend.

This is conservative: a launch within these limits works for any work-group size (as long as the number of work-items in each dimension fits an Int), but some backends accept more work-groups for smaller work-groups (e.g. HIP bounds the number of work-items per dimension). The backend's validation at launch time is authoritative.

max_num_groups(backend::NewBackend)::NTuple{3, Int}

sub_group_size(backend)::Int

The sub-group width of kernels compiled for the active device of backend.

Kernels compiled by kernel_function execute with exactly this width: on the device, get_max_sub_group_size returns it, and full sub-groups have this many work-items. Host code can rely on it, e.g. to pick a Val(N) for a warp-level reduction.

sub_group_size(backend::NewBackend)::Int

multiprocessor_count(backend)::Int

The number of multiprocessors (CUDA SMs, AMD CUs, Intel Xe cores, ...) of the active device of backend, or 0 if unknown. The unit differs between backends, so this is only useful for heuristics, e.g. to choose how many work-groups a grid-stride loop launches.

multiprocessor_count(backend::NewBackend)::Int

Kernel{Backend, Kern}

A kernel compiled by kernel_function for backend, wrapping the backend's own kernel object kern. kernel.backend is the backend value that was passed to kernel_function.

Calling a Kernel launches it:

(kernel::Kernel)(args...; numgroups=(), workgroupsize=(), ndrange=(),
                 max_work_group_size=typemax(Int), kwargs...)

args are the host-side arguments, e.g. a CuArray rather than a CuDeviceArray. The backend converts them with argconvert, and the converted types have to match the argument types the kernel was compiled for.

The launch geometry is given in one of three ways:

• numgroups and workgroupsize: launch exactly that many work-groups of that many work-items. Either defaults to 1.
• ndrange and workgroupsize: launch cld.(ndrange, workgroupsize) work-groups.
• ndrange alone: the work-group size is chosen with launch_configuration, bounded by max_work_group_size and by max_work_group_dims, and filled first dimension first.

Each is an Integer or a tuple of up to 3 Integers; missing dimensions are 1. ndrange and numgroups are mutually exclusive.

ndrange is rounded up to whole work-groups and is not masked: the kernel runs for every work-item of every launched group, and get_global_size returns the padded size. Kernels have to check their own bounds.

A zero anywhere in ndrange or numgroups launches nothing. Work-group sizes must be positive and fit max_work_group_dims and max_work_group_size(kernel), and the number of work-items in each dimension must fit an Int; anything else throws an ArgumentError before the backend sees it. The number of work-groups is validated by the backend.

Other keyword arguments are passed to launch unchanged. They are backend-specific: a backend throws an error for keywords it doesn't support.

A launch is queued on the calling task's queue of the active device, and returns nothing.

kernel_function(backend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...)::Kernel

Compile the callable f for arguments of the (device-side) types tt, for the active device of backend, returning a Kernel. For a higher-level interface, use KernelInterface.@launch.

f is the host-side callable, not converted with argconvert: the backend converts it. For a closure, that matters: it can capture arrays, which its converted form only holds pointers to.

Keyword arguments:

• name: override the name that the kernel will have in the generated code.

Other keyword arguments are backend-specific compiler options (e.g. maxthreads for CUDA.jl); backends throw an error for options they don't support.

The returned kernel keeps f alive, but not the arguments: they are passed again at launch.

kernel_function(backend::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}

argconvert(backend, arg)

Convert arg to its device-side representation, e.g. a CuArray to a CuDeviceArray. Called for every kernel argument, and for the kernel function itself.

It has to be pure: it may be called more than once for the same launch.

argconvert(::NewBackend, arg)

launch(kernel::Kernel, groups::Dims{3}, items::Dims{3}, args::Tuple; kwargs...)

Launch kernel with groups work-groups of items work-items each, passing the host-side arguments args, a tuple. This is what calling a Kernel does after validating and normalizing the launch geometry; users call the kernel instead.

groups and items are positive, items fits max_work_group_dims and max_work_group_size(kernel), and groups .* items doesn't overflow Int. kwargs are the keyword arguments of the call that KernelInterface doesn't know.

launch(kernel::Kernel{<:NewBackend}, groups::Dims{3}, items::Dims{3}, args::Tuple; kwargs...)

KI.@launch backend [launch=true] [numgroups=...] [workgroupsize=...] [ndrange=...] [max_work_group_size=...] [kwargs...] f(args...)

Compile f(args...) for backend and launch it, like @cuda or @metal do.

f is compiled with kernel_function for the types of the arguments converted with argconvert, and the resulting Kernel is called with the launch keywords numgroups, workgroupsize, ndrange and max_work_group_size, whose meaning is documented there. The arguments are kept alive while the launch is being queued.

Other keyword arguments:

• launch: whether to launch the kernel, defaults to true. With launch=false, the kernel is only compiled and returned, and the launch keywords can't be used: launch it by calling it with the arguments and the launch keywords.
• name and any other keyword are passed to kernel_function as compiler options.

Launch options specific to a backend (such as a CUDA stream) can't be passed to @launch; use launch=false and pass them when calling the kernel.

backend is evaluated once. Returns the Kernel.

function vadd(c, a, b)
    i = KI.get_global_id().x
    if i <= length(c)
        @inbounds c[i] = a[i] + b[i]
    end
    return
end

KI.@launch backend ndrange=length(c) vadd(c, a, b)