APIs

ClimaComms

Abstract the computing devices and communication contexts used by CliMA packages, so that the same simulation code can run on a single CPU thread, on multiple CPU threads, on NVIDIA GPUs, and across MPI processes.

The two central abstractions are:

• AbstractDevice: the hardware a computation runs on (e.g., CPUSingleThreaded, CUDADevice).
• AbstractCommsContext: the environment through which processes communicate (e.g., SingletonCommsContext, MPICommsContext).

Devices and contexts are selected at runtime, typically from the CLIMACOMMS_DEVICE and CLIMACOMMS_CONTEXT environment variables via device and context. Backend-specific implementations (CUDA, MPI) live in package extensions and are loaded with @import_required_backends.

ClimaComms.@import_required_backends

Import the backend packages required by the runtime configuration: if the CLIMACOMMS_CONTEXT environment variable requests MPI, import MPI.jl; if the CLIMACOMMS_DEVICE environment variable requests CUDA, import CUDA.jl. The packages must be available in the active Julia environment.

Add this macro to the top of driver scripts, after import ClimaComms, so that the same script works for any device and context.

Examples

import ClimaComms
ClimaComms.@import_required_backends
context = ClimaComms.context()

ClimaComms.cuda_is_required()

Return true if CUDA.jl needs to be loaded, based on the CLIMACOMMS_DEVICE environment variable. See device for more information.

Examples

cuda_is_required() && using CUDA

ClimaComms.mpi_is_required()

Return true if MPI.jl needs to be loaded, based on the CLIMACOMMS_CONTEXT environment variable. See context for more information.

Examples

mpi_is_required() && using MPI

ClimaComms.cuda_ext_is_loaded()

Return true if the ClimaCommsCUDAExt extension is loaded (i.e., if CUDA.jl has been imported).

ClimaComms.mpi_ext_is_loaded()

Return true if the ClimaCommsMPIExt extension is loaded (i.e., if MPI.jl has been imported).

AbstractDevice

The computing device on which code is executed.

Devices are empty structs: they carry no data and exist so that multiple dispatch can select device-specific implementations (e.g., Array vs. CuArray, a serial loop vs. a CUDA kernel).

Subtypes:

• CPUSingleThreaded: a CPU using a single thread.
• CPUMultiThreaded: a CPU using multiple threads.
• CUDADevice: a single CUDA-enabled GPU.

Use device to select a device at runtime from the CLIMACOMMS_DEVICE environment variable.

AbstractCPUDevice

Abstract device type for single-threaded and multi-threaded CPU runs.

Subtypes:

• CPUSingleThreaded: a CPU using a single thread.
• CPUMultiThreaded: a CPU using multiple threads.

CPUSingleThreaded()

Use the CPU with a single thread.

CPUMultiThreaded()

Use the CPU with multiple threads.

CUDADevice()

Use an NVIDIA GPU via CUDA.jl.

CUDA.jl must be loaded for this device to be usable; see @import_required_backends.

ClimaComms.device()

Construct the device specified by the CLIMACOMMS_DEVICE environment variable.

Allowed values of CLIMACOMMS_DEVICE:

• CPU (default): CPUSingleThreaded or CPUMultiThreaded, depending on the number of Julia threads;
• CPUSingleThreaded;
• CPUMultiThreaded;
• CUDA: CUDADevice, which requires CUDA.jl to be loaded (see @import_required_backends).

Examples

device = ClimaComms.device()
ArrayType = ClimaComms.array_type(device)

See also context.

ClimaComms.device_functional(device)

Return true when the device is correctly set up (e.g., for a CUDADevice, when CUDA is available and functional).

ClimaComms.array_type(device::AbstractDevice)

Return the base array type used by the specified device (currently Array or CuArray).

Examples

ArrayType = ClimaComms.array_type(ClimaComms.device())
x = ArrayType([1.0, 2.0, 3.0])

ClimaComms.free_memory(device)

Return the bytes of memory that are currently available for allocation on the device.

ClimaComms.total_memory(device)

Return the bytes of memory that are theoretically available for allocation on the device.

Adapt.adapt_structure(to::Type{<:AbstractArray}, device::AbstractDevice)

Adapt a given device to the device associated with the given array type.

Examples

julia> Adapt.adapt(Array, ClimaComms.CUDADevice())
ClimaComms.CPUSingleThreaded()

