Unstacking Monads for Performance

While reflecting on how I might be able to improve my SAT solver, I discovered that my inner bottleneck loop contained two monad transformers (ErrorT and StateT) on top of a base monad (ST). The two transformers are provided by the Monad Transformer Library (MTL). According to the Haskell Wiki’s page on improving performance with monads:

MTL is an excellent library for programming with monads. However stacked monad transformers do not inline well and the library is in need of an optimization pass. As a result, it can often impose a performance hit of up to 300% (your code will run up to three times slower).

Err … I guess having stacked transformers is a Bad Idea for my inner loop. So I set out to improve the situation. On the same wiki page, there is a report of excellent speedups for the continuation-passing-style (CPS) approach (section 2 on the wiki). The idea is that you implement a custom monad manually combining the features you want, in CPS.

The State-threading-ST-Error Monad

Accordingly, following the advice of that same wiki page, I implemented a new monad supporting state threading, errors, and ST actions, all in continuation-passing style:

> {-# LANGUAGE PolymorphicComponents
>             ,MultiParamTypeClasses
>             ,FunctionalDependencies
>             ,FlexibleInstances
>  #-}
>
> import Control.Monad.Error hiding ((>=>), forM_)
> import Control.Monad.ST.Strict
> import Control.Monad.State.Lazy hiding ((>=>), forM_)
> import Control.Monad.MonadST

Performing an ST action requires a kind of lifting.

> dpllST :: ST s a -> SSTErrMonad e st s a
> {-# INLINE dpllST #-}
> dpllST st = SSTErrMonad (\k s -> st >>= \x -> k x s)
>

SSTErrMonad e st s a: the error type e, state type st, ST thread
s and result type a.

> newtype SSTErrMonad e st s a =
>     SSTErrMonad { unSSTErrMonad :: forall r. (a -> (st -> ST s (Either e r, st)))
>                               -> (st -> ST s (Either e r, st)) }
>
> instance Monad (SSTErrMonad e st s) where
>     return x = SSTErrMonad ($ x)
>     m >>= f  = bindSSTErrMonad m f
>
> bindSSTErrMonad :: SSTErrMonad e st s a -> (a -> SSTErrMonad e st s b) -> SSTErrMonad e st s b
> {-# INLINE bindSSTErrMonad #-}
> bindSSTErrMonad m f = SSTErrMonad (\k -> unSSTErrMonad m (\a -> unSSTErrMonad (f a) k))
>
> instance MonadState st (SSTErrMonad e st s) where
>     get = SSTErrMonad (\k s -> k s s)
>     put s' = SSTErrMonad (\k _ -> k () s')
>
> instance (Error e) => MonadError e (SSTErrMonad e st s) where
>     throwError err =            -- throw away continuation
>         SSTErrMonad (\_ s -> return (Left err, s))
>     catchError action handler = SSTErrMonad
>         (\k s -> do (x, s')                      case x of
>                       Left error -> unSSTErrMonad (handler error) k s'
>                       Right result -> k result s')

The brilliant thing about this implementation is that the monadic bind operator >>= does no case analysis: only if you explicitly attempt to catch an error do you need to do case analysis. In contrast, the ErrorT implementation does:

instance (Monad m, Error e) => Monad (ErrorT e m) where
    m >>= k  = ErrorT $ do
        a <- runErrorT m
        case a of
            Left  l -> return (Left l)
            Right r -> runErrorT (k r)
    ...

The wiki page argues essentially that function calling (in my monad’s >>=) is less expensive than constant case analysis when errors are uncommon. And indeed, in my solver, they are uncommon (outside the inner loop, they are not possible).

Pretty Result Graphs; and, What did I do wrong?

But the result is a letdown:

(Graph produced using a Haskell Chart library, built on top of gtk2hs.)

The graph compares the runtime (in seconds) of the original code, in blue, with the new code, in red, on 52 benchmarks available from SATLIB.

The new monad only improves solving times slightly across my benchmarks. So, possibilities:

• I’ve implemented the monad incorrectly, which doesn’t seem likely since my many tests pass.
• I’ve implemented the monad correctly but inefficiently.
• The time spent doing case analysis isn’t significant compared with function invocation.

If anyone has a suggestion, please post it in the comments. I’ll definitely try it out.

P.S. Rest of SSTErrMonad, if you’re interested

> -- | @runSSTErrMonad m s@ executes a `SSTErrMonad' action with initial state @s@
> -- until an error occurs or a result is returned.
> runSSTErrMonad :: (Error e) => SSTErrMonad e st s a -> (st -> ST s (Either e a, st))
> runSSTErrMonad m = unSSTErrMonad m (\x s -> return (return x, s))
>
> evalSSTErrMonad :: (Error e) => SSTErrMonad e st s a -> st -> ST s (Either e a)
> evalSSTErrMonad m s = do (result, _)                           return result

Update 17 May 2008, 1549: The graph was on too big a scale because of one point, so I removed the point and now the differences are more manifest.