GPU flows

CTFlows can integrate a flow whose state lives on a GPU device: you pass CuArrays instead of Arrays, and — for the flows that need automatic differentiation — you opt in with method=:gpu.

This page covers what works, what it costs, and the one thing that is not optional: the rules your own functions must follow to run on a device.

Backends

Everything below was validated on NVIDIA hardware (CUDA.jl). The device-handling code in CTFlows dispatches on GPUArraysCore.AbstractGPUArray, so it is written to be backend-agnostic, but AMDGPU and Metal are untested.

---

When a GPU helps — and when it does not

There are two very different ways to put an ODE on a GPU, and they have opposite performance profiles.

	Array-level	Ensembles
What is parallel	the operations inside one solve	many independent trajectories
State size it suits	large (≳ 10⁴ components)	small
A typical optimal-control flow (4–20 states)	slower than CPU	the case it is designed for
Status in CTFlows	supported, see below	tooling validated, no Flow API yet — see GPU ensembles

For a state of a handful of components, an array-level GPU solve is slower than the CPU one: kernel-launch latency dominates the arithmetic. That is expected and not a defect.

So why use it? Composability. An array-level GPU flow is what lets a flow sit inside a computation that is already device-resident — a shooting function whose other pieces live on the GPU, or a GPU solve driven by another package — without a round-trip to host memory at every step. Use it for correctness and composition, not to accelerate a small isolated flow.

If what you want is raw throughput over many trajectories (multiple shooting, batched initial guesses, sweeps over a variable), that is the ensemble model — see GPU ensembles.

---

Quickstart

Move the state to the device and call the flow as usual.

using CTFlows
using CTFlows: Flows
using CTBase: Data
using OrdinaryDiffEqTsit5: OrdinaryDiffEqTsit5
using CUDA: CUDA

x0 = CUDA.CuArray([1.0, 2.0])

flow = Flows.Flow(Data.VectorField(x -> -x))   # ẋ = -x
xf = flow(0.0, x0, 1.0)                        # a CuArray: [1, 2] · e⁻¹

No keyword was needed there. For a flow that differentiates your function — anything built from a Hamiltonian, or from an optimal control problem with a costate — add method=:gpu:

h = Data.Hamiltonian((x, p) -> (sum(abs2, x) + sum(abs2, p)) / 2)

flow = Flows.Flow(h; method=:gpu)
xf, pf = flow(0.0, CUDA.CuArray([1.0, 0.0]), CUDA.CuArray([0.0, 1.0]), 1.0)

What method=:gpu actually does

It is not a hardware switch. It tells CTFlows to resolve its strategies in their GPU variants — the integrator (SciML{GPU}) and the AD backend (DifferentiationInterface{GPU}, which selects a GPU-capable backend, see AD backends on device).

Whether your arrays are on the device is decided by you, by what you pass in. That is why an AD-free flow needs no keyword at all: handing it a CuArray is enough.

The rule of thumb: AD-free flow → just pass device arrays. AD-dependent flow → also pass method=:gpu.

The keyword is accepted by every Flow constructor and can be combined freely with the usual integrator options:

flow = Flows.Flow(h; method=:gpu, reltol=1e-10, abstol=1e-10)

method=:gpu needs device arrays

A flow built with method=:gpu refuses a host array (a Vector, a Matrix…): the call throws a PreconditionError before the integration starts, instead of failing deep inside the AD backend. A 1-D state, which is a scalar, is not an array and is accepted. To integrate host arrays, build the flow without method=:gpu (CPU).

Load the packages the backend needs

The AD backend is a marker: its package must be loaded. For method=:gpu the default backend is AutoMooncake(), so run using Mooncake (and using CUDA) before calling the flow; on CPU the default AutoForwardDiff() needs using ForwardDiff. A missing package gives an ExtensionError that names it (Missing Mooncake, Hint: using Mooncake) at the call, before the integration.

---

The GPU contract for your code

This is the part that breaks in practice. CTFlows itself is device-clean; the GPU-safety burden sits almost entirely in the functions you supply — dynamics, Hamiltonians, control laws, and cost terms.

Rule: no scalar indexing