@time device expr

Device-flexible @time.

Lowers to

@time expr

for CPU devices and

CUDA.@time expr

for CUDA devices.

@elapsed device expr

Device-flexible @elapsed.

Lowers to

@elapsed expr

for CPU devices and

CUDA.@elapsed expr

for CUDA devices.

@assert device cond [text]

Device-flexible @assert.

Lowers to

@assert cond [text]

for CPU devices and

CUDA.@cuassert cond [text]

for CUDA devices.

@sync device expr

Device-flexible @sync.

Lowers to

@sync expr

for CPU devices and

CUDA.@sync expr

for CUDA devices.

An example use-case of this might be:

BenchmarkTools.@benchmark begin
    if ClimaComms.device() isa ClimaComms.CUDADevice
        CUDA.@sync begin
            launch_cuda_kernels_or_spawn_tasks!(...)
        end
    elseif ClimaComms.device() isa ClimaComms.CPUMultiThreading
        Base.@sync begin
            launch_cuda_kernels_or_spawn_tasks!(...)
        end
    end
end

If the CPU version of the above example does not leverage spawned tasks (which require using Base.sync or Threads.wait to synchronize), then you may want to simply use @cuda_sync.

@cuda_sync device expr

Device-flexible CUDA.@sync.

Lowers to

expr

for CPU devices and

CUDA.@sync expr

for CUDA devices.

ClimaComms.time(f, device, args...; kwargs...)

Device-flexible version of @time; functional form of @time.

Calls

@time f(args...; kwargs...)

for CPU devices and

CUDA.@time f(args...; kwargs...)

for CUDA devices.

ClimaComms.elapsed(f, device, args...; kwargs...)

Device-flexible version of @elapsed; functional form of @elapsed.

Calls

@elapsed f(args...; kwargs...)

for CPU devices and

CUDA.@elapsed f(args...; kwargs...)

for CUDA devices.

ClimaComms.sync(f, device, args...; kwargs...)

Device-flexible version of @sync; functional form of @sync.

Calls

@sync f(args...; kwargs...)

for CPU devices and

CUDA.@sync f(args...; kwargs...)

for CUDA devices.

An example use-case of this might be:

BenchmarkTools.@benchmark begin
    if ClimaComms.device() isa ClimaComms.CUDADevice
        CUDA.@sync begin
            launch_cuda_kernels_or_spawn_tasks!(...)
        end
    elseif ClimaComms.device() isa ClimaComms.CPUMultiThreading
        Base.@sync begin
            launch_cuda_kernels_or_spawn_tasks!(...)
        end
    end
end

If the CPU version of the above example does not leverage spawned tasks (which require using Base.sync or Threads.wait to synchronize), then you may want to simply use cuda_sync.

ClimaComms.cuda_sync(f, device, args...; kwargs...)

Device-flexible version of CUDA.@sync; functional form of @cuda_sync.

Calls

f(args...; kwargs...)

for CPU devices and

CUDA.@sync f(args...; kwargs...)

for CUDA devices.

ClimaComms.allowscalar(f, device, args...; kwargs...)

Device-flexible version of CUDA.@allowscalar.

Lowers to

f(args...)

for CPU devices and

CUDA.@allowscalar f(args...)

for CUDA devices.

This is usefully written with closures via

allowscalar(device) do
    f()
end

@threaded [device] [coarsen=...] [block_size=...] for ... end

Device-flexible generalization of Threads.@threads, which distributes the iterations of a for-loop across multiple threads, with the option to control thread coarsening and GPU kernel configuration. Coarsening makes each thread evaluate more than one iteration of the loop, which can improve performance by reducing the runtime overhead of launching additional threads (though too much coarsening worsens performance because it reduces parallelization). The device is either inferred by calling ClimaComms.device(), or it can be specified manually, with the following device-dependent behavior:

• When device is a CPUSingleThreaded(), the loop is evaluated as-is. This avoids the runtime overhead of calling Threads.@threads with a single thread, and, when the device type is statically inferrable, it also avoids compilation overhead.

