Memoisation

Sometimes, we have pure functions that will return the exact same output for a given input. Instead of having to re-calculate these values again and again, we may want to store these values so that they can be returned quickly, instead of having to recompute the result again. This only makes sense when we are calling the same method many times with the same inputs, and we have the memory resources to keep track of the results.

This technique is known as memoisation, which caches the result of a function call based on the input arguments. Let’s take a very simple example of calculating a number in the Fibonacci sequence:

function fib(n)
    if n <= 2
        return 1
    end
    return fib(n-1) + fib(n-2)
end

We can create a functor to combine the function and the storage of the results in a dictionary:

struct CachedFib
    cache::Dict{Int, Int}
end

# Create an argument-less constructor
CachedFib() = CachedFib(Dict{Int, Int}())

function (f::CachedFib)(n)
    if n <= 2
        return 1
    end
    cache = f.cache
    # Check if the cache contains the result
    if haskey(cache, n)
        return cache[n]
    end
    # If not, actually calculate the result
    result = f(n-1) + f(n-2)
    # Store the result in the cache
    cache[n] = result
    # Return the result
    return result
end

We can now create this functor and use it to calculate our values:

cached_fib_fn = CachedFib();
println(cached_fib_fn(10))

55

Let’s check that the two implementations agree, and then measure what the cache actually buys us. Note that CachedFib mutates its own cache, so benchmarking a functor that has already been called would only measure dictionary lookups. To measure the cost of computing the answer, we build a fresh functor with an empty cache inside the benchmark:

import BenchmarkTools: @btime
@show fib(30) == CachedFib()(30)
print("Naive recursion: "); @btime fib(30);
print("Cold cache:      "); @btime CachedFib()(30);

fib(30) == (CachedFib())(30) = true
Naive recursion:   2.787 ms (0 allocations: 0 bytes)
Cold cache:        701.250 ns (7 allocations: 1.59 KiB)

The naive version recomputes the same subresults over and over, so its cost grows exponentially in n, whereas the cached version computes each fib(k) exactly once and then reads it back from the dictionary. Even though it has to build and populate a dictionary on every one of those samples, the cached version comes out more than a thousand times faster. Once the cache is warm, a repeated call is nothing more than a single dictionary lookup, which takes tens of nanoseconds:

warm_fib = CachedFib();
warm_fib(30);
print("Warm cache:      "); @btime $warm_fib(30);

Warm cache:        6.360 ns (0 allocations: 0 bytes)

This is the trade-off memoisation makes: we spend memory - one dictionary entry per distinct input, which is where the allocations reported for the cold-cache benchmark come from - in order to avoid repeating work.

Note that the backing cache does not have to be a dictionary, or even stored in main memory, one can just as easily use a disk cache as well for results that take a long time to run. This pattern can be very useful for long-running calculations that may be interrupted, such as a HPC job. Using the disk as a cache can easily allow you to recover from stops, avoiding having to recalculate everything from the beginning.