Design patterns for generic code with reusable caches for intermediate values

General Usage
1

We have an expensive algorithm with is run many, many times for different parameter values and want to design the interface for caching the memory preallocation. A few principles:

  • We should have enough information to create the cache prior to running it through the first time, it would like to reuse allocations made for the model solutions as well.
  • We need to support 2 variations of the algorithms for dense vs. sparse values. The cache would need to support this
  • Performance really matters, and the sizes of these might be very large
  • It would be nice to have a convenience version of the function when you don’t want to worry about any caches.

In order to implement this, I came up with the following MWE to get across the features we need for generic types, caches, etc. I would love feedback on the design conventions, hidden gotchas for performance hits, etc. Obvoiusly N=2 is just for the example and I have been too lazy to constrain the types of the generic types

using SparseArrays, Parameters, LinearAlgebra, Statistics

struct MyModel
    is_sparse::Bool
    N::Int64
end

struct MySolverCache{T}
    temp::T  # e.g. Array{Float64,2}
end

MySolverCache(m::MyModel) = MySolverCache(m, Val(m.is_sparse))
MySolverCache(m::MyModel, is_sparse::Val{false}) = MySolverCache(Array{Float64,2}(undef, m.N, m.N))  # cache for dense
MySolverCache(m::MyModel, is_sparse::Val{true}) = MySolverCache(spzeros(m.N, m.N))  # cache for sparse

mutable struct MySolution{T, T2}
    vec_val::T
    array_val::T2   # e.g. Array{Float64,2}
    val::Float64
end

# Constructors to allocate model solution caches
MySolution(m::MyModel) = MySolution(m, Val(m.is_sparse))
MySolution(m::MyModel, is_sparse::Val{false}) = MySolution(Array{Float64,1}(undef, m.N), Array{Float64,2}(undef, m.N, m.N), 0.0)
MySolution(m::MyModel, is_sparse::Val{true}) = MySolution(Array{Float64,1}(undef, m.N), spzeros(m.N, m.N), 0.0)

function make_model(N; is_sparse = false, allocate_cache = false)
    m = MyModel(is_sparse, N)
    return allocate_cache ? (m, MySolverCache(m), MySolution(m)) : m 
end

function solve_model!(x, m::MyModel, solver_cache::MySolverCache = MySolverCache(m), sol::MySolution = MySolution(m))
    @unpack N = m
    @unpack temp = solver_cache
    @unpack vec_val, array_val = sol

    temp .= 0.0 # clears cache if we needed to... works on sparse and dense.
    
    # some calculations...
    temp[1,1] = x
    temp[1,2] = 3.0
    temp[2,1] = 1.0

    vec_val .= Array(diag(temp))   # return a vector
    array_val .= temp  # return a matrix
    array_val[1,1] += 2.0
    sol.val =  mean(vec_val)

    return sol
end

# usage
m, solver_cache, sol = make_model(2, is_sparse = true, allocate_cache = true)  # allocates and sparse
solve_model!(0.5, m, solver_cache, sol)
solve_model!(0.2, m, solver_cache, sol)  # would reuse the cache

m, solver_cache, sol = make_model(2, allocate_cache = true)  # allocates and dense
solve_model!(0.5, m, solver_cache, sol)
solve_model!(0.2, m, solver_cache, sol)

m = make_model(2)  # no allocations, dense
solution = solve_model!(0.5, m)  # no pre-allocation
1 Like
2

  1. I would suggest putting sparsity (is_sparse) in the type somehow so that the single-argument methods for MySolverCache and MySolution are type-stable. A type parameter can do this, but I would use different types (eg SparseModel, DenseModel).

  2. I prefer explicit functions like allocate_cache and empty_solution to doing everything with constructors, but that’s just a style thing.

  3. I would not preallocate anything and would go for a functional style whenever possible instead of filling existing containers. For anything nontrivial, the calculations should dominate anyway, and the code becomes much more transparent.

2 Likes
3

Thanks. Is type stability really for those two? It seems like this would always be put through the function barrier when solving the model? But maybe I am missing something?

Normally, I am 100% with you and the question needs to always be asked. Here it will be important

4

Type stability may not matter for your application if the bulk of the calculation happens inside the function barrier; but it is not possible to determine this from this MWE.

1 Like