# Large ODE Solver for Metal.jl

**URL:** https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769
**Category:** General Usage
**Tags:** question, gpu, ode, metaljl
**Created:** [December 27, 2025, 9:16pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769 "2025-12-27T21:16:45Z")
**Posts on this page:** 12
**Page:** 1

<div class="post-metadata">

### Author: ![Gavin-Rockwood](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gavin-rockwood/32/219194_2.png) [@Gavin-Rockwood](https://discourse.julialang.org/u/Gavin-Rockwood)
#### Post date: [December 27, 2025, 9:16pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/1 "2025-12-27T21:16:45Z")

</div>

Is there any implementation for ODE solving using Metal.jl backends? I know there is DiffEqGPU, but the metal support is only for ensemble problems. I cannot for instance do

```julia-auto
f(u,p,t) = A*u
tspan = (0.0f0, 1.0f0)
prob = ODEProblem(f, u0, tspan)
solve(prob, Tsit5())

```

or

```julia-auto
solve(prob, GPUTsit5)

```

where A, u, are Metal arrays. Tsit5() copies back and forth to CPU, and GPUTsit5 is designed for the ensemble solver. Am I just going about this in a wrong way or would I need to write my own solving algorithm to do this?

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [December 28, 2025, 10:56am UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/2 "2025-12-28T10:56:33Z")

</div>

