# Memory caching for reducing allocations

**URL:** https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260
**Category:** General Usage
**Created:** [December 29, 2024, 10:19pm UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260 "2024-12-29T22:19:19Z")
**Posts on this page:** 7
**Page:** 1

<div class="post-metadata">

### Author: ![artemsolod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/artemsolod/32/20704_2.png) [@artemsolod](https://discourse.julialang.org/u/artemsolod)
#### Post date: [December 29, 2024, 10:19pm UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260/1 "2024-12-29T22:19:19Z")

</div>

I am experimenting with the idea of automatically reusing memory from previous allocations. For some workloads like parameter search on fixed data I know that allocations will have the same size between iterations. Preallocating often requires modifying functions deep into call stack which is quite inconvenient.

My current attempt is to have a `mem_cache` object with a dictionary of previously `finalized` arrays (or rather their underlying `Memory` objects). However, this requires overriding `Array` (or `Memory`) constructor which I believe leads to invalidations and is considered bad practice.

Is there a better way to achieve this? And is this approach viable overall as a targeted optimization at call site? Thanks!

```julia
using BenchmarkTools
using Random, Statistics

const mem_cache = (; lock=Threads.SpinLock(), dict=Dict{Int, Vector{Memory{Float64}}}())
const cache_on = Base.ScopedValues.ScopedValue(false) # toggle array caching

# reuse cached memoery if available, otherwise use internal `jl_alloc_array_1d` to allocate
# register finalizer that stores memory in cache
function cached_vector_allocator(nels)
    arr = lock(mem_cache.lock) do
        cache_vec = get(mem_cache.dict, nels, nothing)
        if !isnothing(cache_vec) && !isempty(cache_vec)
            Base.wrap(Array, pop!(cache_vec), (nels, ))
        else
            @ccall jl_alloc_array_1d(Vector{Float64}::Any, nels::Csize_t) :: Vector{Float64}
        end
    end
    finalizer(arr) do x
        lock(mem_cache.lock) do
            cache_vec = get!(() -> Vector{Memory{Float64}}(), mem_cache.dict, nels)
            push!(cache_vec, x.ref.mem)
        end
    end
    return arr
end

# use cached vectors for large arrays
@inline function Array{Float64, 1}(::UndefInitializer, m::Int64)
    if cache_on[] && (m >= 100_000)
        cached_vector_allocator(m)
    else
        @ccall jl_alloc_array_1d(Vector{Float64}::Any, m::Csize_t) :: Vector{Float64}
    end
 end

```

and benchmarking

```julia
function test_seq(arr; n_iter=100)
    for _ in 1:n_iter
        median(arr)
    end
end

function test_par(arr; n_iter=100)
    Threads.@threads for _ in 1:n_iter
        median(arr)
    end
end

Random.seed!(42)
arr = randn(5_000_000)

# Sequential benchmark
@btime test_seq($arr) # 5.907 s (900 allocations: 3.79 GiB)

GC.gc()
@btime Base.ScopedValues.with(cache_on => true) do
    empty!(mem_cache.dict)
    test_seq($arr)
end # 4.010 s (940 allocations: 642.37 MiB)

# Parallel benchmark
GC.gc()
@btime test_par(arr) # 2.250 s (922 allocations: 3.79 GiB)

GC.gc()
@btime Base.ScopedValues.with(cache_on => true) do
    empty!(mem_cache.dict)
    test_par($arr)
end # 1.198 s (955 allocations: 527.93 MiB)

```

---

<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: [December 29, 2024, 11:12pm UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260/2 "2024-12-29T23:12:30Z")

</div>

