cpp-dispatch — dispatching across devices

Modern compute nodes are heterogeneous: one or more CPUs alongside one or more GPUs. The work scheduler / dispatcher lets a single compiled model spread a large batched evaluation across devices — slice the batch into sub-batches, run each on a device, and stitch the results back together.

This is a C++ runtime feature: it serves the compiled (AOTI packages) path embedded in a host application (e.g. MOOSE), where the hot loop runs without Python. The Python authoring path (neml2.load_model, neml2-run) stays eager and single-device — to spread work from Python, run your own per-device loop.

Note

Only CPU and CUDA devices are supported. Two scheduling modes are available: synchronous single-device (SimpleScheduler, MPISimpleScheduler) and asynchronous multi-device (StaticHybridScheduler, which runs CPU + GPU(s) concurrently via a thread-per-device pool).

Compile one artifact per device

AOTI graphs are pinned to a device at export time, so a dispatcher that targets several devices needs one artifact per device. neml2-compile takes multiple --device values and emits them side by side:

Listing 6 tests/aoti/forward_single/model.i — the model to compile
# forward_single: a one-leaf forward-only model. Smoke test for the
# bake-by-default path on the simplest possible shape.

[Models]
  [model]
    type = LinearIsotropicElasticity
    strain = 'strain'
    stress = 'stress'
    coefficients = '100 0.3'
    coefficient_types = 'YOUNGS_MODULUS POISSONS_RATIO'
  []
[]

$ neml2-compile tests/aoti/forward_single/model.i --model model --device cpu cuda -d :

(-d : compiles every Jacobian/JVP pair so the dispatched jacobian / jvp work; drop it for a forward-only artifact, or name specific pairs.) A batch-independent block is returned unbatched and the dispatcher passes it through unchanged (it is identical across batch chunks).

produces a standalone stub next to a per-device artifact folder:

aoti/
  model_aoti.i            # standalone stub; points at the folder via artifact_path
  model/
    cpu/   model_meta.json + *.pt2
    cuda/  model_meta.json + *.pt2

The loader resolves <artifact_path>/<device>/ for the device it runs on. The Python shim picks the subfolder matching torch.get_default_device() (so neml2-run --device cuda loads cuda/); the C++ loader picks it from the scheduler (below).

Load and dispatch from C++

neml2::aoti::load_model mirrors Python’s load_model(path, name) and returns a DispatchedModel — a Model-shaped handle exposing the same forward / jvp / jacobian plus the parameter derivatives param_jacobian / param_vjp (each chunked across devices and stitched back — Jacobian blocks concatenated along the batch, VJP adjoints summed). The optional scheduler is the dispatch opt-in; it is supplied in C++ source, never from the .i.

#include "neml2/csrc/dispatchers/factory.h"
#include "neml2/csrc/dispatchers/SimpleScheduler.h"

using namespace neml2::aoti;

// No scheduler -> no dispatch: runs the whole batch on cpu (zero-overhead
// pass-through over the underlying Model).
auto m = load_model("aoti/model_aoti.i", "model");
auto out = m.forward(inputs);

// With a scheduler -> the batch is chunked along its leading axis, each chunk
// moved to the compute device, run, and the results concatenated back on the
// input device.
auto sched = std::make_shared<SimpleScheduler>(SimpleScheduler::Config{"cuda", 4096});
auto md = load_model("aoti/model_aoti.i", "model", sched);
auto out_gpu = md.forward(inputs);   // cpu inputs -> gpu compute -> cpu results

When the scheduler’s device equals the input device and the whole batch fits in one chunk, DispatchedModel short-circuits to a direct Model call — so the no-dispatch case carries no slicing or transfer cost.

Schedulers

A scheduler decides which device(s) a workload runs on and how large each sub-batch chunk is. All are plain C++ objects configured by a Config struct. DispatchedModel picks its execution mode from the scheduler’s type: a synchronous scheduler (SimpleScheduler, MPISimpleScheduler) runs the chunk loop on the calling thread; an asynchronous one (StaticHybridScheduler) drives a thread-per-device pool.

SimpleScheduler

Sends the whole workload to a single device, chunked. Config{device, batch_size} (e.g. {"cuda:0", 1024}; batch_size = 0 means “no chunking”). Use it to:

  • process a batch too large for device memory in fixed-size pieces;

  • empirically tune the per-call batch size for a model + device; or

  • drive one device per process when a host pins devices by hand.

Note

Illustrative schematic: the batch is split into fixed-size chunks and fed to one device, a new chunk dispatched as in-flight ones finish and free capacity.

