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 (unless the layout is linear, see supports_linear_subgroups): 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 (the warp or wavefront size), sub_group_size(backend) on the host.

Backends should make it a constant of the generated code, so that code depending on it, e.g. a loop over the lanes or a shuffle butterfly, is specialized for it; on some, it is only folded late in compilation, or not at all. Don't rely on it being known during type inference: to pick types or Val parameters from the width, e.g. the integer type of a mask with a bit per lane, pass sub_group_size(backend) from 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 (unless the layout is linear, see supports_linear_subgroups). 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, unless the layout is linear (see supports_linear_subgroups, which only covers some work-group shapes).

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. How it relates to get_local_id is unspecified, unless the layout is linear (see supports_linear_subgroups, which only covers some work-group shapes).

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(); see shfl.

@device_override sub_group_barrier()

sub_group_barrier(width::Integer)

Like sub_group_barrier, for the work-items of the caller's segment of width consecutive lanes of the sub-group, as for the shuffles with a width (see shfl): wait until they have reached the barrier, and make their writes to global and local memory before it visible to the segment. It doesn't order memory with other segments.

sub_group_barrier(width::Integer)

Like sub_group_barrier, for the work-items of the caller's segment of width consecutive lanes of the sub-group, as for the shuffles with a width (see shfl): wait until they have reached the barrier, and make their writes to global and local memory before it visible to the segment. It doesn't order memory with other segments.

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(val::T, lane::Integer)::T

Return val of the work-item with get_sub_group_local_id equal to lane in the sub-group. 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 together (not in a divergent branch), but they may read from different lanes. Unless supports_independent_subgroups, all sub-groups of the work-group have to execute it.

shfl doesn't synchronize memory accesses: use sub_group_barrier to communicate through memory.

Types for which supports_shuffle returns true are supported. Besides the types a backend supports natively, that includes other primitive types of 1, 2, 4, 8 or 16 bytes (e.g. Bool, Char or Int64) if the backend supports UInt32, and isbits structs and tuples of such types. Primitive types of up to 4 bytes are shuffled as a UInt32, and larger ones as UInt64 words, each of which is shuffled natively if the backend supports UInt64, and as two UInt32 words otherwise. Structs and tuples are packed into UInt32 words, with fields of less than 4 bytes sharing words, and shuffled word by word.

@device_override shfl(val::T, lane::Integer) where {T <: Union{...}}

shfl(val::T, lane::Integer, width::Integer)::T
shfl_down(val::T, offset::Integer, width::Integer)::T
shfl_up(val::T, offset::Integer, width::Integer)::T
shfl_xor(val::T, mask::Integer, width::Integer)::T

Shuffles within segments of width consecutive lanes of the sub-group, with CUDA's semantics for a width:

• shfl reads from lane mod(lane - 1, width) + 1 of the caller's segment;
• shfl_down and shfl_up read from the lane offset lanes further or earlier, for an offset of at least 0, and return val of the work-item itself where that lane is outside of the caller's segment;
• shfl_xor reads from the lane whose 0-based position in the segment is the caller's xor mask, which has to be between 0 and width - 1.

Reading from a lane of the segment that has no work-item (in a partial sub-group) gives an unspecified value. width has to be a power of two of at most the sub-group width get_max_sub_group_size, and the same for all work-items of the sub-group, so that the segments partition it. It should be a constant of the generated code.

If supports_partitionable_subgroups, only the work-items of the caller's segment have to execute them together: the segments of a sub-group can diverge, e.g. loop a different number of times. Otherwise, all work-items of the sub-group have to execute them together, as for the shuffles without a width. Either way, offset and mask have to be the same for the work-items that execute them together, while lane may differ.

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, for an offset of at least 0. When that lane is past the sub-group width, i.e. get_sub_group_local_id() + offset > get_max_sub_group_size(), the result is val of the work-item itself, as on CUDA, HIP and Metal. When the lane is within the width but has no work-item, in a partial sub-group, 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; see shfl.

shfl_down doesn't synchronize memory accesses. See shfl for the supported types.

@device_override shfl_down(val::T, offset::Integer) where {T <: Union{...}}

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

Return val of the work-item offset lanes earlier in the sub-group, i.e. with get_sub_group_local_id equal to get_sub_group_local_id() - offset, for an offset of at least 0. When there is no such lane, i.e. get_sub_group_local_id() <= offset, the result is val of the work-item itself, as on CUDA, HIP and Metal.

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

shfl_up doesn't synchronize memory accesses. See shfl for the supported types.

@device_override shfl_up(val::T, offset::Integer) where {T <: Union{...}}

shfl_xor(val::T, mask::Integer)::T

Return val of the work-item whose 0-based lane id is the 0-based lane id of this work-item xor mask, i.e. with get_sub_group_local_id equal to ((get_sub_group_local_id() - 1) ⊻ mask) + 1. mask has to be between 0 and get_max_sub_group_size() - 1. When that lane has no work-item, in a partial sub-group, the result is an unspecified value (of type T).