Reading one element of a device array (x[1], v[i]) is a host↔device round-trip. CUDA.jl forbids it under CUDA.allowscalar(false), which is the setting CTFlows' own GPU tests run under — and the setting you should develop against.

Write with reductions and broadcasts instead:

✗ Fails on device	✓ Device-safe
x[1]^2 + x[2]^2	sum(abs2, x)
p[1]*x[1] + p[2]*x[2]	dot(p, x) (using LinearAlgebra)
r[1] = -x[1]; r[2] = -x[2]	r .= .-x
xf[1] in a Mayer term	sum(xf), or any reduction over xf
x[1] * x[2]	prod(x), or restructure the formula

Reductions are the escape hatch: anything you can phrase as sum, prod, dot, norm, mapreduce runs as a device kernel and returns a host scalar you can go on computing with. Genuinely component-wise formulas — a dynamics that couples specific components — have no device-safe rewrite at the flow level; that is a real limitation, not a style preference.

The variable v is a host object

v[1] is fine: the variable stays on the host by design (see Limits and boundaries). The rule above concerns the state, the costate, and the dynamics buffer.

The reassuring part: GPU-safe style is also perfectly good CPU code. The same Hamiltonian runs on both:

# reductions only — no x[1], no p[2]
h = Data.Hamiltonian((x, p) -> (sum(abs2, x) + sum(abs2, p)) / 2)

flow = Flows.Flow(h)                        # CPU here; identical code with method=:gpu
flow(0.0, [1.0, 0.0], [0.0, 1.0], 1.0)      # harmonic oscillator at t = 1

([0.5403023057535162, 0.8414709847374632], [-0.8414709847374632, 0.5403023057535162])

The contract covers every closure of an OCP

Not just the dynamics. A trajectory call evaluates the Mayer and Lagrange terms on the flow's own endpoints and states — which, on a GPU flow, are device arrays. A Mayer term written the canonical way scalar-indexes them and fails:

# ✗ xf is a CuArray here — xf[1] is a scalar read
CTModels.Building.objective!(pre, :min; mayer=(x0, xf, v) -> xf[1])

# ✓ a reduction stays on device
CTModels.Building.objective!(pre, :min; mayer=(x0, xf, v) -> sum(xf))

The same applies to in-place dynamics: the buffer you are handed follows the state's array type, so write it with a broadcast.

# ✗ component-wise writes into a device buffer
CTModels.Building.dynamics!(pre, (r, t, x, u, v) -> (r[1] = -x[1]; r[2] = -x[2]; nothing))

# ✓ broadcast
CTModels.Building.dynamics!(pre, (r, t, x, u, v) -> (r .= .-x; nothing))

The canonical r[1] = … idiom is CPU-only

Most optimal-control examples in the ecosystem write dynamics component by component. That style is fine on CPU and cannot run on a device. Porting a problem to GPU means rewriting those closures in broadcast form.

Why the burden is yours

CTModels never scalar-indexes the dynamics buffer — it only ever hands your closure a contiguous @view of it, which is a valid lazy view on a CuArray. The model layer is structurally device-clean, so if a device run fails, look at your own closures first.

---

What runs on device

Every constructor below has been executed on an NVIDIA device under allowscalar(false):

Flow source	AD needed	Device call
Flow(::VectorField)	no	pass CuArrays
Flow(::HamiltonianVectorField)	no	pass CuArrays
Flow(::ODEFunction), Flow(::ODEProblem)	no	pass CuArrays
Flow(ocp) — state flow (no costate)	no	pass CuArrays
Flow(ocp, ::OpenLoop/::ClosedLoop)	no	pass CuArrays
Flow(::Hamiltonian)	yes	method=:gpu
Flow(ocp) — Hamiltonian (with costate)	yes	method=:gpu
Flow(ocp, ::DynClosedLoop)	yes	method=:gpu
multi-phase concatenation (*), with jumps	inherits	works on device
variable_costate=true	yes	works with a host or device variable

One caveat on the last row

The augmented variable_costate cases were validated on device with an explicit ad_backend=AutoZygote(), not with the method=:gpu default. Both are expected to work; only the former has actually been run on hardware.

