Design patterns for generic code with reusable caches for intermediate values

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

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

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}
    array_val::T2   # e.g. Array{Float64,2}

# 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 

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

# 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
  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.


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

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