• When device is a CPUMultiThreaded(), the loop is passed to Threads.@threads. This supports three different kinds of "schedulers" for determining how many iterations of the loop to evaluate in each thread:

  1. (default) a "dynamic" scheduler that changes the number of iterations as new threads are launched,
  2. a "static" scheduler that evaluates a fixed number of iterations per thread, and
  3. a "greedy" scheduler that uses a small number of threads, continuously evaluating iterations in each thread until the loop is completed (only available as of Julia 1.11).

  Setting coarsen to :dynamic or :greedy launches threads with those schedulers. Setting it to :static or an integer value launches threads with static scheduling (using :static is similar to using 1, but slightly more performant). To read more about multi-threading, see the documentation for Threads.@threads.

• When device is a CUDADevice(), the loop is compiled with CUDA.@cuda and run with CUDA.@sync. Since CUDA launches all threads at the same time, only static scheduling can be used. Setting coarsen to any symbol causes each thread to evaluate a single iteration (default), and setting it to an integer value causes each thread to evaluate that number of iterations (the default is similar to using 1, but slightly more performant). If the total number of iterations in the loop is extremely large, the specified coarsening may require more threads than can be simultaneously launched on the GPU, in which case the amount of coarsening is automatically increased.

  The optional argument block_size is also available for manually specifying the size of each block on a GPU. The default value of :auto sets the number of threads in each block to the largest possible value that permits a high GPU "occupancy" (the number of active thread warps in each multiprocessor executing the kernel). An integer can be used instead of :auto to override this default value. If the specified value exceeds the total number of threads, it is automatically decreased to avoid idle threads.

Any iterator with methods for firstindex, length, and getindex can be used in a @threaded loop. All lazy iterators from Base and Base.Iterators, such as zip, enumerate, Iterators.product, and generator expressions, are also compatible with @threaded. (Although these iterators do not define methods for getindex, they are automatically modified by threadable to support getindex.) Using multiple iterators with @threaded is equivalent to looping over a single Iterators.product, with the innermost iterator of the loop appearing first in the product, and the outermost iterator appearing last.

julia> a = CUDA.CuArray{Int}(undef, 100); b = similar(a);

julia> threaded_copyto!(a, b) = ClimaComms.@threaded for i in axes(a, 1)
           a[i] = b[i]
       end
threaded_copyto! (generic function with 1 method)

julia> threaded_copyto!(a, b)

julia> ClimaComms.@threaded for i in axes(a, 1)
           a[i] = b[i]
       end
ERROR: InvalidIRError: ...

Moreover, type variables are not inferrable across function boundaries, so types used in a threaded loop cannot be precomputed before the loop:

julia> threaded_add_epsilon!(a) = ClimaComms.@threaded for i in axes(a, 1)
           FT = eltype(a)
           a[i] += eps(FT)
       end
threaded_add_epsilon! (generic function with 1 method)

julia> threaded_add_epsilon!(a)

julia> function threaded_add_epsilon!(a)
           FT = eltype(a)
           ClimaComms.@threaded for i in axes(a, 1)
               a[i] += eps(FT)
           end
       end
threaded_add_epsilon! (generic function with 1 method)

julia> threaded_add_epsilon!(a)
ERROR: InvalidIRError: ...

To fix other kinds of inference issues on GPUs, especially ones brought about by indexing into iterators with nonuniform element types, see UnrolledUtilities.jl.

ClimaComms.threaded(f, device, itrs...; kwargs...)

Functional form of @threaded. If there are n iterators and f is a function of n arguments, the threaded function is similar to

@threaded device [kwargs...] for xₙ in itrs[n], ..., x₂ in itrs[2], x₁ in itrs[1]
    f(x₁, x₂, ..., xₙ)
end

On single-threaded CPU devices, the @threaded macro inlines the for-loop without any intermediate function calls, so that it has a lower latency than the threaded function. On other devices, the only difference between the macro and the function is that keyword argument symbols like :dynamic and :auto must be wrapped in Vals for the function.

ClimaComms.threadable(device, itr)

Return a version of the iterator itr that can be used in a @threaded loop; either itr itself or a ThreadableWrapper of itr.

ThreadableWrapper

Wrapper for an iterator from Base or Iterators that can be used in @threaded, with methods for firstindex, length, and getindex. The getindex method only supports linear indices between firstindex and firstindex + length - 1. For the ThreadableWrapper of Iterators.product, getindex converts each linear index to a Cartesian index using regular integer division on CPUs and Base.multiplicativeinverse on GPUs.

