Suspension Analysis and Selective Continuation-Passing Style for Higher-Order Probabilistic Programming Languages

Probabilistic programming languages (PPLs) make encoding and automatically solving statistical inference problems relatively easy by separating models from the inference algorithm. A popular choice for solving inference problems is to use Monte Carlo inference algorithms. For higher-order functional PPLs, these inference algorithms rely on execution suspension to perform inference, most often enabled through a full continuation-passing style (CPS) transformation. However, standard CPS transformations for PPL compilers introduce significant overhead, a problem the community has generally overlooked. State-of-the-art solutions either perform complete CPS transformations with performance penalties due to unnecessary closure allocations or use efficient, but complex, low-level solutions that are often not available in high-level languages. In contrast to prior work, we develop a new approach that is both efficient and easy to implement using higher-order languages. Specifically, we design a novel static suspension analysis technique that determines the parts of a program that require suspension, given a particular inference algorithm. The analysis result allows selectively CPS transforming the program only where necessary. We formally prove the correctness of the suspension analysis and implement both the suspension analysis and selective CPS transformation in the Miking CorePPL compiler. We evaluate the implementation for a large number of Monte Carlo inference algorithms on real-world models from phylogenetics, epidemiology, and topic modeling. The evaluation results demonstrate significant improvements across all models and inference algorithms.