Isn’t [GitHub - MasonProtter/Bumper.jl: Bring Your Own Stack](https://github.com/MasonProtter/Bumper.jl) useful here?

---

<div class="post-metadata">

### Author: ![artemsolod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/artemsolod/32/20704_2.png) [@artemsolod](https://discourse.julialang.org/u/artemsolod)
#### Post date: [December 29, 2024, 11:20pm UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260/3 "2024-12-29T23:20:50Z")

</div>

`Bumper.jl` is useful but requires changes to your code (like placing `@noescape`’s throughout). In my example `median(arr)` makes a copy - and this is hidden in implementation which is a problem for `Bumper.jl`’s approach.

---

<div class="post-metadata">

### Author: ![Benny](https://avatars.discourse-cdn.com/v4/letter/b/49beb7/32.png) [@Benny](https://discourse.julialang.org/u/Benny)
#### Post date: [December 30, 2024, 6:38am UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260/4 "2024-12-30T06:38:05Z")

</div>

So something like CUDA’s memory pool?

[Memory management · CUDA.jl](https://cuda.juliagpu.org/stable/usage/memory/#Memory-pool)

[Using the NVIDIA CUDA Stream-Ordered Memory Allocator, Part 1 | NVIDIA Technical Blog](https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-1/#memory_pools)

> [@artemsolod](#):
>
> a `mem_cache` object with a dictionary of previously `finalized` arrays (or rather their underlying `Memory` objects)

To be accurate, you finalize garbage or actively freed objects; a cache has live references that prevent both. To cache Julia-level garbage, you’d need to change the internal memory management in C and rebuild Julia. However, `Array`s or other containers can be garbage while their associated `Memory` is not, so you could indeed change the constructors for all the core containers to check a `Memory` cache instead of allocating every time. However, you can’t control how third-party developers allocate the public `Memory` in their packages, so getting down to internals has its benefits.

> [@artemsolod](#):
>
> However, this requires overriding `Array` (or `Memory`) constructor which I believe leads to invalidations and is considered bad practice.

You’re modifying Julia’s core implementation, so this would be a fork, not a typical package we just import. I think it’s hypothetically possible for a heavily pirating and invalidating package to interactively change base Julia, but even if that doesn’t break anything, we’d have to recompile a lot. At this point, making in-place functions would be less work, but there are definitely scenarios where preallocation isn’t feasible, and if we can afford to use extra memory for caching then we’ll take the performance boost (why CUDA does it).

---

<div class="post-metadata">

### Author: ![Sukera](https://avatars.discourse-cdn.com/v4/letter/s/ce7236/32.png) [@Sukera](https://discourse.julialang.org/u/Sukera)
#### Post date: [December 30, 2024, 11:26am UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260/5 "2024-12-30T11:26:05Z")

</div>

You could implement a manual memory arena which you request objects from, though this won’t be “hooked into” by your dependencies. There is currently no way to replace the julia GC wholesale for e.g. one function call (and all allocations happening within).

In theory you might be able to cook something up with GPUCompiler.jl (check out AllocCheck.jl for some inspiration) to modify the generated LLVM-IR, which would allow you to replace the internal calls to the GC with calls to your custom arena. Maybe a future capability for Compiler plugins (see [Pull requests · JuliaLang/julia · GitHub](https://github.com/JuliaLang/julia/pulls?q=is%3Apr+is%3Aopen+compiler+plugin) for previous attempts at adding this) would make this easier, but for now, there is only the hard way.

---

<div class="post-metadata">

### Author: ![artemsolod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/artemsolod/32/20704_2.png) [@artemsolod](https://discourse.julialang.org/u/artemsolod)
#### Post date: [December 30, 2024, 3:06pm UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260/6 "2024-12-30T15:06:15Z")

</div>

Thank you, I will have a look in `GPUCompiler.jl` based approaches.

I was wondering about the feasibility of “scoped type piracy” where the methods are redefined within particular scope / context. So in my example I would only override ` Array{Float64, 1}(::UndefInitializer, m::Int64)` **within** my function call. That reminds me of `Cassette.jl` (there was an example changing `sin` to `cos` within a function call) and also about this talk [https://www.youtube.com/watch?v=3fmwk\_Wo788](https://www.youtube.com/watch?v=3fmwk_Wo788) (“methodtable overlays”, starting from ~18:00). But if I am not mistaken `Cassette` is more or less incompatible with recent julia while compiler plugins are not part of julia yet.

Another hack I was thinking of is to remember the method table before doing the piracy and then restore it back (possibly be removing entries with newer world age?). This, however, I do not know if is even possible in principle.

---

<div class="post-metadata">

### Author: ![ericphanson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ericphanson/32/215186_2.png) [@ericphanson](https://discourse.julialang.org/u/ericphanson)
#### Post date: [December 30, 2024, 5:16pm UTC](https://discourse.julialang.org/t/memory-caching-for-reducing-allocations/124260/7 "2024-12-30T17:16:05Z")

</div>

AllocArrays.jl can alleviate this restriction in some cases; `similar` on an AllocArray uses the bump allocator, so you don’t have to modify library code, only to your top-level types/code. It is a much simpler solution in some ways than tools like GPUCompiler (it is just Bumper + dispatch), but more limited.
