General Loops · AcceleratedKernels.jl

General Looping: `foreachindex` / `foraxes`

General workhorses for converting normal Julia for loops into GPU code, for example:

# Copy kernel testing throughputfunction cpu_copy!(dst, src)    for i in eachindex(src)        dst[i] = src[i]    endend

Would be written for GPU as:

import AcceleratedKernels as AKfunction gpu_copy!(dst, src)    AK.foreachindex(src) do i        dst[i] = src[i]    endend

Yes, simply change for i in eachindex(itr) into AK.foreachindex(itr) do i to run it on GPUs / multithreaded - magic! (or just amazing language design)

This is a parallelised for-loop over the indices of an iterable; converts normal Julia code to GPU kernels running one thread per index. On CPUs it executes static index ranges on max_tasks threads, with user-defined min_elems to be processed by each thread; if only a single thread ends up being needed, the loop is inlined and executed without spawning threads. Kokkos and RAJA call this Kokkos::parallel_for and RAJA::forall; Thrust's closest is thrust::transform.

Example:

import AcceleratedKernels as AKfunction f(a, b)    # Don't use global arrays inside a `foreachindex`; types must be known    @assert length(a) == length(b)    AK.foreachindex(a) do i        # Note that we don't have to explicitly pass b into the lambda        if b[i] > 0.5            a[i] = 1        else            a[i] = 0        end        # Showing arbitrary if conditions; can also be written as:        # @inbounds a[i] = b[i] > 0.5 ? 1 : 0    endend# Use any backend, e.g. CUDA, ROCm, oneAPI, Metal, or CPUusing oneAPIv1 = oneArray{Float32}(undef, 100_000)v2 = oneArray(rand(Float32, 100_000))f(v1, v2)

The loop runs on the backend of the iterable. A range does not determine one, so pass backend when a loop over a range accesses GPU arrays:

AK.foreachindex(1:n; backend=AK.get_backend(v)) do i    v[i] = iend

Host arrays and ranges run on Julia threads.

foreachindex, foraxes, map!, reverse! and the binary searches take a GPU block size - this is often a power of two (mostly 64, 128, 256, 512); the optimum depends on the algorithm, input data and hardware - you can try the different values and @time or @benchmark them:

@time AK.foreachindex(f, itr_gpu; block_size=512)

Similarly, for performance on the CPU the overhead of spawning threads should be masked by processing more elements per thread (but there is no reason here to launch more threads than Threads.nthreads(), the number of threads Julia was started with); the optimum depends on how expensive f is - again, benchmarking is your friend:

@time AK.foreachindex(f, itr_cpu; max_tasks=16, min_elems=1000)

The other operations (sorting, reductions, scans, findall, any/all) take these settings through their alg keyword instead; see Algorithms and backends.

AcceleratedKernels.foreachindex — Function

foreachindex(
    f, itr;
    backend=nothing,
    block_size::Int=256,
    max_tasks::Int=Threads.nthreads(),
    min_elems::Int=1,
)

Parallelised for loop over the indices of an iterable: call f(i) for every i in eachindex(itr).

It allows you to run normal Julia code on a GPU over multiple arrays - e.g. CuArray, ROCArray, MtlArray, oneArray - with one GPU thread per index, in blocks of block_size threads.

On the host backend, the loop runs on Julia threads: at most max_tasks tasks, or fewer such that each task processes at least min_elems indices; if a single task ends up being needed, f is inlined and no task is launched. Tune it to your function - the more expensive it is, the fewer elements are needed to amortise the cost of launching a task (which is a few μs).

backend is derived from itr. A range or other index collection does not determine a backend, so loops over one run on the host unless you pass backend, e.g. the backend of the arrays that f accesses.

Examples

Normally you would write a for loop like this:

function f()    x = Array(1:100)    y = similar(x)    for i in eachindex(x)        @inbounds y[i] = 2 * x[i] + 1    endend

Using this function you can have the same for loop body over a GPU array:

using CUDAimport AcceleratedKernels as AKfunction f()    x = CuArray(1:100)    y = similar(x)    AK.foreachindex(x) do i        @inbounds y[i] = 2 * x[i] + 1    endend

A loop over a range needs the backend of the arrays it accesses:

function g!(y, x)    AK.foreachindex(1:length(x) ÷ 2; backend=AK.get_backend(x)) do i        @inbounds y[i] = x[2i]    endend

Note that the above code is pure arithmetic, which you can write directly (and on some platforms it may be faster) as:

using CUDAx = CuArray(1:100)y = 2 .* x .+ 1

To use this function on a GPU, the objects referenced inside the loop body must have known types - i.e. be inside a function. For example:

using oneAPIimport AcceleratedKernels as AKx = oneArray(1:100)# CRASHES - typical error message: "Reason: unsupported dynamic function invocation"# AK.foreachindex(x) do i#     x[i] = i# endfunction somecopy!(v)    # Because it is inside a function, the type of `v` will be known    AK.foreachindex(v) do i        v[i] = i    endendsomecopy!(x)    # This works

source

AcceleratedKernels.foraxes — Function

foraxes(
    f, itr, dims::Union{Nothing, Integer}=nothing;
    backend=nothing,
    block_size::Int=256,
    max_tasks::Int=Threads.nthreads(),
    min_elems::Int=1,
)

Parallelised for loop over the indices along axis dims of an iterable: call f(i) for every i in axes(itr, dims), or in eachindex(itr) for dims=nothing. dims mirrors axes(itr, dims). The keywords are those of foreachindex.

Examples

Normally you would write a for loop like this:

function f()    x = Array(reshape(1:30, 3, 10))    y = similar(x)    for i in axes(x, 2)        for j in axes(x, 1)            @inbounds y[j, i] = 2 * x[j, i] + 1        end    endend

Using this function you can have the same for loop body over a GPU array:

using CUDAimport AcceleratedKernels as AKfunction f()    x = CuArray(reshape(1:3000, 3, 1000))    y = similar(x)    AK.foraxes(x, 2) do i        for j in axes(x, 1)            @inbounds y[j, i] = 2 * x[j, i] + 1        end    endend

As with foreachindex, the objects referenced inside the loop body must have known types on a GPU, i.e. be inside a function.

source

General Looping: foreachindex / foraxes

General Looping: `foreachindex` / `foraxes`