Limits and boundaries

The variable v stays on the host. It parameterises the right-hand side and broadcasts cleanly against device arrays, so there is nothing to gain from moving it. A device variable is accepted too, but host is the intended usage.

Materialisation pulls back to host. A point call flow(t0, x0, tf) returns device arrays — the hot path stays on the device. Building a full CTModels.Solution, plotting, and grid interpolation gather to the host at the boundary, because everything downstream of them (plot backends, printing, interpolation) is host code by nature.

The 1-D = scalar convention costs a synchronisation. CTFlows represents a one-dimensional state as a scalar, not a 1-vector. Collapsing a length-1 device array to a scalar is by definition a scalar read; it is done safely, but it is a host synchronisation. GPU flows are meant for vector states — a 1-D problem on a device is correct, not fast.

Float32 is preserved end-to-end. GPUs prefer single precision and nothing in the chain silently promotes to Float64:

flow = Flows.Flow(Data.VectorField(x -> -x))
xf = flow(0.0f0, CUDA.CuArray(Float32[1, 2]), 1.0f0)   # eltype(xf) == Float32

---

AD backends on device

The default AD backend on CPU is AutoForwardDiff(). It cannot run on a device: forward mode reads the array element by element, which is exactly the scalar indexing a device array rejects. This is why AD-dependent flows need method=:gpu — it swaps in a GPU-capable default.

That default is AutoMooncake(), and the reason it wins is specific: an OCP Hamiltonian fills its buffer through CTModels' in-place dynamics!, and most reverse-mode backends refuse to differentiate array mutation. Mooncake differentiates it, on device, end to end.

The registry carries this as live help text:

Strategies.describe(:di, Flows.flow_registry())

DifferentiationInterface (strategy)
├─ id: :di
├─ hierarchy: DifferentiationInterface → AbstractADBackend → AbstractStrategy
├─ description: AD backend wrapping DifferentiationInterface.jl backends (e.g., AutoForwardDiff).
├─ family: AbstractADBackend
├─ default parameter: CPU
├─ parameters: CPU, GPU
│
├─ computed options for CPU:
│  └─ ad_backend (backend, ad)::ADTypes.AbstractADType (default: AutoForwardDiff() [computed])
│     description: DifferentiationInterface.jl backend (e.g. AutoForwardDiff() on CPU).
│
└─ computed options for GPU:
   └─ ad_backend (backend, ad)::ADTypes.AbstractADType (default: AutoMooncake() [computed])
      description: DifferentiationInterface.jl backend for GPU execution. Default: AutoMooncake(), validated end-to-end on CuArray including through a mutating in-place RHS. AutoForwardDiff() does not work on GPU (scalar-indexes a CuArray). AutoZygote() can be selected explicitly; it was found unreliable in some device call contexts (an unexpected KernelException was observed on a non-mutating call).

You can always override the backend explicitly:

using ADTypes: ADTypes
using Zygote: Zygote   # the package must be loaded: AutoZygote() is only a marker type
flow = Flows.Flow(h; ad_backend=ADTypes.AutoZygote())

AutoZygote on device is not reliably predictable

Zygote works on device in many cases — but its behaviour was found to depend on the call context, not just on the function: the same Hamiltonian formula that drives a full device integration successfully raised a KernelException when its gradient was taken directly, on a call that performs no mutation at all. It remains selectable, and is a reasonable thing to try; it is not something to rely on without testing your own case.

AutoMooncake on CPU: for large dimensions

ad_backend=AutoMooncake() (with using Mooncake) also works on host arrays, on every AD-backed flow, with the same values as AutoForwardDiff(). It pays off only for large states: for the Hamiltonian  , a full flow call took

n	AutoForwardDiff	AutoMooncake
30	0.45 ms	2.5 ms
100	3.0 ms	3.3 ms
400	39 ms	7.5 ms
1000	231 ms	16 ms

so AutoForwardDiff() stays the CPU default. Mooncake does not compose with an outer ForwardDiff: differentiating a Mooncake flow with ForwardDiff (with respect to the initial condition, the final time or the variable, as a shooting method does) throws a PreconditionError. Keep AutoForwardDiff() for shooting, or differentiate the flow with finite differences. Details in #448.