> [@Gavin-Rockwood](#):
>
> Tsit5() copies back and forth to CPU

It does not. Where did you get that?

---

<div class="post-metadata">

### Author: ![Gavin-Rockwood](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gavin-rockwood/32/219194_2.png) [@Gavin-Rockwood](https://discourse.julialang.org/u/Gavin-Rockwood)
#### Post date: [December 28, 2025, 9:17pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/3 "2025-12-28T21:17:57Z")

</div>

If I set up the same problem, with `f(u,p,t) = A*u` as the function, the number of allocations when A, u are MtlArrays is almost 50x higher than the allocations when both are Float32 arrays (and the MtlArray takes significantly longer than the Float32 arrays). My assumption is that there are steps in the Tsit5 algorithm which aren’t handled properly in Metal, and maybe the data is getting saved in non Metal arrays and is getting copied back and forth? I’m not sure what other reason there would be for such high memory allocations. It is not just Tsit5 either, the vern solvers also have the same problem and I’m guessing that all the other solvers will as well.

I could be totally wrong, and the issue is elsewhere. I just don’t have any idea of what I would be.

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [December 28, 2025, 9:24pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/4 "2025-12-28T21:24:46Z")

</div>

What about `f(du,u,p,t) = mul!(du,A,u)`? You’re using the allocating path instead of the non-allocating path.

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [December 28, 2025, 9:25pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/5 "2025-12-28T21:25:10Z")

</div>

> [@Gavin-Rockwood](#):
>
> (and the MtlArray takes significantly longer than the Float32 arrays)

MtlArray operations are generally pretty slow until things get large. Are you testing 50,000x50,000 matrices?

---

<div class="post-metadata">

### Author: ![Gavin-Rockwood](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gavin-rockwood/32/219194_2.png) [@Gavin-Rockwood](https://discourse.julialang.org/u/Gavin-Rockwood)
#### Post date: [December 28, 2025, 9:38pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/6 "2025-12-28T21:38:38Z")

</div>

I just gave that a try and it still has very large allocations and is orders of magnitude slower.  
Also no, I am testing with smaller matrices, the largest I did was a couple ~4000x4000 (it was becoming marginally equivalent). Is there a way to speed up the MtlArray operations or is this some optimization that needs to happen in Metal.jl or on apples end before it is fast for more intermediate scale problems?

Though even if the operations themselves are slower, I’m still a bit confused by the whole memory allocations thing. I feel like it should be comparable.

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [December 28, 2025, 9:59pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/7 "2025-12-28T21:59:00Z")

</div>

Metal operations are much slower at that size. See for example:

 ![image](https://global.discourse-cdn.com/julialang/original/3X/6/0/6099d27ab293ad756f9cf1d2f75c0229fc1274f3.png)

> <https://github.com/SciML/LinearSolve.jl/issues/357#issuecomment-1671779339>
>
> This is a thread for investigating changes to the LU defaults, based off of benc…hmarks like https://github.com/SciML/LinearSolve.jl/pull/356 .
> 
> (Note: there's a Mac-specific version 3 posts down)
> 
> \`\`\`julia
> using BenchmarkTools, Random, VectorizationBase
> using LinearAlgebra, LinearSolve, MKL\_jll
> nc = min(Int(VectorizationBase.num\_cores()), Threads.nthreads())
> BLAS.set\_num\_threads(nc)
> BenchmarkTools.DEFAULT\_PARAMETERS.seconds = 0.5
> 
> function luflop(m, n = m; innerflop = 2)
> sum(1:min(m, n)) do k
> invflop = 1
> scaleflop = isempty((k + 1):m) ? 0 : sum((k + 1):m)
> updateflop = isempty((k + 1):n) ? 0 :
> sum((k + 1):n) do j
> isempty((k + 1):m) ? 0 : sum((k + 1):m) do i
> innerflop
> end
> end
> invflop + scaleflop + updateflop
> end
> end
> 
> algs = \[LUFactorization(), GenericLUFactorization(), RFLUFactorization(), MKLLUFactorization(), FastLUFactorization(), SimpleLUFactorization()\]
> res = \[Float64\[\] for i in 1:length(algs)\]
> 
> ns = 4:8:500
> for i in 1:length(ns)
> n = ns\[i\]
> @info "$n × $n"
> rng = MersenneTwister(123)
> global A = rand(rng, n, n)
> global b = rand(rng, n)
> global u0= rand(rng, n)
>     
> for j in 1:length(algs)
> bt = @belapsed solve(prob, $(algs\[j\])).u setup=(prob = LinearProblem(copy(A), copy(b); u0 = copy(u0), alias\_A=true, alias\_b=true))
> push!(res\[j\], luflop(n) / bt / 1e9)
> end
> end
> 
> using Plots
> \_\_parameterless\_type(T) = Base.typename(T).wrapper
> parameterless\_type(x) = \_\_parameterless\_type(typeof(x))
> parameterless\_type(::Type{T}) where {T} = \_\_parameterless\_type(T)
> 
> p = plot(ns, res\[1\]; ylabel = "GFLOPs", xlabel = "N", title = "GFLOPs for NxN LU Factorization", label = string(Symbol(parameterless\_type(algs\[1\]))), legend=:outertopright)
> for i in 2:length(res)
> plot!(p, ns, res\[i\]; label = string(Symbol(parameterless\_type(algs\[i\]))))
> end
> p
> 
> savefig("lubench.png")
> savefig("lubench.pdf")
> \`\`\`
> 
> \[lubench.pdf\](https://github.com/SciML/LinearSolve.jl/files/12287070/lubench.pdf)
> !\[lubench\](https://github.com/SciML/LinearSolve.jl/assets/1814174/84d091a9-f9f6-4e5e-a208-0a1c65365cec)
> 
> The justification for RecursiveFactorization.jl still looks very strong from the looks of this.
> 
> \`\`\`
> julia\> versioninfo()
> Julia Version 1.9.1
> Commit 147bdf428c (2023-06-07 08:27 UTC)
> Platform Info:
> OS: Windows (x86\_64-w64-mingw32)
> CPU: 32 × AMD Ryzen 9 5950X 16-Core Processor
> WORD\_SIZE: 64
> LIBM: libopenlibm
> LLVM: libLLVM-14.0.6 (ORCJIT, znver3)
> Threads: 32 on 32 virtual cores
> Environment:
> JULIA\_IMAGE\_THREADS = 1
> JULIA\_EDITOR = code
> JULIA\_NUM\_THREADS = 32
> \`\`\`
> 
> Needs examples on other systems.

> [@Gavin-Rockwood](#):
>
> Though even if the operations themselves are slower, I’m still a bit confused by the whole memory allocations thing. I feel like it should be comparable.

It should be, it’s worth looking into. But first make it in-place like I showed, and then it should be not allocating in the steps.

---

<div class="post-metadata">

### Author: ![Gavin-Rockwood](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gavin-rockwood/32/219194_2.png) [@Gavin-Rockwood](https://discourse.julialang.org/u/Gavin-Rockwood)
#### Post date: [December 28, 2025, 10:08pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/8 "2025-12-28T22:08:30Z")

</div>

Ah, interesting. I have seen that before but I wasn’t sure how much that would translate into ode solving.

I did try the in-place, the number of allocations went down, but nowhere near as much as it did for the CPU version. In fact, the gap in allocations got worse actually, ~120x more for the MtlArray.

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [December 28, 2025, 10:19pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/9 "2025-12-28T22:19:09Z")

</div>

What does the allocations profiler say the source is?

---

<div class="post-metadata">

### Author: ![Gavin-Rockwood](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gavin-rockwood/32/219194_2.png) [@Gavin-Rockwood](https://discourse.julialang.org/u/Gavin-Rockwood)
#### Post date: [December 28, 2025, 10:44pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/10 "2025-12-28T22:44:19Z")

</div>

Oh, uh, I have just been using `@benchmark`. I’m not very familiar with the memory allocation profiler. That might take some time for me to figure out how to interpret.

---

<div class="post-metadata">

### Author: ![Gavin-Rockwood](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gavin-rockwood/32/219194_2.png) [@Gavin-Rockwood](https://discourse.julialang.org/u/Gavin-Rockwood)
#### Post date: [December 28, 2025, 11:10pm UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/11 "2025-12-28T23:10:31Z")

</div>

![image](https://global.discourse-cdn.com/julialang/original/3X/0/7/074469bf92dfdb1d85e484e889bb6c0028927a2e.png)  
Is this the right thing to look at? (following [[Profiling · The Julia Language](https://docs.julialang.org/en/v1/manual/profile/#Memory-allocation-analysis)] using PProf.)

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [December 29, 2025, 2:49am UTC](https://discourse.julialang.org/t/large-ode-solver-for-metal-jl/134769/12 "2025-12-29T02:49:29Z")

</div>

Use the VS Code profiler to get it into a flamegraph? Usually I find that easier to find the lines of code that it can be attributed to.
