Inlining Monad Bayes

• Analysing the dumped core of Population.hs and Weighted.hs of monad-bayes:

  

  It turns out that we should expect to see these definitions, however if we use optimization:2, then we should also expect them to not actually be used. This is because GHC generates a lot of garbage that it doesn’t seem to want to remove, but when we go to places where they would have been used, we should be able to find that they aren’t actually there.

• Here is scoreModel from PmmhHmm.hs of probabilistic-programming.

  

  We can see that there aren’t any of the abstractions in there; this core looks pretty good, however logProbability is a numerical computation which could be inlined.

• The hope is that we won’t find any >>=s if the optimiser has successfully removed them from the function. If we look at the pmmh function from PMMH.hs of monad-bayes:

  

  This function itself won’t be able to optimise because it needs a concrete instance of MonadInfer, but this isn’t necessarily a problem. What we need to look at is when we use this function with something concrete. So let’s look at inferModel from PmmhHmm.hs of probabilistic-programming which uses pmmh; the MonadInfer m concretises to IO because we called sampleIO of monad-bayes on it.

  

  The hope is that the call to pmmh (represented by the green-highlighted lvl) would have been inlined - however this is not the case, as it simply references pmmh from the monad-bayes library, which has not been marked as INLINE or INLINABLE. This means GHC isn’t specialising it, so the program still performs dictionary lookups when that code is compiled.

• The monad-bayes library from hackage was removed as a build dependency, and the local clone of monad-bayes is now used as a dependency instead.

  

• The INLINE pragma was added to pmmh of monad-bayes.

  

  Revisiting the inferModel function, we can now see that the call to pmmh (previously represented by lvl) has been inlined with what is now called $w$spmmh.

  • $w is usually prepended to functions that the compiler has unwrapped some of the arguments for (which eliminates a pattern match).
  • $s is usually prepended to functions that the compiler has specialised.

• Let’s now visit the new version of pmmh, called $w$spmmh.

  

  This immediately looks a lot better than before, as it has been specialised and has had its arguments unwrapped. There still exists some possible optimisations. When seeing occurrences of calls such as $s$fMonadStateT_$c>>=, we can see it’s being passed to mhTrans - this is concerning and suggests that mhTrans along with other functions in similar situations such as pushEvidence and hoist can be inlined.

  Placing an INLINE pragma on these functions results in the following core, with the previous calls being inlined:

  

  Additionally, the $w$c>>= (from Bayes.Traced.Static) at the top is suspicious - it looks like one of the arguments is a monadic value, and so it’s forced to use a >>= because it doesn’t know what the monadic value is.

• Let’s then inline the definition of >>= for the Traced monad inside Bayes.Traced.Static.

  

  This successfully inlines the call to >>= inside inferModel.