---

GPU ensembles

No CTFlows API yet

Solving many trajectories in parallel is where GPUs actually pay off for optimal control, but CTFlows does not wrap it yet — there is no batched Flow call. What follows is the raw SciML pattern, validated on device, that you can use directly today. A user-facing batch API is future work.

Two algorithms from DiffEqGPU.jl, with different requirements on the right-hand side:

using CUDA: CUDA
using DiffEqGPU: DiffEqGPU
using SciMLBase: SciMLBase
using OrdinaryDiffEqTsit5: OrdinaryDiffEqTsit5
using StaticArrays: StaticArrays

N = 8
init(i) = [Float64(i), 2.0 * i]   # trajectory i starts at [i, 2i]

# --- EnsembleGPUArray: in-place RHS, host arrays moved by the algorithm ---
prob = SciMLBase.ODEProblem((du, u, p, t) -> (du .= .-u), init(1), (0.0, 1.0))
eprob = SciMLBase.EnsembleProblem(
    prob; prob_func=(p, ctx) -> SciMLBase.remake(p; u0=init(ctx.sim_id))
)
sol = SciMLBase.solve(
    eprob, OrdinaryDiffEqTsit5.Tsit5(), DiffEqGPU.EnsembleGPUArray(CUDA.CUDABackend());
    trajectories=N,
)

# --- EnsembleGPUKernel: out-of-place SVector RHS, no host closures ---
sprob = SciMLBase.ODEProblem(
    (u, p, t) -> -u, StaticArrays.SVector{2}(init(1)), (0.0, 1.0)
)
seprob = SciMLBase.EnsembleProblem(
    sprob;
    prob_func=(p, ctx) -> SciMLBase.remake(
        p; u0=StaticArrays.SVector{2}(init(ctx.sim_id)),
    ),
)
ssol = SciMLBase.solve(
    seprob, DiffEqGPU.GPUTsit5(), DiffEqGPU.EnsembleGPUKernel(CUDA.CUDABackend());
    trajectories=N,
)

⚠️ The scalar-indexing rule inverts here

This is the single most confusing thing about the two models, and getting it backwards costs hours:

	Flow code (everything above)	EnsembleGPUArray right-hand side
Where it runs	on the host, operating on device arrays	inside a compiled kernel
Governed by	CUDA.allowscalar(false)	what the GPU compiler accepts
du[1] = u[2]	✗ forbidden	✓ correct idiom
du .= A * u	✓ fine	✗ InvalidIRError — BLAS and heap allocation do not exist in a kernel

Inside an ensemble kernel, du and u are device-array slices and component-wise access is ordinary device code. What you must avoid there is array-level linear algebra: a matrix–vector product routes through BLAS and allocates, neither of which a kernel can do. Write the coupled system out component by component.

Two more measured facts

Results are not device-resident. Both algorithms gather to the host before returning — EnsembleGPUArray yields a view into a host Matrix, EnsembleGPUKernel a host SVector. Do not assume you can chain another device computation onto the result without moving it back.

Read final states through .u. Use sol.u[i].u[end], never sol[i][end]: the latter is Cartesian indexing into the (component × time) shape and silently returns a scalar rather than the final state vector — a bug that looks like a numerical error.

---

Testing your own GPU code

The pattern CTFlows uses for its own device tests transfers directly:

using CUDA: CUDA
CUDA.functional() || return          # skip cleanly on machines without a device
CUDA.allowscalar(false)              # make every pass genuinely scalar-index-free

Setting allowscalar(false) is what turns "it ran" into "it ran without secretly synchronising element by element". Develop against it from the start.

---

See also

• GPU internals — design decisions, the routing path, the coercion helpers, and how the device tests are run in CI.

• Integrating — the solver options method=:gpu composes with.

• SciML flows — Flow(::ODEFunction) / Flow(::ODEProblem), both usable on device.

• Optimal control, Control laws — the OCP constructors whose closures the GPU contract applies to.

• CTBase.Strategies — the CPU/GPU strategy parameters behind method=:gpu.