All work-items of the sub-group have to execute shfl_xor together (not in a divergent branch), with the same mask; see shfl. For a reduction over the sub-group, use sub_group_reduce, which also handles partial sub-groups.

shfl_xor doesn't synchronize memory accesses. See shfl for the supported types.

@device_override shfl_xor(val::T, mask::Integer) where {T <: Union{...}}

shfl(val::T, lane::Integer, width::Integer)::T
shfl_down(val::T, offset::Integer, width::Integer)::T
shfl_up(val::T, offset::Integer, width::Integer)::T
shfl_xor(val::T, mask::Integer, width::Integer)::T

Shuffles within segments of width consecutive lanes of the sub-group, with CUDA's semantics for a width:

• shfl reads from lane mod(lane - 1, width) + 1 of the caller's segment;
• shfl_down and shfl_up read from the lane offset lanes further or earlier, for an offset of at least 0, and return val of the work-item itself where that lane is outside of the caller's segment;
• shfl_xor reads from the lane whose 0-based position in the segment is the caller's xor mask, which has to be between 0 and width - 1.

Reading from a lane of the segment that has no work-item (in a partial sub-group) gives an unspecified value. width has to be a power of two of at most the sub-group width get_max_sub_group_size, and the same for all work-items of the sub-group, so that the segments partition it. It should be a constant of the generated code.

If supports_partitionable_subgroups, only the work-items of the caller's segment have to execute them together: the segments of a sub-group can diverge, e.g. loop a different number of times. Otherwise, all work-items of the sub-group have to execute them together, as for the shuffles without a width. Either way, offset and mask have to be the same for the work-items that execute them together, while lane may differ.

sub_group_any(pred::Bool)::Bool

Whether pred is true for any work-item of the sub-group. All work-items of the sub-group get the same result.

All work-items of the sub-group have to execute sub_group_any together (not in a divergent branch); see shfl.

It doesn't synchronize memory accesses, see sub_group_barrier.

@device_override sub_group_any(pred::Bool)::Bool

sub_group_all(pred::Bool)::Bool

Whether pred is true for all work-items of the sub-group. All work-items of the sub-group get the same result.

All work-items of the sub-group have to execute sub_group_all together (not in a divergent branch); see shfl.

It doesn't synchronize memory accesses, see sub_group_barrier.

@device_override sub_group_all(pred::Bool)::Bool

sub_group_ballot(pred::Bool)::UInt64

A mask of the work-items of the sub-group for which pred is true: bit i - 1 (counting from the least significant bit) is set if and only if the sub-group has a work-item with get_sub_group_local_id equal to i and its pred is true. All other bits, e.g. those of lanes without a work-item in a partial sub-group, are zero. All work-items of the sub-group get the same result. count_ones(mask) counts the work-items, and for a nonzero mask, trailing_zeros(mask) + 1 is the lane of the first one.

All work-items of the sub-group have to execute sub_group_ballot together (not in a divergent branch); see shfl. Only sub-groups of at most 64 work-items are supported.

It doesn't synchronize memory accesses, see sub_group_barrier.

@device_override sub_group_ballot(pred::Bool)::UInt64

sub_group_any(pred::Bool, width::Integer)::Bool
sub_group_all(pred::Bool, width::Integer)::Bool
sub_group_ballot(pred::Bool, width::Integer)::UInt64

The votes within segments of width consecutive lanes of the sub-group, as for the shuffles with a width (see shfl): whether pred is true for any or all work-items of the caller's segment, or a mask of those for which it is, relative to the segment: bit i - 1 for the i-th lane of the segment. Lanes without a work-item don't count.

For a width smaller than the sub-group width, they need a sub-group width of at most 64.

sub_group_reduce(op, val::T)::T

Reduce val over the work-items of the sub-group with the binary operator op, in the order of the lanes: op(…op(op(v₁, v₂), v₃)…, vₙ) for the values vᵢ of lanes 1:n, with n = get_sub_group_size(), up to associativity. op has to be associative, but needn't be commutative. All work-items of the sub-group get the result.

All work-items of the sub-group have to execute sub_group_reduce together (not in a divergent branch), with the same op; see shfl. Values of the types that the shuffles support are supported, see supports_shuffle, and others for which the backend has a native reduction.

It doesn't synchronize memory accesses, see sub_group_barrier.

sub_group_reduce(op, val::T, width::Integer)::T
sub_group_scan(op, val::T, width::Integer)::T
sub_group_exclusive_scan(op, val::T, init::T, width::Integer)::T

The collectives within segments of width consecutive lanes of the sub-group, as for the shuffles with a width (see shfl): the reduction and scans of sub_group_reduce, sub_group_scan and sub_group_exclusive_scan, over the work-items of the caller's segment, in the order of its lanes. init has to be the same for the work-items of a segment.

sub_group_scan(op, val::T)::T

