I am working on optimizing a 3x3 solver in Rheology. Here is my function. My question relates to the amount of memory allocation necessary to store a function in a dictionary. Here is a MWE:
using DifferentialEquations
using StaticArrays
#using InteractiveUtils
function dudt_giesekus(u, p, t, gradv)
η0, τ, α = @view p[:]
# Governing equations are for components of the 3x3 stress tensor
σ = SA_F32[u[1] u[4] 0.; u[4] 0. 0.; 0. 0. u[3]]
∇v = SA_F32[0. 0. 0. ; gradv(t) 0. 0. ; 0. 0. 0.]
D = 0.5 .* (∇v .+ transpose(∇v))
T1 = (η0/τ) .* D
T2 = (transpose(∇v) * σ) + (σ * ∇v)
coef = α / (τ * η0)
F = coef * (σ * σ)
du = σ / τ .+ T1 .+ T2 . F .+ T2 # 9 equations (static matrix)
end
function run()
du = [0. 0. 0.; 0. 0. 0.; 0. 0. 0.]
u = [0. 0. 0.; 0. 0. 0.; 0. 0. 0.]
u0 = SA[0. 0. 0.; 0. 0. 0.; 0. 0. 0.]
t = 0.
p = SA[1.,1.,1.]
dct = Dict()
dct[:γ_protoc] = convert(Vector{Float32}, [1, 2, 1, 2, 1, 2, 1, 2])
dct[:ω_protoc] = convert(Vector{Float32}, [1, 1, 0.5, 0.5, 2., 2., 1/3., 1/3.])
tspan = (0., 5.)
σ0 = SA[0. 0. 0.; 0. 0. 0.; 0. 0. 0.]
#v21_protoc = [ (t) > 2. .* cos(3*t) for i in 1:8] # Requires no or almost no memory allocation
# Memory allocation is 12k per call to solve(). WHY? <<<<<<
v21_protoc = (t) > dct[:γ_protoc][1]*cos(dct[:ω_protoc][1]*t)
#fct = (t) > 2f0*cos(1.5*t) # Very efficient. No memory allocation
dudt(u,p,t) = dudt_giesekus(u, p, t, v21_protoc)
prob_giesekus = ODEProblem(dudt, σ0, tspan, p)
sol_giesekus = solve(prob_giesekus, Tsit5())
sol_giesekus = solve(prob_giesekus, Tsit5())
end
run()
When I run the above as is, the line
v21_protoc = (t) > dct[:γ_protoc][1]*cos(dct[:ω_protoc][1]*t)
allocates about 12k of memory every time solve
is called.
First question: why doesn’t the allocation only occur once? After all, the function is not changing. It is as if the line executes inside dudt_giesekus
.
Second question: why 12k for such a simple statement?
Third question: memory allocation goes to zero when I define fct
according to:
fct = (t) > 2f0*cos(1.5*t) # Very efficient. No memory allocation
which is what I would expect.
I use a dictionary because I want the flexibility to update the parameters of the function from a driver script.
Here are some additional interesting tidbits that I’d like to understand better. The memory allocation map when defining fct
via a dictionary and calling solve
twice is:
20640 ∇v = SA_F32[0. 0. 0. ; gradv(t) 0. 0. ; 0. 0. 0.]
33024 D = 0.5 .* (∇v .+ transpose(∇v))

24768 T1 = (η0/τ) .* D
74160 T2 = (transpose(∇v) * σ) + (σ * ∇v)

0 coef = α / (τ * η0)
0 F = coef * (σ * σ)
61920 du = σ / τ .+ T1 .+ T2 . F .+ T2 # 9 equations (static matrix)
However, the memory allocation goes to zero when defining the function without dictionaries. In both cases, the function is anonymous.
Thanks.