Automatically Batching Control-Intensive Programs for Modern Accelerators
This enables efficient execution of control-intensive programs on modern hardware, benefiting fields like Bayesian statistics, but it is incremental as it builds on existing batching concepts.
The paper tackles the challenge of batching computations with data-dependent control flow and recursion for accelerators like GPUs, achieving orders-of-magnitude speedups on the No U-Turn Sampler (NUTS) algorithm in Bayesian statistics.
We present a general approach to batching arbitrary computations for accelerators such as GPUs. We show orders-of-magnitude speedups using our method on the No U-Turn Sampler (NUTS), a workhorse algorithm in Bayesian statistics. The central challenge of batching NUTS and other Markov chain Monte Carlo algorithms is data-dependent control flow and recursion. We overcome this by mechanically transforming a single-example implementation into a form that explicitly tracks the current program point for each batch member, and only steps forward those in the same place. We present two different batching algorithms: a simpler, previously published one that inherits recursion from the host Python, and a more complex, novel one that implemenents recursion directly and can batch across it. We implement these batching methods as a general program transformation on Python source. Both the batching system and the NUTS implementation presented here are available as part of the popular TensorFlow Probability software package.