Memory usage in DifferentialEquations.jl

General Usage
,
1

Hi ,

I am encountering a problem with DifferentialEquations.jl running out of memory when using threads. So I have created a function which is non-allocating and then used threads to speed up the process. The function is just copying terms from one matrix and then pasting it to another with some changes and hence ideally no race conditions. The matrix describes a field in which each element is governed by a differential equation. What ends up happening is the function with threads allocates much more memory and in my problem straight out runs out of memory than the one without threads. I am attaching a mwe which I think captures the essence of the problem.

using Parameters
using DifferentialEquations
using BenchmarkTools

ψ = fill(1.0+0im,20,20,2)
@with_kw struct Parameters_St
     lenx::Int64
     leny::Int64
     im::ComplexF64
end
para = Parameters_St(
    lenx=size(ψ,1),
    leny=size(ψ,2),
    im=1.0,
)

function f_mwethread!(du,u,p,t)
    @unpack lenx,leny= p
     Threads.@threads for j=1:lenx
        for k=1:2
            for i=1:leny
                du[i,j,k] = (exp(-10)-1)*(p.im*u[i,j,k]) #Linear terms
    #            du[i,j,k] += -p.im*u[i,j,k]  #nonlinear terms
             end
         end
     end
    return du
end

function f_mwe!(du,u,p,t)
    @unpack lenx,leny= p
     for i=1:lenx
         for k=1:2
             for j=1:leny
                 du[i,j,k] = (exp(-10)-1)*(p.im*u[i,j,k]) #Linear terms
    #             du[i,j,k] += -p.im*u[i,j,k]  #nonlinear terms
              end
          end
      end
    return du
end

imt=200000
tspan = (0.0,imt)
prob=ODEProblem(f_mwe!,ψ,tspan,para)
probwth=ODEProblem(f_mwethread!,ψ,tspan,para)


p=@benchmark solve($prob,BS3(),dense=false,progress=true,reltol=1e-10,abstol=1e-10,saveat=$(imt/1000))
q=@benchmark solve($probwth,BS3(),dense=false,progress=true,reltol=1e-10,abstol=1e-10,saveat=$(imt/1000))

If one prints out p

BenchmarkTools.Trial:
  memory estimate:  13.41 MiB
  allocs estimate:  5861
  --------------
  minimum time:     1.096 s (0.00% GC)
  median time:      1.100 s (0.00% GC)
  mean time:        1.108 s (0.00% GC)
  maximum time:     1.140 s (0.00% GC)
  --------------
  samples:          5
  evals/sample:     1

and one prints out q

BenchmarkTools.Trial:
  memory estimate:  903.73 MiB
  allocs estimate:  10243937
  --------------
  minimum time:     7.436 s (2.17% GC)
  median time:      7.436 s (2.17% GC)
  mean time:        7.436 s (2.17% GC)
  maximum time:     7.436 s (2.17% GC)
  --------------
  samples:          1
  evals/sample:     1

Now the only thing that occurs to me is that allocating threads also takes some memory and maybe this is not being cleared after each iteration of the ode solver and this is why the memory blows out of proportion. I am new user of DifferentialEquations.jl and would be grateful for any help.

2

You’re probably hitting the closure inference issue. Put a let block around the threads, or if the arrays are too small you may need to use Polyester.jl to reduce the thread overhead (and be careful to not compose the threads if you do that).