Hello,
I’m trying to further reduce memory allocations when solving an ODE using DifferentialEquations.jl. I have already implemented most of the tricks I’m aware of, but I still get a ton of memory allocation when calling min..
Here’s the du function I’m trying to optimize:
function SIA2D!(dH::Matrix{F}, H::Matrix{F}, simulation::SIM, t::F) where {F <: AbstractFloat, SIM <: Simulation}
# Retrieve parameters
@timeit get_timer("ODINN") "Variable assignment" begin
SIA2D_model::SIA2Dmodel = simulation.model.iceflow
glacier::Glacier = simulation.glaciers[simulation.model.iceflow.glacier_idx[]]
params::Parameters = simulation.parameters
int_type = simulation.parameters.simulation.int_type
A::Ref{F} = SIA2D_model.A
B::Matrix{F} = glacier.B
S::Matrix{F} = SIA2D_model.S
dSdx::Matrix{F} = SIA2D_model.dSdx
dSdy::Matrix{F} = SIA2D_model.dSdy
D::Matrix{F} = SIA2D_model.D
dSdx_edges::Matrix{F} = SIA2D_model.dSdx_edges
dSdy_edges::Matrix{F} = SIA2D_model.dSdy_edges
∇S::Matrix{F} = SIA2D_model.∇S
Fx::Matrix{F} = SIA2D_model.Fx
Fy::Matrix{F} = SIA2D_model.Fy
Δx::F = glacier.Δx
Δy::F = glacier.Δy
Γ::Ref{F} = SIA2D_model.Γ
n::int_type = simulation.parameters.physical.n
ρ::F = simulation.parameters.physical.ρ
g::F = simulation.parameters.physical.g
end
# First, enforce values to be positive
@timeit get_timer("ODINN") "H to zero" begin
H[H.<0.0] .= 0.0
# Update glacier surface altimetry
S .= B .+ H
end
# All grid variables computed in a staggered grid
# Compute surface gradients on edges
@timeit get_timer("ODINN") "Surface gradients" begin
diff_x!(dSdx, S, Δx)
diff_y!(dSdy, S, Δy)
∇S .= (avg_y(dSdx).^2 .+ avg_x(dSdy).^2).^((n - 1)/2)
end
@timeit get_timer("ODINN") "Diffusivity" begin
Γ[] = 2.0 * A[] * (ρ * g)^n / (n+2) # 1 / m^3 s
D .= Γ[] .* avg(H).^(n + 2) .* ∇S
end
# Compute flux components
@timeit get_timer("ODINN") "Gradients edges" begin
@views diff_x!(dSdx_edges, S[:,2:end - 1], Δx)
@views diff_y!(dSdy_edges, S[2:end - 1,:], Δy)
end
# Cap surface elevaton differences with the upstream ice thickness to
# imporse boundary condition of the SIA equation
@timeit get_timer("ODINN") "Capping flux" begin
η₀ = params.physical.η₀
dSdx_edges .= min.(dSdx_edges, η₀ * H[1:end-1, 2:end-1]./Δx, η₀ * H[2:end, 2:end-1]./Δx)
dSdy_edges .= min.(dSdy_edges, η₀ * H[2:end-1, 1:end-1]./Δy, η₀ * H[2:end-1, 2:end]./Δy)
dSdx_edges .= max.(dSdx_edges, -η₀ * H[1:end-1, 2:end-1]./Δx, -η₀ * H[2:end, 2:end-1]./Δx)
dSdy_edges .= max.(dSdy_edges, -η₀ * H[2:end-1, 1:end-1]./Δy, -η₀ * H[2:end-1, 2:end]./Δy)
end
@timeit get_timer("ODINN") "Flux" begin
Fx .= .-avg_y(D) .* dSdx_edges
Fy .= .-avg_x(D) .* dSdy_edges
end
# Flux divergence
@timeit get_timer("ODINN") "dH" begin
inn(dH) .= .-(diff_x(Fx) ./ Δx .+ diff_y(Fy) ./ Δy)
end
end
And these are the main results in terms of memory allocation:
to = ────────────────────────────────────────────────────────────────────────────────────
Time Allocations
─────────────────────── ────────────────────────
Tot / % measured: 22.1s / 93.6% 37.1GiB / 99.9%
Section ncalls time %tot avg alloc %tot avg
────────────────────────────────────────────────────────────────────────────────────
Run prediction 1 20.7s 100.0% 20.7s 37.1GiB 100.0% 37.1GiB
Simulate iceflow PDE 1 19.8s 95.7% 19.8s 37.1GiB 100.0% 37.1GiB
Capping flux 13.2k 9.65s 46.7% 729μs 25.4GiB 68.6% 1.97MiB
Surface gradients 13.2k 3.22s 15.6% 243μs 3.21GiB 8.6% 254KiB
Diffusivity 13.2k 2.49s 12.0% 188μs 1.61GiB 4.3% 127KiB
Mass balance 120 1.17s 5.7% 9.75ms 390MiB 1.0% 3.25MiB
Flux 13.2k 969ms 4.7% 73.2μs 3.18GiB 8.6% 252KiB
dH 13.2k 965ms 4.7% 72.9μs 3.15GiB 8.5% 249KiB
Gradients edges 13.2k 377ms 1.8% 28.5μs 0.00B 0.0% 0.00B
H to zero 13.2k 292ms 1.4% 22.1μs 84.9MiB 0.2% 6.56KiB
Variable assignment 13.2k 122ms 0.6% 9.18μs 828KiB 0.0% 64.0B
────────────────────────────────────────────────────────────────────────────────────
What other tricks could I use the further reduce memory allocation? Particularly for the min. part. Thanks a lot in advance!
