# Typestable caching of multiple arrays with different eltype

**URL:** <https://discourse.julialang.org/t/typestable-caching-of-multiple-arrays-with-different-eltype/64197>\
**Category:** Performance\
**Tags:** diffeq\
**Created:** [July 7, 2021, 7:13am UTC](https://discourse.julialang.org/t/typestable-caching-of-multiple-arrays-with-different-eltype/64197 "2021-07-07T07:13:42Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![hexaeder](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/hexaeder/32/24403_2.png) [@hexaeder](https://discourse.julialang.org/u/hexaeder)\
**Post date:** [July 7, 2021, 7:13am UTC](https://discourse.julialang.org/t/typestable-caching-of-multiple-arrays-with-different-eltype/64197/1 "2021-07-07T07:13:42Z")

</div>

Hello,  
i was wondering, whether it is possible to code some kind of cache pool, where i can request cache arrays of various types. Each type of cache array should be only created once. I’ve tried

```julia
# cache arbitrary arrays of dim N
struct CachePool{N}
    caches::Dict{Tuple{DataType, NTuple{N, Int}}, AbstractArray}
    CachePool(N) = new{N}(Dict{Tuple{DataType, NTuple{N,Int}}, AbstractArray}())
end

# request a cache of specific type and size
# allocate new cache if there is none yet
function getcache(c::CachePool, T, size...)::T
    key = (T, size)
    if haskey(c.caches, key)
        cache = c.caches[key]::T
    else
        cache = T(undef, size...)
        c.caches[key] = cache
    end
    return cache
end

pool = CachePool(2)
@code_warntype getcache(pool, Matrix{Int}, 3, 1_000)

```

but (kinda as expected) the `c.caches[key]` is not typestable.

For some context here is the problem i’m trying to solve (maybe there is a much better solution?).  
I have a function for differential equations which needs a cache.

```julia
struct ODEFun{T}
    _cache::T
end
function (ode::ODEFun)(du, u, p, t)
    f!(ode._cache, u, ...) # update cache with expensive calculations
    g!(du, ode._cache, ...) # update du with cache
end

```

This works fine for explicit solvers. However for implicit solvers the types of `du` and `u` change to some AD types. I’d love to have a generic version

```julia
struct ODEFun{T}
    _cachepool::T
end
function (ode::ODEFun)(du::T, u::T, p, t) where {T}
    cache = getcache(ode._cachepool, T)
    f!(cache, u, ...) # update cache with expensive calculations
    g!(du, cache, ...) # update du with cache
end

```

which would work for different AD types as well as CuArrays for example.

Any ideas? 🙂

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [July 7, 2021, 9:50am UTC](https://discourse.julialang.org/t/typestable-caching-of-multiple-arrays-with-different-eltype/64197/2 "2021-07-07T09:50:02Z")

</div>

See @stevengj’s implementation at [FAQ suggestion: lazily allocate buffers for use with automatic differentiation · Issue #769 · SciML/DifferentialEquations.jl · GitHub](https://github.com/SciML/DifferentialEquations.jl/issues/769), and see the diffcache discussion at [ForwardDiff caches usage](https://discourse.julialang.org/t/forwarddiff-caches-usage/64058) . Once JuliaCon videos are done I think it would be good to spawn out a package around this.

---

<div class="post-metadata">

**Author:** ![hexaeder](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/hexaeder/32/24403_2.png) [@hexaeder](https://discourse.julialang.org/u/hexaeder)\
**Post date:** [July 7, 2021, 10:19am UTC](https://discourse.julialang.org/t/typestable-caching-of-multiple-arrays-with-different-eltype/64197/3 "2021-07-07T10:19:31Z")

</div>

Uh thanks, it seems like a simple change to the `get!` function as in the linked issue did the job. It is much nicer anyway 😃

```julia
function getcache(c::CachePool, T, size...)::T
    key = (T, size)
    return get!(c.caches, key) do
	   T(undef, size...)
    end::T
end

```

`@code_typed` still reports the unstability but the allocation is gone and it is like 10 times faster 🤔

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [July 7, 2021, 12:48pm UTC](https://discourse.julialang.org/t/typestable-caching-of-multiple-arrays-with-different-eltype/64197/4 "2021-07-07T12:48:05Z")

</div>

You’ll get a dynamic dispatch, but the type inference issue won’t “leak” out to the rest of your code. This means you pay a price of about 100ns or so extra. For most models that probably doesn’t matter since it’s about as expensive as a few `^` calculations. Because it’s dynamic, it does mean that some static tooling like JET will have a hard time with it as well, but for the most part it should be simpler without a noticeable performance loss on most models.