Impulse: Momentously Fast, General, and Portable Probabilistic Programming via Compiler Augmentation
Probabilistic programming frameworks automate Bayesian inference, but their performance is limited by a structural mismatch: the code a scientist writes to naturally express a model is often not the code an efficient sampler should execute. In practice, inference is often dominated by expensive computations entirely determined by fixed data and model structure, not by the sampled parameters. Because this structure is not explicit to a general-purpose compiler, the inference computation appears parameter-dependent and is therefore redundantly repeated across iterations.
To eliminate this redundancy, we present Impulse, a probabilistic programming system built on MLIR that preserves probabilistic semantics as first-class compiler primitives. We introduce \emph{Sample-Invariant Code Motion} (SICM), a compiler pass that performs dataflow analysis to understand operation dependencies on sampled parameters, applies algebraic rewrites to factor out invariant computation, and hoists the invariant computation out of the inference process. On standard benchmarks, Impulse achieves a geomean speedup of $1.7\times$ over NumPyro, $6.4\times$ over Stan, and $53\times$ over Turing. SICM enables speedups from $8.5\times$ to $190\times$ on models with sample-invariant structure.
| (pldi26src-brantqian.pdf) | 591KiB |
I am a second-year PhD student in the Department of Computer Science at the University of Illinois Urbana-Champaign (UIUC).
