Compiling Universal Probabilistic Programming Languages with Efficient Parallel Sequential Monte Carlo Inference

TL;DR

This paper introduces a novel compilation approach for universal probabilistic programming languages, enabling efficient parallel Sequential Monte Carlo inference on low-level platforms like GPUs, with significant speedups demonstrated.

Contribution

It presents PPL control-flow graphs (PCFGs) for efficient checkpoint handling and a compiler from high-level PPLs to low-level GPU-based implementations, a first in the field.

Findings

Up to 6x speedups over existing PPLs with SMC inference.

First compilation of a universal PPL to GPUs with SMC.

Effective handling of checkpoints enabling high-performance inference.

Abstract

Probabilistic programming languages (PPLs) allow users to encode arbitrary inference problems, and PPL implementations provide general-purpose automatic inference for these problems. However, constructing inference implementations that are efficient enough is challenging for many real-world problems. Often, this is due to PPLs not fully exploiting available parallelization and optimization opportunities. For example, handling probabilistic checkpoints in PPLs through continuation-passing style transformations or non-preemptive multitasking -- as is done in many popular PPLs -- often disallows compilation to low-level languages required for high-performance platforms such as GPUs. To solve the checkpoint problem, we introduce the concept of PPL control-flow graphs (PCFGs) -- a simple and efficient approach to checkpoints in low-level languages. We use this approach to implement RootPPL: a low-level PPL built on CUDA and C++ with OpenMP, providing highly efficient and massively parallel SMC inference. We also introduce a general method of compiling universal high-level PPLs to PCFGs and illustrate its application when compiling Miking CorePPL -- a high-level universal PPL -- to RootPPL. The approach is the first to compile a universal PPL to GPUs with SMC inference. We evaluate RootPPL and the CorePPL compiler through a set of real-world experiments in the domains of phylogenetics and epidemiology, demonstrating up to 6x speedups over state-of-the-art PPLs implementing SMC inference.