Optimising Density Computations in Probabilistic Programs via Automatic Loop Vectorisation

TL;DR

This paper introduces an automatic loop vectorisation technique for probabilistic programming languages, significantly improving inference performance by enabling efficient parallel execution of loops with complex control flow.

Contribution

It presents a sound, formal method for automatically vectorising loops in probabilistic programs, handling nested and data-dependent loops without user intervention.

Findings

Achieves 1.1-6x speedups over baseline

Reduces GPU memory usage in many models

Handles all tested models effectively

Abstract

Probabilistic programming languages (PPLs) are a popular tool for high-level modelling across many fields. They provide a range of algorithms for probabilistic inference, which analyse models by learning their parameters from a dataset or estimating their posterior distributions. However, probabilistic inference is known to be very costly. One of the bottlenecks of probabilistic inference stems from the iteration over entries of a large dataset or a long series of random samples. Vectorisation can mitigate this cost, but manual vectorisation is error-prone, and existing automatic techniques are often ad-hoc and limited, unable to handle general repetition structures, such as nested loops and loops with data-dependent control flow, without significant user intervention. To address this bottleneck, we propose a sound and effective method for automatically vectorising loops in probabilistic programs. Our method achieves high throughput using speculative parallel execution of loop iterations, while preserving the semantics of the original loop through a fixed-point check. We formalise our method as a translation from an imperative PPL into a lower-level target language with primitives geared towards vectorisation. We implemented our method for the Pyro PPL and evaluated it on a range of probabilistic models. Our experiments show significant performance gains against an existing vectorisation baseline, achieving $1.1$--$6\times$ speedups and reducing GPU memory usage in many cases. Unlike the baseline, which is limited to a subset of models, our method effectively handled all the tested models.