MPISimpleScheduler

For MPI jobs that drive several devices from many ranks. Config{devices, batch_sizes, comm} lists the devices to choose from (CPU or CUDA — e.g. {"cuda:0", "cuda:1"}, or {"cpu"} for a pure-CPU run); each rank is assigned one based on its rank within its node (ranks are grouped by hostname, then local_rank % devices.size() indexes into the list), after which it chunks exactly like SimpleScheduler. With m ranks on a node and n devices:

  • m == n — one device per rank;

  • m > n — round-robin, so a device serves several ranks;

  • m < n — an error: idle devices are not allowed (launch more ranks, or pass fewer devices).

comm is optional (nullptr ⇒ MPI_COMM_WORLD); to confine the scheduler to a subcommunicator, point it at your MPI_Comm (read only during construction):

MPI_Comm sub = /* ... */;
MPISimpleScheduler::Config cfg;
cfg.devices     = {"cuda:0", "cuda:1"};
cfg.batch_sizes = {4096};
cfg.comm        = &sub;   // omit for MPI_COMM_WORLD

The per-node split is collective over the chosen communicator, so every rank in it must construct the scheduler. Requires NEML2 built with -DNEML2_MPI=ON and the host to have called MPI_Init; otherwise the constructor throws.

StaticHybridScheduler

Spreads one batch across several devices concurrently — a single DispatchedModel runs CPU + GPU(s) at once via a thread-per-device pool. Config{devices, batch_sizes, capacities, priorities} (the last two optional; each broadcasts from length 1). Assignment is greedy: each chunk goes to the highest-priority device that still has spare capacity (load + batch_size <= capacity), so faster devices stay filled; capacity controls how many chunks may be in flight per device (overlapping the next chunk’s host→device copy with the current chunk’s compute).

#include "neml2/csrc/dispatchers/StaticHybridScheduler.h"

StaticHybridScheduler::Config cfg;
cfg.devices     = {"cpu", "cuda:0", "cuda:1"};
cfg.batch_sizes = {512, 4096, 4096};   // tune per device, e.g. via SimpleScheduler
auto m = load_model("aoti/model_aoti.i", "model",
                    std::make_shared<StaticHybridScheduler>(cfg));
auto out = m.forward(inputs);          // dispatched across all three, gathered back

Note

Illustrative schematic: each chunk goes to the highest-priority device that still has spare capacity, so faster devices stay filled; chunks process concurrently across devices and are gathered back as they finish.

A hybrid pool admits at most one CPU plus distinct GPUs: each device’s AOTI graph already saturates torch’s intra-op (OpenMP) thread pool, so two CPU workers would only oversubscribe the same cores.

Promoted parameters under hybrid. named_parameters() is a single master map; mutating it in place is broadcast to every device copy before the next dispatch, so the usual single-device idiom keeps working:

m.named_parameters().at("model.E").fill_(150e3);  // reflected on every device next call

Error handling

Every exception that leaves forward / jvp / jacobian — on both the synchronous and asynchronous paths — is a neml2::aoti::Exception (itself a std::runtime_error) carrying a recoverable() flag. That flag is the contract a downstream consumer branches on:

  • ConvergenceError (recoverable() == true) — the nonlinear solve diverged or hit its iteration cap. A time-stepping consumer can cut the step and retry.

  • FatalError (recoverable() == false) — a shape / device mismatch, a missing input, a malformed artifact. A retry would fail identically, so it must hard-fail. Foreign errors (a torch c10::Error, std::bad_alloc, …) are normalized to this at the boundary, so a single catch covers everything.

try
{
  auto out = m.forward(inputs);
}
catch (const neml2::aoti::Exception & e)
{
  if (e.recoverable()) { /* e.g. dt *= 0.5; retry */ }
  else                 { throw; } // fatal: give up
}

Under asynchronous dispatch this stays well-defined even when several chunks run at once. A failing chunk is caught inside its worker (a C++ exception escaping a std::thread would call std::terminate), the scheduler is still drained so the pool can never deadlock, and only then does the dispatcher decide what to throw:

  • one failure → it is re-thrown verbatim (its dynamic type, e.g. ConvergenceError, is preserved);

  • several at once → an AggregateError carrying them all. It reports recoverable() only if every sub-error is recoverable, so a lone fatal among otherwise-recoverable failures still forces a hard stop. The individual errors are available via AggregateError::errors().

Either way the DispatchedModel and its scheduler are left clean and reusable for the next call.

See also