Ugly memoization

memo :: (a->b) -> (a->b)

enumerate

memo f = let table = [ (x, f x) | x <- enumerate ] in \ y -> let Just r = lookup y table in r

memo (Enumerate a, Eq a) => (a->b) -> (a->b)

Enumerate

enumerate

1 2 5 10 none 10.7 10.7 10.7 10.7 comb 62.6 52.2 45.8 43.4 beau 27.6 17.0 10.4 8.1

import Data.IORef import qualified Data.Map as M memoIO :: (Ord a) => (a -> b) -> IO (a -> IO b) memoIO f = do v <- newIORef M.empty let f' x = do m <- readIORef v case M.lookup x m of Nothing -> do let { r = f x }; writeIORef v (M.insert x r m); return r Just r -> return r return f'

memoIO

f'

unsafePerformIO

memo :: (Ord a) => (a -> b) -> (a -> b) memo f = let f' = unsafePerformIO (memoIO f) in \ x -> unsafePerformIO (f' x)

unsafePerformIO

1 2 5 10 none 10.7 10.7 10.7 10.7 comb 62.6 52.2 45.8 43.4 beau 27.6 17.0 10.4 8.1 ugly 13.9 7.7 3.9 2.7

A snag

readIORef

writeIORef

memoIO f = do v <- newMVar M.empty let f' x = do m <- takeMVar v case M.lookup x m of Nothing -> do let { r = f x }; putMVar v (M.insert x r m); return r Just r -> do putMVar v m; return r return f'

MVar

takeMVar

putMVar

(f x)

memoIO

IORef

Another snag

memo

sid :: String ->; String sid = memo id fcn s = sid (sid s)

sid

takeMVar

x

sid

takeMVar

unsafePerformIO

memoIO

unsafePerformIO

x

memoIO f = do v <- newMVar M.empty let f' x = do m <- readMVar v case M.lookup x m of Nothing -> do let { r = f x }; m <- takeMVar v; putMVar v (M.insert x r m); return r Just r -> return r return f'

(f x)

x

Moral

Here's a problem that I recently ran into. I have a function taking a string and computing some value. I call this function a lot, but a lot of the time the argument has occurred before. The function is reasonably expensive, about 10 us. Only about 1/5 of the calls to the function has a new argument.So naturally I want to memoize the function. Luckily Hackage has a couple packages for memoization. I found data-memocombinators and MemoTrie and decided to try them. The basic idea with memoization is that you have a function likeI.e., you give a function to memo and you get a new function of the same type back. This new function behaves like the original one, but it remembers every time it is used and the next time it gets the same argument it will just return the remembered result. This is only safe in a pure language, but luckily Haskell is pure.In an imperative language you can use a mutable memo table that stores all the argument-result pairs and updates the memo table each time the function is used. But how is it even possible to implement that in a pure language? The idea is to lazily construct the whole memo table in the call to memo, and it will then be lazily filled in.Assume that all values of the argument type a can be enumerated by the method, we could then write memo like this:Note how the memo table is constructed given just f, and this memo table is then used in the returned function.The type of this function would be something likeassuming that the classhas the magic methodThis just a very simplified example, if you tried to use this it would be terrible because the returned function does linear lookup in a list. Instead we want some kind of search tree, which is what the two packages I mention implement. The MemoTrie package does this in a really beautiful way, I recommend reading Conal's blog post about it.OK, enough preliminaries. I used criterion to perform the benchmarking, and I tried with no memoization (none), memo-combinators (comb), and MemoTrie (beau). I had a test function taking about 10us, and then i called this functions with different number of repeated arguments: 1, 2, 5, and 10. I.e., 5 means that each argument occurred 5 times as the memoized function was called.So with no memoization the time per call was 10.7 us all the time, no surprise there. With the memo combinators it was much slower than no memoization; the overhead for looking something up is bigger than the cost of computing the result. So that was a failure. The MemoTrie does better, at about an argument repetition of five it starts to break even, and at ten it's a little faster to memoize.Since I estimated my repetition factor in the real code to be about five even the fastest memoization would not be any better then recomputation. So now what? Give up? Of course not! It's time to get dirty.Once you know a function can be implemented in a pure way, there's no harm in implementing the same function in an impure way as long as it presents the pure interface. So lets write the memo function the way it would be done in, e.g., Scheme or ML. We will use a reference to hold a memo table that gets updated on each call. Here's the code, with the type that the function gets.Theallocated a reference with an empty memo table. We then define a new function,, which when it's called with get the memo table and look up the argument. If the argument is in the table then we just return the result, if it's not then we compute the result, store it in the table, and return it. Good old imperative programming (see below why this code is not good imperative code).But, horror, now the type is all wrong, there's IO in two places. The function we want to implement is actually pure. So what to do? Well, if you have a function involving the IO type, but you can prove it is actually pure, then (and only then) you are allowed to useI'll wave my hands instead of a proof (but more later), and here we goWow, twoon the same line. It doesn't get much less safe than that.Let's benchmark again:Not too shabby, using the ugly memoization is actually a win already at two, and just a small overhead if the argument occurs once. We have a winner!No so fast, there'sMy real code can actually be multi-threaded, so the memo function had better work in a multi-threaded setting. Well, it doesn't. There's no guarantee aboutandwhen doing multi-threading.So we have to rewrite it. Actually, the code I first wrote is the one below; I hardly ever use IORef because I want it to work with multi-threading.So now we use aninstead. This makes it thread safe. Only one thread can execute between theand the. This guarantees than only one thread can update the memo table at a time. If two threads try at the same time one has to wait a little. How long? The time it takes for the lookup, plus some small constant. Remember that Haskell is lazy, the theis not actually computed with the lock held, which is good.So I think this is a perfectly reasonable. And we can do the same unsafe trick as before and make it pure. Performance of this version is the same as with theAhhhh, bliss. But wait, there'sThat might look reasonable, but in fact thefunction is broken now. It appears to work, but here's a simple use that failsWhat will happen here? The outer call towill execute theand then do the lookup. Doing the lookup with evaluate the argument,. But this argument is another call to, this will try to execute the. Disaster has struck, deadlock.What happened here? The introduction ofruined the sequencing guaranteed by the IO monad that would have prevented the deadlock if we had used. I got what I deserved for usingCan it be repaired? Well, we could make sureis fully evaluated before grabbing the lock. I settled for a different repair, where the lock is held in a shorter portion of the code.This solution has its own problem. It's now possible for several threads to computefor the sameand the result of all but one of those will be lost by overwriting the table. This is a price I'm willing to pay for this application.Yes, you can use imperative programming to implement pure functions. But the onus is on you to prove that it is safe. This is not as easy as you might think. I believe my final version is correct (with the multiple computation caveat), but I'm not 100% sure.