# Minimizing allocations from memoized functions

**URL:** <https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221>\
**Category:** Performance\
**Tags:** memory-allocation, memoize\
**Created:** [October 29, 2020, 5:23am UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221 "2020-10-29T05:23:28Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![ash](https://avatars.discourse-cdn.com/v4/letter/a/f19dbf/32.png) [@ash](https://discourse.julialang.org/u/ash)\
**Post date:** [October 29, 2020, 5:23am UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/1 "2020-10-29T05:23:28Z")

</div>

I have the following MWE using a memoized function:

```julia
using Memoization
using FFTW

struct Operators{F}
    K::F
    # More operators
end 

function run(N, α)
    M = 512
    ω = rand(M)

    # Cache the kinetic factor
    function K(α)
        fun = if α == 0 
            @memoize function K_cubic(dx::Real)
                ifftshift(cis.(dx*ω.^2/2))
            end
            K_cubic
        elseif α > 0
            @memoize function K_hirota(dx::Real)
                ifftshift(cis.(dx*(ω.^2/2 - α*ω.^3)))
            end
            K_hirota
        end
        return fun
    end

    ops = Operators(K(α))
    
    ψ = Array{Complex{Float64}}(undef, M, N)
    dx = 1/N
    myloop(ψ, dx, ops, N)

end

function myloop(ψ, dx, ops, N)
    for i = 1:N-1
        @time @views T2A!(ψ[:, i+1], ψ[:, i], dx, ops)
        # or 
       # @time @views T4A!(ψ[:, i+1], ψ[:, i], dx, ops)
    end
end

function T2A!(ψₒ, ψᵢ, dx, ops) # 3 allocs total 
    # Do some non allocating operations using stuff in ops
    ψₒ .= ops.K(dx) .* ψᵢ # 3 allocs from memoized K(dx)
    # Do some more non allocating operations using stuff in ops
end

function T4A!(ψₒ, ψᵢ, dx, ops) # 3 allocs total 
    T2A!(ψₒ, ψᵢ, 0.123*dx, ops) # 3 allocs from memoized K(dx)
    T2A!(ψₒ, ψₒ, -0.456*dx, ops) # 3 allocs from memoized K(dx)
    T2A!(ψₒ, ψₒ, 0.123*dx, ops) # 3 allocs from memoized K(dx)
end

run(1000, 0.1)

```

After the first run of `T2A!`, every iteration of the loop has 3 allocations coming from the line `ψₒ .= ops.K(dx) .* ψᵢ`. Is it possible to get rid of them? The allocations are quite small (a few tens or hundreds of bytes).

Note that I can’t just cache `K` at the given `dx` and pass it as an array to `T2A!` due to how `T4A!` works (my production code has a function `T6A!` that calls `T4A!` with different values of `dx`, and functions that call `T6A!` etc). Without the “higher order” functions that change the `dx`, this approach worked fantastically, but now I’m stuck.

---

<div class="post-metadata">

**Author:** ![lmiq](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lmiq/32/18314_2.png) [@lmiq](https://discourse.julialang.org/u/lmiq)\
**Post date:** [October 29, 2020, 5:56am UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/2 "2020-10-29T05:56:26Z")

</div>

I do not know if that is related, buy why your K function returns a function instead of the value of the corresponding ifftshift? Could you memoise K directly?

---

<div class="post-metadata">

**Author:** ![ash](https://avatars.discourse-cdn.com/v4/letter/a/f19dbf/32.png) [@ash](https://discourse.julialang.org/u/ash)\
**Post date:** [October 29, 2020, 6:25am UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/3 "2020-10-29T06:25:28Z")

</div>

It’s to avoid writing something like this:

```julia
    function K(α, dx)
        if α == 0 
                ifftshift(cis.(dx*ω.^2/2))
        elseif α > 0
                ifftshift(cis.(dx*(ω.^2/2 - α*ω.^3)))
        end
    end

```

which would force me to unnecessarily carry \alpha around when it’s not needed. It’s also 18 allocations per memoized call instead of 3. Not quite sure why, probably has something to do with the branches inside the function.

---

<div class="post-metadata">

**Author:** ![marius311](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marius311/32/3953_2.png) [@marius311](https://discourse.julialang.org/u/marius311)\
**Post date:** [October 29, 2020, 8:23am UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/4 "2020-10-29T08:23:43Z")

</div>

I’m guessing the 3 allocations could be from the memoization machinery when memoizing a closure like is the case in your code, e.g.

```julia
julia> foo = (x -> (@memoize closure(y) = x+y))(1)
(::var"#closure#23"{Int64}) (generic function with 1 method)

julia> @time foo(2) # second time
  0.000013 seconds (3 allocations: 80 bytes)

```

Not sure if those can be gotten rid of but if you wanted to try I’d be happy to take a PR to Memoization.jl (I’m the author of that package). Note that memoizing top-level functions is a bit simpler, so there there’s no allocations, so a workaround could be to not use closures, assuming you really want to get rid of these allocations.

```julia
julia> @memoize bar(x) = x
bar (generic function with 1 method)

julia> @time bar(2) # second time
  0.000002 seconds

```

---

<div class="post-metadata">

**Author:** ![lmiq](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lmiq/32/18314_2.png) [@lmiq](https://discourse.julialang.org/u/lmiq)\
**Post date:** [October 29, 2020, 5:09pm UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/5 "2020-10-29T17:09:26Z")

</div>

I have just learnt about this type of construct, and although it does not solve that specific allocation problem, it does solve one type instability you have on your test. I do not know if this will help you in general, but anyway (this is a function-like object):

```julia
struct Ks
  α :: Float64
  ω :: Vector{Float64}
end

@memoize function (K :: Ks)(dx)
  if K.α == 0
    ifftshift(cis.(dx*K.ω.^2/2))
  else
    ifftshift(cis.(dx*(K.ω.^2/2 - K.α*K.ω.^3)))
  end
end

function run2(N, α)
    M = 512
    ω = rand(512)
    # Cache the kinetic factor
    ops = Operators(Ks(α,ω))
    ψ = Array{Complex{Float64}}(undef, M, N)
    dx = 1/N
    myloop(ψ, dx, ops, N)
end

```

The rest of the code is identical to yours. I am not completely sure that this exact implementation gives you the same result, but it does solve the a type-instability you had:

Your code:

```julia
julia> @code_warntype run(5,0.1)
Variables
  #self#::Core.Compiler.Const(run, false)
  N::Int64
  α::Float64
  M::Int64
  ω::Array{Float64,1}
  K::var"#K#118"{Array{Float64,1}}
  ops::Operators{_A} where _A <<<< RED
  ψ::Array{Complex{Float64},2}
  dx::Float64

Body::Nothing
1 ─ (M = 512)
│ (ω = Main.rand(M::Core.Compiler.Const(512, false)))
│ %3 = Main.:(var"#K#118")::Core.Compiler.Const(var"#K#118", false)
│ %4 = Core.typeof(ω)::Core.Compiler.Const(Array{Float64,1}, false)
│ %5 = Core.apply_type(%3, %4)::Core.Compiler.Const(var"#K#118"{Array{Float64,1}}, false)
│ (K = %new(%5, ω))
│ %7 = (K)(α)::Any <<<<< RED
│ (ops = Main.Operators(%7))
│ %9 = Core.apply_type(Main.Complex, Main.Float64)::Core.Compiler.Const(Complex{Float64}, false)
│ %10 = Core.apply_type(Main.Array, %9)::Core.Compiler.Const(Array{Complex{Float64},N} where N, false)
│ %11 = M::Core.Compiler.Const(512, false)::Core.Compiler.Const(512, false)
│ (ψ = (%10)(Main.undef, %11, N))
│ (dx = 1 / N)
│ %14 = Main.myloop(ψ, dx, ops, N)::Core.Compiler.Const(nothing, false)
└── return %14

```

With the function-like object you get:

```julia
julia> @code_warntype run2(5,0.1)
Variables
  #self#::Core.Compiler.Const(run2, false)
  N::Int64
  α::Float64
  M::Int64
  ω::Array{Float64,1}
  ops::Operators{Ks}
  ψ::Array{Complex{Float64},2}
  dx::Float64

Body::Nothing
1 ─ (M = 512)
│ (ω = Main.rand(512))
│ %3 = Main.Ks(α, ω)::Ks
│ (ops = Main.Operators(%3))
│ %5 = Core.apply_type(Main.Complex, Main.Float64)::Core.Compiler.Const(Complex{Float64}, false)
│ %6 = Core.apply_type(Main.Array, %5)::Core.Compiler.Const(Array{Complex{Float64},N} where N, false)
│ %7 = M::Core.Compiler.Const(512, false)::Core.Compiler.Const(512, false)
│ (ψ = (%6)(Main.undef, %7, N))
│ (dx = 1 / N)
│ %10 = Main.myloop(ψ, dx, ops, N)::Core.Compiler.Const(nothing, false)
└── return %10

```

---

<div class="post-metadata">

**Author:** ![ash](https://avatars.discourse-cdn.com/v4/letter/a/f19dbf/32.png) [@ash](https://discourse.julialang.org/u/ash)\
**Post date:** [October 29, 2020, 7:00pm UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/6 "2020-10-29T19:00:59Z")

</div>

Thank you! I was not aware of this construct and that type instability has been driving me crazy trying to fix it. I’ve tried a similar approach in my actual code and it works quite well. Is there some documentation I can refer to that explains this a little? I’m very new to Julia so I am a little confused by this construct.

---

<div class="post-metadata">

**Author:** ![ash](https://avatars.discourse-cdn.com/v4/letter/a/f19dbf/32.png) [@ash](https://discourse.julialang.org/u/ash)\
**Post date:** [October 29, 2020, 7:11pm UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/7 "2020-10-29T19:11:38Z")

</div>

Thank you for the explanation. The problem is that the loop can have a few billion iterations in long simulations, and when using higher order algorithms similar to `T4A!` I showed above, some of them call  
T2A up to 40 times per loop iteration, so these allocations add up (and they just bug me in general).

I will take a look at the code and see if I can do something, but I’m brand new to Julia so I doubt I can contribute at the moment.

> [@marius311](#):
>
> Note that memoizing top-level functions is a bit simpler, so there there’s no allocations, so a workaround could be to not use closures, assuming you really want to get rid of these allocations.

I have tried a simplified version of Leandro’s suggestion:

```julia
using Memoization
using FFTW

struct Ks
    α :: Float64
    ω :: Vector{Float64}
end
  
@memoize function (K::Ks)(dx)
    if K.α == 0
      ifftshift(cis.(dx*K.ω.^2/2))
    else
      ifftshift(cis.(dx*(K.ω.^2/2 - K.α*K.ω.^3)))
    end
end
  
function run(N, α)
      M = 512
      ω = rand(512)
      ψ = Array{Complex{Float64}}(undef, M, N)
      dx = 1/N
      myloop(ψ, dx, Ks(α,ω), N)
end

function myloop(ψ, dx, K, N)
    for i = 1:N-1
        @views T2A!(ψ[:, i+1], ψ[:, i], dx, K)
    end
end

function T2A!(ψₒ, ψᵢ, dx, K) # 3 allocs total 
    @time ψₒ .= K(dx) .* ψᵢ # 3 allocs from memoized K(dx)
end

run(4, 0.1)  

```

This still has 3 allocs per call of `K(dx)`. Is `K` not considered a top level function in this implementation or is it because it returns a function “handle”? (not sure what they’re called in Julia)

P.S. nice coincidence seeing another Berkeley physicist on here!

---

<div class="post-metadata">

**Author:** ![lmiq](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lmiq/32/18314_2.png) [@lmiq](https://discourse.julialang.org/u/lmiq)\
**Post date:** [October 29, 2020, 7:22pm UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/8 "2020-10-29T19:22:39Z")

</div>

> [@ash](#):
>
> Is there some documentation I can refer to that explains this a little?

The docs are here: [Methods · The Julia Language](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects)

But I have learnt this only a couple of days ago thanks to this post:

> [@Define data-dependent function, various ways](https://discourse.julialang.org/t/define-data-dependent-function-various-ways/48988/4):
>
> I think [function-like objects](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) (sometimes called “functors”) would be yet another option to achieve the same kind of things. (AFAIU, this is what’s internally used to implement closures, and I think it requires a bit less work for the compiler to optimize. Maybe someone more knowledgeable will chime in and confirm or correct this) struct Fun data :: Vector{Int} end function (f::Fun)(x) s = 0. for i in 1:length(x) s += (x[i]-f.data[i])^2 end s end println(" Functor: ") functor …

In this particular example of yours it seems to make the code more modular, and you can feed omega and alpha parameters to the object as well.

---

<div class="post-metadata">

**Author:** ![lmiq](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lmiq/32/18314_2.png) [@lmiq](https://discourse.julialang.org/u/lmiq)\
**Post date:** [October 29, 2020, 7:24pm UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/9 "2020-10-29T19:24:58Z")

</div>

> [@ash](#):
>
> Is `K` not considered a top level function

I think it is a closure anyway (actually the other way around, closures are function-like objects). You wont solve those allocations of memoise using this.

---

<div class="post-metadata">

**Author:** ![marius311](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marius311/32/3953_2.png) [@marius311](https://discourse.julialang.org/u/marius311)\
**Post date:** [October 30, 2020, 1:15am UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/10 "2020-10-30T01:15:54Z")

</div>

> [@ash](#):
>
> Is `K` not considered a top level function

Correct, its not “top-level” because in Memoization.jl each _instance_ of a `Ks` object gets its own independent memoizaiton cache, and its looking up the right cache for a given object that is causing the allocations (same as for a closure).

---

<div class="post-metadata">

**Author:** ![mikkoku](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikkoku/32/16274_2.png) [@mikkoku](https://discourse.julialang.org/u/mikkoku)\
**Post date:** [October 31, 2020, 7:05am UTC](https://discourse.julialang.org/t/minimizing-allocations-from-memoized-functions/49221/11 "2020-10-31T07:05:17Z")

</div>

How about writing `K_cubic` and `K_hirota` as ordinary functions and K/Ks as a wrapper for them? For me this gets allocations from 3 to 1.

```julia
struct Ks
  α :: Float64
  ω :: Vector{Float64}
end

@memoize function K_cubic(dx::Real, ω)
    ifftshift(cis.(dx*ω.^2/2))
end
@memoize function K_hirota(dx::Real, ω, α)
    ifftshift(cis.(dx*(ω.^2/2 - α*ω.^3)))
end
function (K :: Ks)(dx)
  if K.α == 0
    K_cubic(dx, K.ω)
  else
    K_hirota(dx, K.ω, K.α)
  end
end

```