# Minimising DifferentialEquations GPU memory allocation

**URL:** <https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050>\
**Category:** Modelling & Simulations\
**Tags:** question\
**Created:** [May 18, 2023, 9:45am UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050 "2023-05-18T09:45:57Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![Tommy\_Fischer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tommy_fischer/32/47412_2.png) [@Tommy\_Fischer](https://discourse.julialang.org/u/Tommy_Fischer)\
**Post date:** [May 18, 2023, 9:45am UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/1 "2023-05-18T09:45:57Z")

</div>

I’m simulating the Gross-Pitaevskii Equation for a three-dimensional array with M^3 values, and am running into some issues with GPU memory for large M values. In short, I’m wondering if there is anything more I can do to reduce how much memory solve() is allocating?

Currently I’m using CUDA.jl so that the dynamics are solved on the GPU, and am using a callback to intermittently save data back to CPU directly (see MWE below). I’ve managed to push N to ~ 300, but anything much bigger than that and it runs out of memory very quickly.

```julia
using FFTW, CUDA, DifferentialEquations, LinearAlgebra, Plots

    function kfunc_opt!(dψ,ψ)
        mul!(dψ,Pf,ψ)
        dψ .*= k2
        Pi!*dψ
        return nothing
    end

    function GPE!(dψ,ψ,var,t) # GPE Equation 
        kfunc_opt!(dψ,ψ)
        @. dψ = -(im + γ)*(0.5*dψ + (V_0 + abs2(ψ) - 1)*ψ)
    end

    function GPU_Solve(EQ!, ψ, tspan) 
        
        savepoints = tspan[2:end]  
        condition(u, t, integrator) = t ∈ savepoints
    
        function affect!(integrator)                 
            i += 1
            ret[:,:,:,i] .= Array(integrator.u)
        end
    
        ret = zeros(M,M,M,length(tspan)) .|> complex 
        ret[:,:,:,1] .= Array(ψ)
    
        cb = DiscreteCallback(condition, affect!)   
        i = 1                                           
        
        prob = ODEProblem(EQ!,ψ,(tspan[1],tspan[end]))   
        solve(prob, callback=cb, tstops = savepoints, save_start = false, save_everystep = false, save_end = false)
   
        return ret
    end

    L = 8
    M = 60
    
    x = LinRange(-L,L,M) |> cu;
    dx = x[2] - x[1]
    kx = fftfreq(M,2π/dx) |> collect |> cu;
    dkx = kx[2] - kx[1]

    k2 = kx.^2 .+ kx'.^2 .+ reshape(kx,(1,1,M)).^2;
    V_0 = 0.3*[i^2 + j^2 + k^2 for i in x, j in x, k in x] |> cu;

    const Pf = Float32(dx^3/(2π)^1.5)*plan_fft(cu((rand(M,M,M) + im*rand(M,M,M))));
    const Pi! = Float32(M^3*dkx^3/(2π)^1.5)*plan_ifft!(cu((rand(M,M,M) + im*rand(M,M,M))));

γ = 1
tspan = LinRange(0.0,3,30); 

CUDA.memory_status()
res_GS = GPU_Solve(GPE!,cu((randn(M,M,M) + im*randn(M,M,M))),tspan);

begin
    t = 7 # Change this to look at different times
    heatmap(abs2.(res_GS[:,:,30,t]))
end

```

If any other info is required, or if there are any other glaring issues with my code please let me know. Cheers.

---

<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:** [May 18, 2023, 10:07am UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/2 "2023-05-18T10:07:00Z")

</div>

Is the equation stiff? If you don’t need the Jacobian then you can specialize on that, and if you do the need the Jacobian you should declare its sparsity pattern.

---

<div class="post-metadata">

**Author:** ![rveltz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rveltz/32/2707_2.png) [@rveltz](https://discourse.julialang.org/u/rveltz)\
**Post date:** [May 18, 2023, 11:57am UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/3 "2023-05-18T11:57:15Z")

</div>

Dont you have globals? Pf and Pi!

---

<div class="post-metadata">

**Author:** ![Tommy\_Fischer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tommy_fischer/32/47412_2.png) [@Tommy\_Fischer](https://discourse.julialang.org/u/Tommy_Fischer)\
**Post date:** [May 18, 2023, 9:11pm UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/4 "2023-05-18T21:11:25Z")

</div>

My bad, have fixed the MWE and made sure it’s actually running properly. Let me know if it’s still a bit unclear

---

<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:** [May 18, 2023, 10:06pm UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/5 "2023-05-18T22:06:08Z")

</div>

The biggest allocation is still going to be the Jacobian

---

<div class="post-metadata">

**Author:** ![Tommy\_Fischer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tommy_fischer/32/47412_2.png) [@Tommy\_Fischer](https://discourse.julialang.org/u/Tommy_Fischer)\
**Post date:** [May 18, 2023, 10:20pm UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/6 "2023-05-18T22:20:10Z")

</div>

I’m not 100% sure if it’s stiff or not, the main (dimensionless) equation of motion is:

\frac{\partial \psi (\textbf{r},t)}{\partial t} = -\frac{i}{h} \left[\frac{\nabla^2}{2} + V\_0(\textbf{r}) + | \psi |^2 \right] \psi

Where V\_0 is just a number at each point. Would this be stiff because of the -|\psi|^2 \psi term?

---

<div class="post-metadata">

**Author:** ![apo383](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/apo383/32/11272_2.png) [@apo383](https://discourse.julialang.org/u/apo383)\
**Post date:** [May 19, 2023, 5:04am UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/7 "2023-05-19T05:04:04Z")

</div>

> [@Tommy\_Fischer](#):
>
> ```julia
> Pi!*dψ
> return nothing
> 
> ```

Is that a typo? Doesn’t make sense to multiply by `Pi! ` without doing anything with it.

---

<div class="post-metadata">

**Author:** ![Tommy\_Fischer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tommy_fischer/32/47412_2.png) [@Tommy\_Fischer](https://discourse.julialang.org/u/Tommy_Fischer)\
**Post date:** [May 19, 2023, 5:51am UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/8 "2023-05-19T05:51:50Z")

</div>

Nah not a typo, Pi!\*d\psi does an in-place inverse fourier transform on d\psi

---

<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:** [May 19, 2023, 10:08am UTC](https://discourse.julialang.org/t/minimising-differentialequations-gpu-memory-allocation/99050/9 "2023-05-19T10:08:46Z")

</div>

Just try choosing an explicit RK method like Tsit5