The inclusive scan of val over the work-items of the sub-group with the associative binary operator op, in the order of the lanes: the work-item in lane i gets the reduction of the values of lanes 1 to i, as with sub_group_reduce.

All work-items of the sub-group have to execute sub_group_scan together (not in a divergent branch), with the same op; see shfl. Values of the types that the shuffles support are supported, see supports_shuffle, and others for which the backend has a native scan.

It doesn't synchronize memory accesses, see sub_group_barrier.

sub_group_exclusive_scan(op, val::T, init::T)::T

The exclusive scan of val over the work-items of the sub-group with the associative binary operator op, in the order of the lanes, starting from init: the work-item in lane 1 gets init itself, and the one in lane i > 1 gets op(init, r) for the reduction r of the values of lanes 1 to i - 1, as with sub_group_scan. init has to be the same for all work-items of the sub-group, and of the type of val; it needn't be an identity of op.

All work-items of the sub-group have to execute sub_group_exclusive_scan together (not in a divergent branch), with the same op; see shfl. Values of the types that the shuffles support are supported, see supports_shuffle, and others for which the backend has a native scan.

It doesn't synchronize memory accesses, see sub_group_barrier.

sub_group_reduce(op, val::T, width::Integer)::T
sub_group_scan(op, val::T, width::Integer)::T
sub_group_exclusive_scan(op, val::T, init::T, width::Integer)::T

The collectives within segments of width consecutive lanes of the sub-group, as for the shuffles with a width (see shfl): the reduction and scans of sub_group_reduce, sub_group_scan and sub_group_exclusive_scan, over the work-items of the caller's segment, in the order of its lanes. init has to be the same for the work-items of a segment.

_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 with a fixed width sub_group_size: the sub-group queries (get_sub_group_size etc.), sub_group_barrier, and the votes sub_group_any, sub_group_all and, for widths of at most 64, sub_group_ballot. Which types the shuffles support is queried separately with supports_shuffle.

Which work-items form a sub-group is unspecified, unless supports_linear_subgroups returns true, and the communication functions have to be executed in control flow that is uniform over the work-group, unless supports_independent_subgroups returns true. See Sub-groups for the layout and participation guarantees.

supports_linear_subgroups(::Backend)::Bool

Whether the active device forms sub-groups from consecutive work-items, for every kernel and launch with a 1-D work-group or one whose x extent is a multiple of the sub-group width W (sub_group_size): then the work-item with the linear local index lin (x fastest, see get_local_id) is in sub-group fld(lin - 1, W) + 1, at lane mod(lin - 1, W) + 1, the sub-group ids are 1:cld(n, W) for a work-group of n work-items, and only the last sub-group can be partial. In particular, a 1-D work-group of at most W work-items is a single sub-group. It promises nothing about other work-group shapes. Implies supports_subgroups.

Code needs this when it relates local ids to lanes: e.g. stencils that shuffle values between neighbouring work-items, transposes through shuffles, or work-group scans built from sub-group scans. Code that doesn't can index storage by the sub-group slot (get_sub_group_id() - 1) * W + get_sub_group_local_id(), which is unique within the work-group, but not dense: it can exceed the work-group size, so storage has to have room for get_num_sub_groups() * W values.

false means that the layout isn't guaranteed, not that it is different: OpenCL, Vulkan and other APIs leave it to the implementation.

supports_independent_subgroups(::Backend)::Bool

Whether the sub-group communication functions (shuffles, votes, collectives and sub_group_barrier) can be executed in control flow that is uniform over the sub-group, but not over the work-group: e.g. a loop whose trip count differs between sub-groups, or after an early return of whole sub-groups. Otherwise, all sub-groups of a work-group have to execute the same communication functions, in the same order.

This covers whole sub-groups, not the segments of the functions with a width (see supports_partitionable_subgroups), and doesn't promise forward progress: a sub-group spin-waiting for another one can still hang. Implies supports_subgroups.

supports_partitionable_subgroups(::Backend)::Bool

Whether a sub-group can be partitioned into segments of width consecutive lanes that diverge independently, for the sub-group communication functions with a width (the shuffles, votes, collectives and sub_group_barrier with a width, e.g. shfl): those only have to be executed by the work-items of the caller's segment, rather than by the whole sub-group, so the segments of a sub-group can take different branches, loop a different number of times, or skip the functions altogether.

This extends supports_independent_subgroups to segments, and implies it. Like that query, it doesn't promise forward progress between segments.

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

Whether kernels on the active device support the shuffles shfl, shfl_down, shfl_up and shfl_xor for values of type T.

Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if UInt32 is, and so are isbits structs and tuples of such types (nested ones included). The fallbacks shuffle these as words; see shfl for how.

This reports whether values of type T can be moved between work-items, not whether the device can compute with them: e.g. Float64 values can be shuffled (as words) on a device without supports_float64, as long as they are only stored and loaded, e.g. through an array of UInt64.

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)