[xtalax]

Status: ready for review

• Tests: CI is green (88/88 after rerunning PDESystemLibrary Core, which master had since fixed). Locally on current master, GROUP=GPU passes 17/17 and Sol_Interface 26/26 (Julia 1.12.6, RTX 4080 SUPER).
• Wrapper build time is addressed separately by Build PDE solution wrappers array-wise #718.
• Device parameters: host_solution copies the states and leaves prob.p where it is. Moving device-resident parameters to the host is left out until a real case needs it (see "Not done").
• GPU tests: the GPU group stays opt-in and local-only (GROUP=GPU, not in "All" or test_groups.toml), because CI has no GPU runner.

Problem

After a GPU solve, the default wrap = Val(true) builds a PDETimeSeriesSolution. That constructor scalar-indexes the CuArray states, so it throws Scalar indexing is disallowed under CUDA.allowscalar(false). Until now the only workaround was solve(...; wrap = Val(false)).

Root cause

PDETimeSeriesSolution(sol, ::MOLMetadata) in src/interface/solution/timedep.jl builds each grid point's time series from the raw solution states, and it does so in three ways, all of which touch device memory:

1. Unknowns: sol[i, :] does a per-element VectorOfArray getindex on every saved state.
2. Unknowns that mtkcompile eliminated: SciMLBase.observed(sol, sym, :) runs the generated observed function on each state.
3. Array observed equations: array_observed_solution evaluates the observed function on state_values(sol), then does getindex.(values, (I,)) per element.

PDENoTimeSolution in timeindep.jl has the same pattern, using sol.u[i] and SciMLBase.observed(sol, sym).

Fix

A new host_solution(sol) in solution_utils.jl copies non-Array states to the host once, using SciMLBase.build_solution(sol.prob, sol.alg, sol.t, Array.(sol.u); retcode, calculate_error = false). The NonlinearSolution method works the same way with Array(sol.u). Both constructors run all extraction and observed evaluation on that host copy. They still store the original sol as original_sol and take every other field from it. The rebuilt solution keeps prob, so observed functions, parameters and the symbolic container are the same objects as before.

When the states are already Arrays, host_solution returns sol itself. The CPU path does exactly what it did before and makes no copy. A new assertion in Sol_Interface pins that (host_solution(sol.original_sol) === sol.original_sol).

Tests

New opt-in GPU group

The GPU tests live in test/GPU, with their own Project.toml (CUDA, NonlinearSolve, and so on), wired into runtests.jl the same way as NeuralNets. The group is not part of "All" and is not listed in test_groups.toml, so CI never runs it. Run it with GROUP=GPU. If CUDA is not functional, it records a @test_skip.

The test file sets CUDA.allowscalar(false) for the whole file. It uses a 1D heat equation with 41 grid points, above the 16-element scalar fill_arr! threshold. Each case compares the GPU wrap with the host wrap of the same states, using exact equality (==):

Case	How the device states were produced	Paths exercised	max |GPU wrap − CPU wrap|
ODEProblem from mtkcompile	Real GPU solve: Tsit5 on CuArray u0, run under CUDA.@allowscalar	Unknowns, plus the Dirichlet boundary points eliminated into observed equations	0.0
DAEProblem from default discretize	Solved on the CPU, states mapped to CuArray with build_solution	Unknowns	0.0
NonlinearProblem (1D Poisson)	Solved on the CPU, u mapped to CuArray	PDENoTimeSolution, unknowns and observed	0.0
Array observed equation (reduced-model field w = basis * reduced)	Hand-built CuArray states	array_observed_solution	0.0

The real GPU solve needs scalar indexing for two reasons:

• The MTK-generated RHS writes element by element. This is the known codegen issue (MTK#5242 and boundary equations).
• The default AutoDespecialize function wrapper and MTK initialization convert u0 back to a Vector. To keep the states on the device, the test builds an ODEFunction{true, FullSpecialize} from the compiled function and solves with initializealg = NoInit().

So only the wrapping runs under allowscalar(false). A GPU solve of the DAEProblem fails inside the solver with a GPU kernel compilation error, which is unrelated to this change. That is why the DAE case maps CPU states to the device instead.

Against unpatched master, all four cases error: three with Scalar indexing is disallowed, and the array-observed case with a Matrix * CuArray BLAS MethodError. With the patch, all 17 assertions pass.

Test group results (local, Julia 1.12.6)

Group	Result
Sol_Interface	20/20 pass
Components	444 pass, 4 broken (the existing @test_broken in utils_test.jl)
DAE	22/22 pass
Stationary	8/8 pass
GPU (RTX 4080 SUPER, CUDA.jl 6.4.2)	17/17 pass

Runic --check . exits 0.

Not done

• Device parameters. host_solution copies the states but not prob.p. MOL's MTKParameters are host arrays, and MTK has no Adapt rule for them. If someone builds a problem with device-resident tunables, the observed evaluation would still scalar-index them. I left that alone rather than guess at a conversion with no real case to test it against.
• No change to SciMLBase wrap_sol. The fix is entirely in MOL's constructors.
• No GPU support for the solve itself. Making the MTK-generated RHS GPU-clean is upstream work (MTK#5242).

Note: the DAE case uses mapped states because a GPU DAE solve currently fails inside the solver. That failure comes from ModelingToolkit's scalarized du reads in the DAE residual, which produce host array_literals inside GPU broadcasts. It is fixed by the draft SciML/ModelingToolkit.jl#5242, not by MethodOfLines.

🤖 Posted by an AI agent · harness: Claude Code 2.1.296 · model: claude-opus-5-5
Review: unreviewed
Session: 644ed0ff-c8c9-4cc3-b4c6-f326231a4030

[xtalax]

1 parameters also belong on gpu
2 local only for now