AbstractCommsContext

The environment through which processes communicate.

A context wraps an AbstractDevice and, for distributed runs, the information needed for processes to exchange data. Contexts make code independent of the form of parallelism: communication primitives such as reduce, gather, and barrier dispatch on the context and become no-ops in single-process runs.

Subtypes:

• SingletonCommsContext: a single process; all communication primitives are no-ops.
• MPICommsContext: distributed runs via MPI.

Use context to select a context at runtime from the CLIMACOMMS_CONTEXT environment variable.

SingletonCommsContext(device = device())

A communications context for single-process runs. All communication primitives (e.g., reduce, gather, barrier) are no-ops. AbstractCPUDevice and CUDADevice device options are currently supported.

Fields

• device: the AbstractDevice on which computations run.

MPICommsContext()
MPICommsContext(device)
MPICommsContext(device, comm)

An MPI communications context, used for distributed runs. AbstractCPUDevice and CUDADevice device options are currently supported. The comm argument defaults to MPI.COMM_WORLD.

MPI.jl must be loaded for this context to be usable; see @import_required_backends.

Fields

• device: the AbstractDevice on which computations run.
• mpicomm: the MPI communicator (an MPI.Comm).

ClimaComms.context(device = device())

Construct the communication context specified by the CLIMACOMMS_CONTEXT environment variable.

Allowed values of CLIMACOMMS_CONTEXT:

• SINGLETON (default): a SingletonCommsContext, for single-process runs;
• MPI: an MPICommsContext, for distributed runs, which requires MPI.jl to be loaded (see @import_required_backends).

The context wraps the given device; by default, the device is also read from an environment variable (see device).

Examples

context = ClimaComms.context()
device = ClimaComms.device(context)

ClimaComms.local_communicator(ctx::MPICommsContext)
ClimaComms.local_communicator(f, ctx::MPICommsContext)

Create a new MPI communicator containing the processes on the same physical node as the caller. In the single-argument form, the caller is responsible for freeing the communicator; the two-argument (do-block) form calls f on the communicator and then frees it.

Called from init to assign GPUs to the MPI ranks on each node.

Examples

ClimaComms.local_communicator(ctx) do local_comm
    ClimaComms._assign_device(
        ClimaComms.device(ctx),
        MPI.Comm_rank(local_comm),
    )
end

Adapt.adapt_structure(to::Type{<:AbstractArray}, ctx::AbstractCommsContext)

Adapt a given context to a context whose device is associated with the given array type.

Examples

julia> Adapt.adapt(Array, ClimaComms.context(ClimaComms.CUDADevice()))
ClimaComms.SingletonCommsContext{ClimaComms.CPUSingleThreaded}(ClimaComms.CPUSingleThreaded())

ClimaComms.init(ctx::AbstractCommsContext)

Perform any necessary initialization for the specified backend (e.g., initializing MPI and assigning GPUs to MPI ranks). Return a tuple (pid, nprocs) of the process ID and the number of participating processes.

Call this once, before any other communication operations on ctx.

ClimaComms.mypid(ctx::AbstractCommsContext)

Return the process ID of the calling process, an integer between 1 and nprocs. The root process has mypid(ctx) == 1.

ClimaComms.iamroot(ctx::AbstractCommsContext)

Return true if the calling process is the root process (the process with ID 1).

ClimaComms.nprocs(ctx::AbstractCommsContext)

Return the number of participating processes.

ClimaComms.abort(ctx::AbstractCommsContext, status::Int)

Terminate the caller and all participating processes with the specified exit status.

ClimaComms.barrier(ctx::AbstractCommsContext)

Perform a global synchronization across all participating processes: each process blocks until every process has reached the barrier.

ClimaComms.reduce(ctx::AbstractCommsContext, val, op)

Perform a reduction across all participating processes, using op as the reduction operator and val as this process's contribution. The result is only valid on the root process.

See also allreduce to make the result available on all processes.

ClimaComms.reduce!(ctx::AbstractCommsContext, sendbuf, recvbuf, op)
ClimaComms.reduce!(ctx::AbstractCommsContext, sendrecvbuf, op)

Perform an elementwise reduction across all participating processes, using op as the reduction operator and sendbuf as this process's contribution, and store the result in the root process's recvbuf. If a single sendrecvbuf buffer is provided, the reduction is performed in-place. Return nothing.

See also allreduce! to make the result available on all processes.

ClimaComms.allreduce(ctx::AbstractCommsContext, sendbuf, op)

Perform an elementwise reduction across all participating processes, using op as the reduction operator and sendbuf as this process's contribution, and return the result in a newly allocated array on every process. sendbuf can also be a scalar, in which case the result is a value of the same type.

ClimaComms.allreduce!(ctx::AbstractCommsContext, sendbuf, recvbuf, op)
ClimaComms.allreduce!(ctx::AbstractCommsContext, sendrecvbuf, op)

Perform an elementwise reduction across all participating processes, using op as the reduction operator and sendbuf as this process's contribution, and store the result in the recvbuf of every process. If a single sendrecvbuf buffer is provided, the reduction is performed in-place. Return nothing.

allreduce! is equivalent to reduce! followed by bcast, but can achieve better performance.

ClimaComms.bcast(ctx::AbstractCommsContext, object)

Broadcast object from the root process to all other processes, and return it on every process. The value of object on non-root processes is ignored.

ClimaComms.gather(ctx::AbstractCommsContext, array)

Gather an array from each participating process into a single array on the root process, concatenating along the last dimension. The arrays must have the same size on every process except possibly in the last dimension. The result is only valid on the root process.

AbstractGraphContext

A context for exchanging data between neighboring processes in a graph, such as the ghost (halo) regions of a domain decomposition.

Construct with graph_context; exchange data with start, progress, and finish.

ClimaComms.graph_context(
    context::AbstractCommsContext,
    sendarray, sendlengths, sendpids,
    recvarray, recvlengths, recvpids,
)

Construct an AbstractGraphContext for exchanging neighbor data via a graph.

Arguments

• context: the communication context on which to construct the graph context.
• sendarray: array containing the data to send, ordered by destination process.
• sendlengths: list of the number of elements to send to each process in sendpids.
• sendpids: list of process IDs to send to.
• recvarray: array to receive data into, ordered by source process.
• recvlengths: list of the number of elements to receive from each process in recvpids.
• recvpids: list of process IDs to receive from.

Notes

For MPICommsContext, the keyword argument persistent = true selects persistent MPI send/receive requests instead of MPI.Isend / MPI.Irecv!, which reduces the overhead of repeated exchanges.

ClimaComms.start(ctx::AbstractGraphContext)

Initiate the graph data exchange: post the receives and sends for the data currently in the send buffers.

ClimaComms.progress(ctx::AbstractGraphContext)

Drive communication. Call after start to ensure that communication proceeds asynchronously while other work is performed.

ClimaComms.finish(ctx::AbstractGraphContext)

Complete the communication step begun by start. After this returns, the data received from all neighbors is available in the receive buffers.

OnlyRootLogger()
OnlyRootLogger(ctx::AbstractCommsContext)

Return a logger that prints to the console on the root process and silences all other processes.

If no context is passed, obtain the default context via context. For MPI runs, this logger is installed as the global logger by the first call to init.

MPILogger(ctx::AbstractCommsContext)
MPILogger(iostream, ctx::AbstractCommsContext)

Return a logger that prefixes each log message with the process ID (e.g., [P1]). Output goes to stdout if no iostream is given.

Examples

using Logging
logger = ClimaComms.MPILogger(ClimaComms.context())
global_logger(logger)
@info "Hello"   # prints "[P1]  Info: Hello" on the root process

FileLogger(ctx, log_dir; log_stdout = true, min_level = Logging.Info)

Return a logger that writes each process's log messages to a separate file in log_dir (rank_1.log, rank_2.log, ...), with log_dir/output.log a symbolic link to the root process's file. For single-process runs, all messages go directly to log_dir/output.log.

Keyword Arguments

• log_stdout = true: if true, the root process also logs to stdout.
• min_level = Logging.Info: the minimum level a message must have to be logged.

Examples

using Logging
logger = ClimaComms.FileLogger(ClimaComms.context(), "logs")
with_logger(logger) do
    @info "Written to logs/output.log and stdout"
end

with_tempdir(f::Function, ctx::AbstractCommsContext)

Create a temporary directory on the root process, broadcast its path to all processes, and call f on the path. All processes receive the same path, so the directory can be used for files shared across the run.