# Optimising Density Computations in Probabilistic Programs via Automatic Loop Vectorisation

Sangho Lim, Hyoungjin Lim, Wonyeol Lee, Xavier Rival, Hongseok Yang

► To cite this version:

Sangho Lim, Hyoungjin Lim, Wonyeol Lee, Xavier Rival, Hongseok Yang. Optimising Density Computations in Probabilistic Programs via Automatic Loop Vectorisation. Proceedings of the ACM on Programming Languages, 2026, 10 (POPL), pp.597 - 627. 10.1145/3776663. hal-05468306

HAL Id: hal-05468306

https://hal.science/hal-05468306v1

Submitted on 20 Jan 2026

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
Optimising Density Computations
in Probabilistic Programs via Automatic Loop Vectorisation

SANGHO LIM*, KAIST, Republic of Korea

HYOUNGJIN LIM*, KAIST, Republic of Korea

WONYEOL LEE, POSTECH, Republic of Korea

XAVIER RIVAL, Inria, France and DIENS - École Normale Supérieure de Paris, CNRS, PSL University, France

HONGSEOK YANG, KAIST, Republic of Korea

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× 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.

CCS Concepts: • Mathematics of computing → Probabilistic inference problems; • Computing methodologies → Parallel algorithms; Vector / streaming algorithms.

Additional Key Words and Phrases: probabilistic programming, loop vectorisation, loop parallelism

**ACM Reference Format:**

Sangho Lim, Hyoungjin Lim, Wonyeol Lee, Xavier Rival, and Hongseok Yang. 2026. Optimising Density Computations in Probabilistic Programs via Automatic Loop Vectorisation. *Proc. ACM Program. Lang.* 10, POPL, Article 21 (January 2026), 31 pages. https://doi.org/10.1145/3776663

# 1 Introduction

Over the last two decades, probabilistic programming languages (PPLs) have emerged as important tools for expressing and reasoning about models in machine learning, statistics, and computational science [1, 9, 17, 25, 52, 63, 64]. They are used to analyse datasets from epidemiology [53, 71] and ecol- ogy [33, 54] to phylogenetics [58] and robotics [16, 32]. PPLs enable users to model probabilistic phe- nomena naturally using primitives for probabilistic sampling and conditioning, as well as features

Authors' Contact Information: Sangho Lim, KAIST, Daejeon, Republic of Korea, lim.sang@kaist.ac.kr; Hyoungjin Lim, KAIST, Daejeon, Republic of Korea, lmkmkrr@kaist.ac.kr; Wonyeol Lee, POSTECH, Pohang, Republic of Korea, wonyeol.lee@postech.ac.kr; Xavier Rival, Inria, Paris, France and DIENS - École Normale Supérieure de Paris, CNRS, PSL University, Paris, France, rival@di.ens.fr; Hongseok Yang, KAIST, Daejeon, Republic of Korea, hongseok.yang@kaist.ac.kr.

<figure><img src="image_2.png" /></figure>

This work is licensed under a Creative Commons Attribution 4.0 International License.

© 2026 Copyright held by the owner/author(s).

ACM 2475-1421/2026/1-ART21

https://doi.org/10.1145/3776663

<footnote>*Both authors contributed equally to this paper.</footnote>
from general-purpose programming languages, such as loops, branches, and more recently, trainable parameters (e.g., neural networks). These models are then analysed using generic inference engines of PPLs, which approximate posterior distributions or learn model parameters by performing algorithms from Bayesian statistics and deep learning, such as Hamiltonian Monte Carlo (HMC) [24, 43], Stochastic Variational Inference (SVI) [29], and Stochastic Gradient Descent (SGD) [34, 56].

This paper aims to reduce the high computational cost of PPL inference engines, a challenge that persists despite advances in statistical algorithms and compilation. Our work focuses on scenarios requiring repeated evaluation of probabilistic programs during parameter learning or posterior inference, particularly those involving intensive computations of model densities and their gradients arising from repetition structures, often implemented as for-loops. We develop a sound automatic method for optimising such loops through tensor-based vectorisation, thereby enabling the underlying inference engines to analyse these models more efficiently.

Models with repetition structures are common in practice. Graphical models from statistics even employs a special notation, called *plate*, to compactly describe such structures. The simplest repetition structure is a loop over a dataset where each iteration is independent of the others in the sense that it does not have data dependency from the other iterations. Our optimisation allows the inference engines to process such loops in parallel using tensors, yielding significant benefits for medium-to-large datasets. More challenging examples come from sequence models and hierarchical models, such as Bayesian hierarchical models [4, 20] and hidden Markov models [73], where repetition structures are implemented as nested loops, or they lack independence and are coded as loops whose iterations are data-dependent on the previous ones. Even in these complex scenarios, our optimisation enables the inference engines to identify the available parallelism and exploit it through tensor-based vectorisation, achieving significant speedups in posterior inference and parameter learning.

Speeding up PPL inference engines using tensor-based vectorisation is a known idea in the PPL community, supported by a range of existing primitives, such as `vmap` [13], `vectorized_map` [2], `vectorized_markov` [47], and the MarkovChain distribution [23]. However, these primitives fall short of fully automatic vectorisation of the general repetition structures that we target. Concretely, they require users to write repetition structures explicitly in vectorised form, or to provide vectorisation-related information, e.g., dependency information in the repetition structures. This task is both cumbersome and error-prone; for example, manually vectorising the loop in Fig. 1 requires care due to inter-iteration data dependencies, as detailed in §2. Moreover, these primitives are limited, and cannot handle nested repetition structures or dynamically varying dependencies arising from internal random choices; they often support only repetition without data dependencies. In contrast, our vectorisation is fully automatic, requires no user input, and correctly handles general repetition structures, including nested ones and those with complex, randomly-varying dependencies.

When viewed as a code transformer, our optimisation can be understood as a *non-standard loop vectoriser*. Given a model with a repetition structure expressed as a for-loop over `{0, 1, ..., n-1}`, our optimisation transforms the loop to parallelism-friendly tensor-based vectorised code. The vectorised code here is subtle, which goes beyond the simple standard idea of executing all the *n* iterations of the original loop in parallel. To handle the case that the iterations of the loop are data-dependent, the code implements two novel ideas: (i) *speculative execution* of the loop body, and (ii) the use of a *fixed-point check* for testing the validity of speculation.

(i) The code runs all the *n* iterations of the original loop in parallel, but some iterations are *speculative* in the sense that their starting states may differ from those of the corresponding iterations in the original loop. This speculation lets the code avoid waiting until the starting states of all the iterations of the original loop are known. But the outcomes of these speculative computations may diverge from those of the corresponding iterations of the original loop.
(ii) To address this divergence issue, the code repeats the parallel (and partly speculative) execution of all the *n* iterations of the loop multiple times, and terminates this repetition when it finds out that all the speculative executions in the current repetition are correct, i.e., they correspond to the *n* iterations of the original loop. The termination mechanism here employs a *fixed-point check* to detect the correctness of speculation, and it is guaranteed to stop the repetition within at most *n* iterations (due to a particular type of speculation we use).

For instance, as we explain in §3, our optimisation transforms the loop of Fig. 1 into vectorised code with a fixed-point check, which ensures that the parallel execution of the original loop is repeated just twice. If the reader is familiar with abstract interpretation, a good intuition is that the vectorised code implements an executable version of the collecting semantics from the abstract interpretation literature; an execution in the collecting semantics corresponds to a set of executions in the standard semantics, and loops in the collecting semantics are interpreted in terms of fixed-points.

When a model with a small degree of data dependency (e.g., a low-order Markov model) is expressed in a PPL that supports a powerful tensor framework, tensor broadcasting, and efficient parallel tensor operations, our optimisation can vectorise the computation of the model's density. The vectorised model runs more efficiently than the original one, thanks to parallelism and early termination enabled by our fixed-point check. Moreover, our vectorisation also accelerates gradient computations over the model's density through the automatic differentiation provided by the PPL's backend. In this case, automatic differentiation operates on the vectorised computation graph of the density, thereby computing gradients in a vectorised manner as well. Such models and PPLs are common, and the resulting speedup directly improves the performance of PPL inference engines that repeatedly calculate or differentiate the model's density, such as Stochastic Variational Inference (SVI) and Maximum a Posteriori (MAP) estimation.

We prove the correctness of our optimisation using the formalism of expressing and analysing tensor-based vectorised computations, which we also develop in the paper. To evaluate the effectiveness of our optimisation, we implemented it for the Pyro PPL, and applied it to posterior inference and parameter learning by SVI and MAP estimation as well as to Markov Chain Monte Carlo (MCMC) sampling.<fnref n="1" /> Our experiments on eight Pyro models show that our optimisation significantly enhances inference performance, achieving 1.1–6× speedups over an existing vectorisation baseline, and reducing GPU memory usage to 14–80% for seven out of eight models, with one outlier exhibiting the increase to 109–112%. The vectorisation baseline fails on some of these eight models, while our optimisation was successfully applied to all the models.

We summarise the contributions of this paper below:

(1) We propose a language with specialised primitives that enables tensor-based vectorised execution of probabilistic programs. It is designed to facilitate vectorisation of loops in PPLs and support a broader class of models with control flow and nested structures (§3 and §4).

(2) We present a transformation that converts standard probabilistic programs into our proposed language for vectorisation. We prove the correctness of this transformation (§3 and §4).

(3) We implemented our transformation on top of the Pyro PPL, and evaluated it on several families of models, demonstrating that the transformation preserves program semantics while significantly improving execution time and in some cases dramatically reducing memory usage especially when the models have short-range data dependencies. (§5 and §6).

We clarify that our work focuses on optimising density and gradient computations in probabilistic programs, but does not address the sample-generation cost in sampling-based inference algorithms. The extended version of the paper is available as a preprint [40].

<footnote n="1">¹Our implementation is available at https://github.com/Lim-Sangho/auto-vectorise-ppl public.</footnote>
$$ HMM = \left[ \begin{array}{l} x := 0.0; y := 0.0; \\ \textbf{for } t \textbf{ in range}(10) \textbf{ do } \{ \\ \quad x := \textbf{fetch}((\text{"z"}, t)); \\ \quad \textbf{score}(\log p_N(x; f(y), 1.0)); \\ \quad \textbf{score}(\log p_N(o(t); g(x), 1.0)); \\ \quad y := x \} \end{array} \right], \quad HMM' = \left[ \begin{array}{l} x := 0.0; y := 0.0; \\ \textbf{for } t \textbf{ in range}(10) \textbf{ do } \{ \\ \quad x := \textbf{sample}((\text{"z"}, t], N(f(y), 1.0)); \\ \quad \textbf{observe}(N(g(x), 1.0), o(t)); \\ \quad y := x \} \end{array} \right] $$

<caption>Fig. 1. Example programs describing a hidden Markov model. Here, *HMM* is written in the language presented in §2, and *HMM'* is written in a typical probabilistic programming language.</caption>

## 2 A Simple Programming Language for Unnormalised Probability Densities

This paper is concerned with a class of PPL inference engines that view probabilistic programs as *evaluators* of unnormalised posterior densities, instead of *generators* of weighted samples, and repeatedly run these evaluators during inference or parameter learning. This class includes Stochastic Variational Inference (SVI), Hamiltonian Monte Carlo (HMC), Maximum Likelihood Estimation (MLE), and Maximum a Posteriori (MAP) estimation, forming the core of popular PPLs such as Pyro and Stan. Our goal is to optimise the density evaluators expressed by probabilistic programs.

In this section, we formalise this density-evaluator view of probabilistic programs. We start with an informal overview (§2.1), and then formally define a simple programming language, where a program assumes a fixed collection of random variables, takes the values of these random variables as inputs, and then returns a real number, called total score, as an output (§2.2 and §2.3).

### 2.1 Overview of a Score-Computing Language

In the imperative language to be introduced in §2.2, programs are intended to compute the unnormalised (posterior) probability densities of some probabilistic models, where being unnormalised means the integrals of the densities at all inputs are not 1 (which comes mostly from conditioning in those models). Their inputs are the values of the random variables in the models, and their outputs are the logarithms of the unnormalised densities at the given input values.

For instance, the program *HMM* in Fig. 1 describes the unnormalised posterior density of a hidden Markov model (HMM), and corresponds to *HMM'* in the same figure written with the more traditional *sample* and *observe* constructs. The program defines a probabilistic model over random variables *z*<sub>*t*</sub> and *o*<sub>*t*</sub> for *t* ∈ {0, ..., 9}, where *z*<sub>*t*</sub> (which is referred to by the notation [("z", *t*)] in *HMM* with the string "z") denotes the hidden state of a stochastic machine at step *t*, and *o*<sub>*t*</sub> (which is used only implicitly by the second *score* command in *HMM*) means the observation on the state *z*<sub>*t*</sub> made at step *t*. The probability density of the model is

$$ p_{HMM}(z_0, \dots, z_9, o_0, \dots, o_9) = \prod_{t=0}^{9} (p_N(z_t; f(z_{t-1}), 1.0) \cdot p_N(o_t; g(z_t), 1.0)) $$

where *f*, *g* are the functions in *HMM*, *p*<sub>*N*</sub> the density of the normal distribution, and *z*<sub>-1</sub> the constant 0.0. The program conditions this density with observations *o*<sub>*t*</sub> = *o*(*t*) for all *t*, where the *o*(*t*) denote fixed observed values. Then, it computes the following quantity using the *score* command:

$$ \log p_{HMM}(z_0, \dots, z_9, o(0), \dots, o(9)) = \sum_{t=0}^{9} (\log p_N(z_t; f(z_{t-1}), 1.0) + \log p_N(o(t); g(z_t), 1.0)), \quad (1) $$

where the values of *z*<sub>*t*</sub> are given to the program as inputs. In each iteration of the for-loop, the program reads the input value of *z*<sub>*t*</sub> using the *fetch* construct and computes the two terms in the summand of Eq. (1) using the *score* construct. The outcomes of *score* get accumulated, and the final sum is returned as the output: the logarithm of the unnormalised density in Eq. (1).

In the rest of this section, we formally describe the syntax and semantics of this language of score-computing programs, which will be used to define our optimisation and prove its correctness.
**Syntax:**

Integer Expressions $Z ::= n \mid x^{\text{int}} \mid \text{op}_i(Z_1, \dots, Z_k, E_1, \dots, E_m)$

Real Expressions $E ::= r \mid x^{\text{real}} \mid \text{op}_r(Z_1, \dots, Z_k, E_1, \dots, E_m)$

Index Expressions $I ::= [(\alpha_1, Z_1); \dots; (\alpha_k, Z_k)]$ (for distinct $\alpha_i$'s)

Commands $C ::= \text{score}(E) \mid x^{\text{real}} ::= \text{fetch}(I) \mid \text{skip} \mid x^{\text{int}} ::= Z \mid x^{\text{real}} ::= E$
$\mid C_1; C_2 \mid \text{ifz} Z C_1 C_2 \mid \text{for } x^{\text{int}} \text{ in range}(n) \text{ do } C$ (for $n \ge 1$)

**Semantics:**

$$ [[n]]\sigma = n \qquad [[x^{\text{int}}]]\sigma = \sigma(x^{\text{int}}) $$

$$ [[\text{op}_i(Z_1, \dots, Z_k, E_1, \dots, E_m)]\sigma = [[\text{op}_i]]_a([Z_1]\sigma, \dots, [Z_k]\sigma, [E_1]\sigma, \dots, [E_m]\sigma) $$

$$ [[r]]\sigma = r \qquad [[x^{\text{real}}]]\sigma = \sigma(x^{\text{real}}) $$

$$ [[\text{op}_r(Z_1, \dots, Z_k, E_1, \dots, E_m)]\sigma = [[\text{op}_r]]_a([Z_1]\sigma, \dots, [Z_k]\sigma, [E_1]\sigma, \dots, [E_m]\sigma) $$

$$ [[[(\alpha_1, Z_1); \dots; (\alpha_k, Z_k)]]\sigma = [(\alpha_1, [Z_1]]\sigma; \dots; (\alpha_k, [Z_k]]\sigma)] $$

$$ (x^{\text{real}} := \text{fetch}(I)), D, \sigma \downarrow \sigma[x^{\text{real}} \mapsto D([I]]\sigma), 0.0 \qquad \text{score}(E), D, \sigma \downarrow \sigma, [[E]]\sigma $$

$$ (x^{\text{int}} := Z), D, \sigma \downarrow \sigma[x^{\text{int}} \mapsto [[Z]]\sigma], 0.0 \qquad (x^{\text{real}} := E), D, \sigma \downarrow \sigma[x^{\text{real}} \mapsto [[E]]\sigma], 0.0 $$

$$ \frac{C_1, D, \sigma \downarrow \sigma', r_1 \quad C_2, D, \sigma' \downarrow \sigma'', r_2}{(C_1; C_2), D, \sigma \downarrow \sigma'', (r_1 + r_2)} \qquad \frac{[[Z]]\sigma = 0 \quad C_1, D, \sigma \downarrow \sigma', r}{(\text{ifz} Z C_1 C_2), D, \sigma \downarrow \sigma', r} \qquad \frac{[[Z]]\sigma \neq 0 \quad C_2, D, \sigma \downarrow \sigma', r}{(\text{ifz} Z C_1 C_2), D, \sigma \downarrow \sigma', r} $$

$$ \frac{\sigma_0 = \sigma \quad C, D, \sigma_k[x \mapsto k] \downarrow \sigma_{k+1}, r_{k+1} \text{ for all } k \in \{0, \dots, n-1\}}{(\text{skip}, D, \sigma \downarrow \sigma, 0.0) \qquad (\text{for } x \text{ in range}(n) \text{ do } C), D, \sigma \downarrow \sigma_n, \sum_{k=1}^{n} r_k} $$

<caption>Fig. 2. Syntax and semantics of the score-computing language.</caption>

## 2.2 Syntax

Let $\mathbb{S}, \mathbb{Z}, \mathbb{R}$, Var be the sets of strings, integers, real numbers, and variables, respectively. We use the following symbols to range over elements in these sets: $\alpha, \beta \in \mathbb{S}$, $n, m \in \mathbb{Z}$, $r \in \mathbb{R}$, and $x, y, z \in \text{Var}$. Each variable $x \in \text{Var}$ has a type $\tau \in \{\text{int}, \text{real}\}$. We often write $x^{\tau}$ to denote that $x$ has type $\tau$.

The syntax of the score-computing language is given in Fig. 2, based on the above notations. The language has three types of expressions, which denote state-dependent integers, reals, and indices. By indices, we mean finite sequences of string-integer pairs that do not use the same string more than once, that is, elements $i$ in the following set:

$$ i \in \text{Index} = \{[(\alpha_1, m_1); \dots; (\alpha_k, m_k)] \in (\mathbb{S} \times \mathbb{Z})^* : \alpha_j \text{’s are distinct for } j \in \{1, \dots, k\}, k \ge 0\}. $$

For instance, when evaluated at a state, an index expression $[(\alpha_1, Z_1); \dots; (\alpha_k, Z_k)]$ returns the index $[(\alpha_1, m_1); \dots; (\alpha_k, m_k)]$, i.e., the length-$k$ sequence of string-integer pairs $(\alpha_j, m_j)$, where $m_j$ is the integer value of $Z_j$ at the state. These indices serve as identifiers of random variables in the language. When a command in the language needs the value of a random variable, it uses the index for the variable as a key and looks up an entry at this key in a table, called random database, which stores the values of all the random variables. The grammars for the expressions are standard, except for the use of $\text{op}_i$ and $\text{op}_r$, the former being an integer-returning operation with appropriate input types and the latter being a real-returning operation.

Commands in the language denote variants of state transformers, which take a random database and a starting state as inputs, and return a final state and a real number called total score as outputs. A random database $D$ here refers to a map from all indices (i.e., finite sequences of string-integer pairs with no repetition of a string) to real numbers, i.e., $D \in \text{RDB} = [\text{Index} \to \mathbb{R}]$. It specifies the
values of random variables (identified by the indices). Thus, besides updating the state, a command *C* computes the total score of the given random database. When *C* represents a probabilistic model, the computed score is the logarithm of an unnormalised density of the model.

The command *score(E)* evaluates *E* in the current state, and returns the result as the total score, while keeping the state unchanged. The next command *x<sup>real</sup>* := *fetch(I)* reads the value at the index *I* in the random database, stores the read value in *x<sup>real</sup>*, and returns the resulting state together with zero score. The other commands are the standard constructs of the imperative languages: *skip*, assignments, sequencing, conditionals, and for-loops. Regarding the state transformation, they have the standard meanings, except that the conditionals use an implicit zero check on their *Z* argument, so that their true branch *C*<sub>1</sub> gets executed if *Z* evaluates to zero, and the false branch *C*<sub>2</sub> is followed otherwise. Regarding the score computation, these commands implement the following two principles: (i) when a command *C* is run, it may execute the *score* command multiple times, and the results of these multiple invocations of *score* get added to become the final total score of *C*; (ii) none of the other atomic commands, such as assignments and *skip*, contributes to the total score. Note that the language is restricted in the sense that it does not have general while-loops. This restriction ensures that all commands in the language always terminate, and this termination property is used when we prove the correctness of our optimisation. We point out that despite this restriction, the language is expressive enough to allow a wide range of typical probabilistic models, as indicated by the eight models used in our experiments. Lifting this restriction is, however, an interesting and important future work. We write *Com* for the set of all commands in the language.

## 2.3 Semantics

Let *State* be the set of states defined by the type-preserving maps from variables to integers or reals:

$$ \mathit{State} = \{ \sigma \in [\mathit{Var} \to \mathbb{Z} \cup \mathbb{R}] : \sigma(x^{\mathit{int}}) \in \mathbb{Z} \text{ and } \sigma(y^{\mathit{real}}) \in \mathbb{R} \text{ for all variables } x^{\mathit{int}} \text{ and } y^{\mathit{real}} \}. $$

We interpret expressions as maps from states to the values of appropriate types: $[[Z]] : \mathit{State} \to \mathbb{Z}$, $[[E]] : \mathit{State} \to \mathbb{R}$, and $[[I]] : \mathit{State} \to \mathit{Index}$. We assume that the semantics of integer-returning or real-returning atomic operations in the language are given, i.e., we have $[[\mathit{op}_i]]_a : \mathbb{Z}^k \times \mathbb{R}^m \to \mathbb{Z}$ and $[[\mathit{op}_r]]_a : \mathbb{Z}^k \times \mathbb{R}^m \to \mathbb{R}$ for all $\mathit{op}_i$ and $\mathit{op}_r$ taking $k$ integer and $m$ real arguments. Using this assumed semantics of atomic operations, we interpret expressions in the standard way (Fig. 2).

We define the meanings of commands using a big-step operational semantics, which specifies the relation $\downarrow \subseteq (\mathit{Com} \times \mathit{RDB} \times \mathit{State}) \times (\mathit{State} \times \mathbb{R})$. Intuitively, $(C, D, \sigma \downarrow \sigma', r)$ means that running *C* on the state *σ* under the random database *D* terminates successfully, producing the state *σ'* and the total score *r*. The score *r* is the logarithm of the unnormalised density at *D* according to a model that *C* implements. Since the language does not have general while-loops, and no commands get stuck or generate an error, every tuple $(C, D, \sigma)$ of command, random database, and input state leads to a unique pair $(\sigma', r)$ made of an output state and a total score. We define $\downarrow$ in the standard manner using the inference-rule notation (Fig. 2).

## 3 Loop Vectorisation

Our optimisation operates on commands in the language of §2 by vectorising for-loops in those commands. In this section, we describe this loop vectorisation as a translation from the language to another language that supports vectorised operations while continuing to have the primitives for describing unnormalised densities. We informally explain and formally define this *target language* of the translation, where all the atomic commands are implicitly vectorised and new language constructs control various aspects of vectorisation (§3.1 and §3.2). The target language is closely related to the popular tensor-based programming languages, such as Python with PyTorch tensors, which our implementation is based on (see §5). After defining the target language, we describe
$$C_0 = \left[ \begin{array}{l} x^{\text{real}\dagger} := \text{fetch}([\text{"z"}, t^{\text{int}\dagger}]); \\ \text{score}(\log p_N(x^{\text{real}\dagger}; f(y^{\text{real}\dagger}), 1.0)); \\ \text{score}(\log p_N(o(t^{\text{int}\dagger}); g(x^{\text{real}\dagger}), 1.0)); \\ y^{\text{real}\dagger} := x^{\text{real}\dagger} \end{array} \right], \quad C_1 = \left[ \begin{array}{l} \text{extend\_index}(\text{"vec"}, 10) \{ \\ \qquad t^{\text{int}\dagger} := \text{lookup\_index}(\text{"vec"}); \\ \qquad C_0 \} \end{array} \right]$$

<caption>Fig. 3. Example commands in the target language.</caption>

the translation from the language of §2 to the target language, and explain the use of speculative computations and fixed-point check in the translated commands, which we mentioned in the introduction (§3.3). In the rest of the paper, we refer to the language of §2 as the *source language*.

## 3.1 Overview of a Target Language for Loop Vectorisation

The target language is an imperative language with a restricted form of recursion, and includes all the constructs from the source language. Its two main features are as follows.

First, variables and expressions in the target language have *lifted types* rather than the original types in the source language, where lifting means changing a type $\tau$ to the type of maps from *Index* to $\tau$. Concretely, variables in the target language store (finite representations of) maps from indices to integers or reals, instead of just integers or reals. Similarly, expressions in the target language evaluate to maps from indices to integers, reals, or indices.

Second, commands in the target language cannot use indices explicitly in order to access entries in these maps. Instead, they maintain a *set of indices* implicitly, and these indices are then used when expressions and variables are accessed during execution. Note that an index set, instead of a single index, is maintained here. This means that a single execution of a command in the target language corresponds to multiple standard executions of the same command, one per each index in the index set, where these executions access different entries of the lifted variables using their indices. This execution model of the target language is similar to the SIMD (single instruction multiple data) model, and can also be understood as an executable special case of the collecting semantics in the abstract interpretation literature.

To help the reader to gain a high-level intuition about how this target language works, we consider two commands in the language: $C_0$ and $C_1$ in Fig. 3.

**Example 1** ($C_0$ in Fig. 3). The command $C_0$ corresponds to a variant of the loop body of *HMM* in Fig. 1, where all occurrences of variables are annotated with their types.

*Lifted types and index sets.* The annotations in $C_0$ use the *lifted types*, real$^\dagger$ and int$^\dagger$, instead of real and int; these *lifted types* indicate that the variables store not real numbers or integers, but maps to them. For instance, the variable $t$ can store the following map from indices to integers:

$$m_t = \lambda i \in \text{Index.} \text{ if } (i = [\text{"vec"}, \ell] \dotplus i' \text{ for some } \ell \in \{0, 1, \dots, 9\} \text{ and } i') \text{ then } \ell \text{ else } 0. \quad (2)$$

Here $j \dotplus j'$ means the concatenation of indices $j$ and $j'$ which are viewed as finite sequences of string-integer pairs. For this $m_t$, we have $m_t([\text{"vec"}, \ell]) = \ell$ for all $\ell \in \{0, \dots, 9\}$.

Some entries of the maps stored in the variables $t, x, y$ are read and updated during the execution of $C_0$. But which entries are read or updated at each line of $C_0$ is not specified explicitly. Instead, it is determined implicitly by an *index set* that is given in the beginning of the execution.

Concretely, assume that $C_0$ is run under the index set $A_{\text{vec}} = \{i_0, i_1, \dots, i_9\}$ for $i_\ell = [\text{"vec"}, \ell]$, as well as a random database $D$ and a starting state $\sigma$ of the target language such that $D([\text{"z"}, \ell]) = z_\ell$, $\sigma(t) = m_t$, $\sigma(x) = m_x$, and $\sigma(y) = m_y$. Here $z_\ell$ is the value of the random variable $z_\ell$ of *HMM* that we used in §2, $m_t$ is the map in Eq. (2), and $m_x, m_y$ are the maps defined by $m_x(i) = m_y(i) = 0.0$ for all $i$. In this setup, $C_0$ is run as follows.
*First command.* The first command $x^{\text{real}\dagger} := \text{fetch}([("z", t^{\text{int}\dagger})])$ reads the ten entries of the map $m_t$ at $i_0, \dots, i_9 \in A_{\text{vec}}$, and forms the ten keys [("z", $m_t(i_0)$)] = [("z", 0)], \dots, [("z", $m_t(i_9)$)] = [("z", 9)] for the random database $D$. Then, it accesses $D$ using these keys, and updates the map $m_x$ with the results of these accesses: for each index $i_t \in A_{\text{vec}}$, the value $D([("z", t)]) = z_t$ is copied to all the entries of $m_x$ whose indices have the form $i_t \leftrightarrow i'$ for some $i'$. Note that the update here is *non-standard* in that the copy of the read value $z_t$ is not restricted to the $i_t$-th entry of $m_x$. More precisely, the resulting state is

$$ \sigma[x \mapsto m_x'] \quad \text{where} \quad m_x'(i) = \text{if } (i = i_t \leftrightarrow i' \text{ for some } l \in \{0, \dots, 9\} \text{ and } i') \text{ then } z_l \text{ else } m_x(i). \quad (3) $$

As we will explain in §4, this non-standard variable update models *tensor broadcasting* in PyTorch and other tensor libraries, and is one of the factors that let us have an effective tensor-based implementation of maps stored in variables. Except for the use of this update, $\sigma[x \mapsto m_x']$ is precisely the state that we would obtain by vectorising the corresponding *fetch* command in the source language over $A_{\text{vec}}$ in the usual sense, i.e., running the command component-wise over a vector of size ten whose components are accessed by indices $i_0, \dots, i_9$.

The first command $x^{\text{real}\dagger} := \text{fetch}([("z", t^{\text{int}\dagger})])$ also returns scores for the ten indices in $A_{\text{vec}}$, which happen to be all zero in this case. The returned scores are packaged into the constant zero map from $A_{\text{vec}}$ to reals, written as $[i_0 \mapsto 0.0, \dots, i_9 \mapsto 0.0]$.

*Remaining commands.* The remaining three commands are executed in a similarly $A_{\text{vec}}$-vectorised way, and give rise to the following final state $\sigma'$ and score map $T' : A_{\text{vec}} \to \mathbb{R}$:

$$ \sigma' = \sigma[x \mapsto m_x', y \mapsto m_x'], \quad T'(i_t) = \log p_N(z_t; f(0.0, 1.0)) + \log p_N(o(l); g(z_t), 1.0), \quad (4) $$

where $m_x'$ is the map in Eq. (3). These outputs correspond to the combined outcomes of all the ten iterations of the loop in *HMM* in the source language, where for each $l \in \{0, \dots, 9\}$, the $l$-th iteration is run after $t$ is set to $l$ and both $x$ and $y$ are set to 0.0.

Note that the sum of the scores in $T'$ differs from $r_{HMM}$, the total score computed by *HMM* in the source language with random variables $z_l$ fixed to the values $z_l$ (see Eq. (1)):

$$ \begin{aligned} \sum_{l=0}^{9} T'(i_l) &= \sum_{l=0}^{9} \left( \log p_N(z_l; f([0.0]), 1.0) + \log p_N(o(l); g(z_l), 1.0) \right) \\ &\neq \sum_{l=0}^{9} \left( \log p_N(z_l; f([z_{l-1}]), 1.0) + \log p_N(o(l); g(z_l), 1.0) \right) = r_{HMM}, \end{aligned} \quad (5) $$

where $z_{-1} = 0.0$. The different parts are highlighted with boxes. They originate from the fact that the execution of $C_0$ under $A_{\text{vec}}$ is *speculative* about the starting value of the variable $y$; at each index $i_t \in A_{\text{vec}}$, the $i_t$ part of the execution uses 0.0 as the starting value of $y$, instead of $z_{l-1}$ used by *HMM*. This type of speculation over starting values of variables at each index is exploited in our loop vectorisation, as we will explain in §3.3. $\square$

**Example 2** ($C_1$ in Fig. 3). The command $C_1$ extends $C_0$ by adding the *lookup_index* command and then wrapping the whole with the *extend_index* construct. The *extend_index* construct *locally changes* the index set given in the beginning of the execution, and in so doing, it alters the level of vectorisation of commands inside its scope temporarily.

Concretely, assume that $C_1$ is executed under the index set $A_{\llbracket} = \{\llbracket\}\llbracket\}$ containing only the empty index (i.e., empty sequence) and the starting state $\sigma_1$ where $\sigma_1(t) = (\lambda i. 0)$ and $\sigma_1(x) = \sigma_1(y) = (\lambda i. 0.0)$ (i.e., $t, x, y$ store constant-zero maps modulo return types). Also, assume that the execution is given the same random database $D$ as in Example 1. In this setup, $C_1$ is run as follows.

*Entry phase and body of extend_index.* The *extend_index* construct in $C_1$ first *extends* the index $\llbracket$ in $A_{\llbracket}$ by adding the string-integer pair ("vec", $l$) at the end for every $l \in \{0, \dots, 9\}$. The updated index set is $A_{\text{vec}} = \{i_0, \dots, i_9\}$ from Example 1.
$$ \textit{Lifted Integer Expr.} Z ::= n | x^{\text{int}\dagger} | \text{op}_i(Z_1, \dots, Z_k, E_1, \dots, E_m) $$

$$ \textit{Lifted Real Expr.} E ::= r | x^{\text{real}\dagger} | \text{op}_r(Z_1, \dots, Z_k, E_1, \dots, E_m) $$

$$ \textit{Lifted Index Expr.} I ::= [(\alpha_1, Z_1); \dots; (\alpha_k, Z_k)] \quad (\text{for distinct } \alpha_i\text{'s}) $$

$$ \begin{array}{l} \textit{Commands} C ::= \textbf{score}(E) \mid x^{\text{real}\dagger} := \textbf{fetch}(I) \mid \textbf{skip} \mid x^{\text{int}\dagger} := Z \mid x^{\text{real}\dagger} := E \mid C_1; C_2 \\ \phantom{\textit{Commands} C} | \textbf{for} \ x^{\text{int}\dagger} \ \textbf{in range}(n_+) \ \textbf{do} \ C \mid \textbf{if} \ Z \ C_1 \ C_2 \mid \textbf{loop\_fixpt\_noacc}(n_+) \ C \\ \phantom{\textit{Commands} C} | \textbf{extend\_index}(\alpha, n_+) \ C \mid x^{\text{int}\dagger} := \textbf{lookup\_index}(\alpha) \mid \textbf{shift}(\alpha) \end{array} $$

<caption>Fig. 4. Syntax of the target language. $n, m$ are non-negative integers, $n_+$ a positive integer, and $\alpha$ a string.</caption>

The **extend_index** construct then executes the commands in its scope under $A_{vec}$, $D$, and $\sigma_1$. The first command in the scope is $t^{\text{int}\dagger} := \text{lookup\_index}(\text{"vec"})$. For every index $i_t \in A_{vec}$, it finds the value paired with "vec" in $i_t$, and copies the value to all the entries of the map $\sigma_1(t)$ whose indices are of the form $i_t \leftrightarrow i'$ for some $i'$. That is, the resulting state is $\sigma_1[t \mapsto m_t]$ where $m_t$ is the map in Eq. (2), so that the state is the same as the starting state $\sigma$ of $C_0$ from Example 1. Regarding the score map, the command does not perform anything special, and simply returns the zero map $[i_0 \mapsto 0.0, \dots, i_9 \mapsto 0.0]$. The remaining command in the scope is $C_0$. It is run under $A_{vec}$, $D$, and $\sigma$, and results in the final state $\sigma'$ and the score map $T'$ from Example 1 (see Eq. (4)).

**Exit phase of extend_index.** The **extend_index** construct finally restores the original index set $A_{\}$ so that the subsequent computation is run under $A_{\}$]. Moreover, the construct post-processes $\sigma'$ and $T'$ to $\sigma''$ and $T''$, which are then returned as the final state and score map. Intuitively, $\sigma''$ is obtained by making the result at the last index $i_9 \in A_{vec}$ the final outcome of the entire computation, and $T''$ is obtained by summing over all the entries in $A_{vec}$. Concretely, for every variable $w$, the postprocessing builds the new map $m_w'' = \sigma''(w)$ from the old map $m_w' = \sigma'(w)$, by reading the $i_9$-th entry of $m_w'$ and storing the read value at every index: $m_w'' = (\lambda i. m_w'(i_9))$. The score map $T''$ is the singleton function mapping the unique index $\emptyset \in A_{\}$] to $\sum_{t=0}^{9} T'(i_t)$. This post-processing comes from our loop vectorisation; when the vectorised parallel execution of all the loop iterations is completed, we need to pick only the final state computed by the last iteration but we have to accumulate the scores computed by all the iterations. The details will appear in §3.3. $\square$

## 3.2 Syntax of the Target Language

Fig. 4 shows the syntax of the target language, which is obtained from that of the source language by annotating variables with lifted types real† and int†, instead of real and int, and including four new constructs, one for a special type of looping and three for controlling vectorisation. Commands in the target language operate on states with the lifted variables, i.e., variables that store maps from indices to reals or integers. As input, they receive a random database $D$, a state $\sigma$ of lifted variables, and a finite set $A$ of indices. Then, they output a final state $\sigma'$ and a score map $T'$ that assigns a score to each index $i \in A$. We write $Com_{tgt}$ for the set of commands in the target language.

Let us go through the cases of a command $C$ in the grammar and explain its behaviour assuming that $C$ is run under $(D, \sigma, A)$.

**Existing commands.** If the command $C$ is one of the first five cases in the grammar, i.e., an atomic command that also appears in the source language (modulo type annotations of variables), then its execution in the target language corresponds to the *vectorised parallel execution* of the same command in the source language over indices in $A$. For each index $i \in A$, $C$ is run according to the semantics of the source language, except that all variable reads are now on the $i$ entries of the maps stored in those variables, and similarly variable updates are on the $i$-related entries; precisely,
all the entries at the indices $i \mapsto i''$ for some $i''$ are updated. The final score map of $C$ is the function $T'$ that maps each $i \in A$ to the total score computed by the execution under $i$.

If the command $C$ is the sequential composition $(C_1; C_2)$, then $C_1$ is run first under given $(D, \sigma, A)$ and $C_2$ is run afterwards under $(D, \sigma'_1, A)$ where $\sigma'_1$ is the final state of $C_1$. The final state $\sigma'$ of the entire $C$ is the one from the execution of $C_2$, and the final score map $T'$ of $C$ is the point-wise sum of the maps $T'_1$ and $T'_2$ resulting from the executions of $C_1$ and $C_2$. If the command $C$ is (for $x^{\text{int}\dagger}$ in range$(n_+)$ do $C$), then $C$ behaves the same as its loop-unrolled version.

If the command $C$ is the conditional (if $Z C_1 C_2$), then its execution starts by partitioning the index set $A$ into $A_1$ and $A_2$ such that $A_1$ contains all the indices $i \in A$ under which $Z$ evaluates to zero, and $A_2$ contains the remaining indices. Then, $C_1$ is run under $(D, \sigma, A_1)$, and $C_2$ is run under $(D, \sigma'_1, A_2)$ afterwards where $\sigma'_1$ is the final state of $C_1$. The final state $\sigma'$ of the conditional command $C$ is the one from the execution of $C_2$, and the final score map $T'$ of $C$ is the following combination of $T'_1$ and $T'_2$ resulting from the runs of $C_1$ and $C_2$: $T'(i) = T'_1(i)$ for all $i \in A_1$ and $T'(i) = T'_2(i)$ for all $i \in A_2$. If the reader is familiar with abstract interpretation, she or he may have noticed that our conditional command behaves according to (a finite executable version of) the standard collecting semantics in abstract interpretation.

*New commands.* The remaining four cases are new. The first case enables a fine-grained control over looping and score management, and the other cases enable that over vectorised parallel execution.

If the command $C$ is (loop\_fixpt\_noacc$(n_+)$ $C_1$), it behaves the same as the sequential composition of $n_+$-many copies of $C_1$'s, except for two points. First, the execution of $C$ skips some copies of $C_1$ and stops early if its fixed-point check gives a positive answer. Concretely, after the run of each copy of $C_1$, the execution checks whether the state has been unchanged during the run of this copy, that is, whether the state is a fixed point of the run of the copy. If the check says yes, the execution stops, skipping the remaining copies of $C_1$. This early stopping based on fixed-point check leads to the performance improvement in our optimisation. Second, the execution of $C$ does not follow the score accumulation scheme used by the sequential composition, but it instead keeps the score of the last executed copy of $C_1$ only, ignoring the scores computed by all the other executed copies of $C_1$. This unusual treatment of scores is closely tied to the use of speculation in our optimisation; ignoring amounts to discarding the results of certain unconfirmed speculative computations.

If the command $C$ is (extend\_index$(\alpha, n_+)$ $C_1$), it extends every index $i$ in $A$ by adding $(\alpha, k)$ at the end for all $k \in \{0, \dots, n_+-1\}$; recall that $C$ is assumed to be run under $(D, \sigma, A)$. This extension may fail if $i \mapsto [(\alpha, k)]$ for some $i \in A$ and $k \in \{0, \dots, n_+-1\}$ violates the definition of Index because the string $\alpha$ is already present in $i$. If this happens, the execution of $C$ gets stuck. Otherwise, we have a new index set $A'$ of size $|A| \times n$ consisting of the extended indices, and $C_1$ is then run under $(D, \sigma, A')$. When the execution of $C_1$ is completed, the original index set $A$ is restored, and the final state $\sigma'_1$ and score map $T'_1$ of $C_1$ are post-processed to yield the final results, $\sigma'$ and $T'$, for $C$. Here $T'$ is obtained from $T'_1$ via summation: $T'(i) = \sum_{k=0}^{n_+-1} T'_1(i \oplus [(\alpha, k)])$ for every $i \in A$. The state $\sigma'$ is obtained by going through every variable $x$ and updating the map $\sigma'_1(x)$ stored in $x$ as follows: $\sigma'(x)(i') = \sigma'_1(x)(i \oplus [(\alpha, n_+-1)])$ if $i' = i \oplus i''$ for some $i \in A$ and $i''$, and $\sigma'(x)(i') = \sigma'_1(x)(i')$ otherwise. Intuitively, this post-processing step ensures that for all $i \in A$, $\sigma'$ keeps only the result of the last ($i \oplus [(\alpha, n_+-1)]$) execution of $C_1$, while $T'$ keeps the scores computed by all the ($i \oplus [(\alpha, \ell)]$) executions of $C_1$ for $\ell \in \{0, \dots, n_+-1\}$.

If the command $C$ is ($x^{\text{int}\dagger} := \text{lookup\_index}(\alpha)$), it first checks whether every index $i$ in the set $A$ includes a pair $(\alpha, \ell_i)$ for some $\ell_i \in \mathbb{Z}$. If the answer is no, the execution of $C$ gets stuck. Otherwise, for every $i \in A$, $C$ stores $\ell_i$ at the ($i \oplus i''$) entry of the map of $x$, for all $i''$. This construct allows the $i$ execution to depend on information specific to $i$, thereby enabling different components of vectorised execution to perform different computations.
<figure><img src="image_3.png" /><figcaption>Fig. 5. Translation of commands in the target language.</figcaption></figure>

If the command C is `shift(α)`, it updates all the maps stored in the variables by *moving* values inside each map in a manner analogous to the bitwise-shift operation. More precisely, assume that C is run under `(D, σ, A)` and `A = {i → [(α, k)] : i ∈ A₀, k ∈ {0, ..., n⁺1}}` for some `n⁺ ≥ 1` and `A₀` such that `α` is not already present in every index `i ∈ A₀`. Then, `shift(α)` results in the state `σ'` by going through every variable `x` and updating the map `σ(x)` stored in `x` as follows:

$$ \sigma'(x)(i') = \begin{cases} \sigma(x)(i \oplus [(\alpha, k-1)]) & \text{if } i' = i \oplus [(\alpha, k)] \oplus i'' \text{ for some } k \in \{1, \dots, n_+ - 1\} \text{ and } i''', \\ \sigma(x)(i) & \text{if } i' = i \oplus [(\alpha, 0)] \oplus i'' \text{ for some } i''', \\ \sigma(x)(i') & \text{otherwise.} \end{cases} $$

A good way to understand this definition is to ignore the `i''` part, and to focus on where the entries of `σ'` come from: the `(i \oplus [(α, k)])` entry is obtained from the `(i \oplus [(α, k-1)])` entry for all `k ∈ {1, ..., n⁺1}`, and the `(i \oplus [(α, 0)])` entry is obtained from the `i` entry. This `shift` command lets different components of a vectorised parallel execution communicate with each other by moving values between them, and enables us to vectorise loops that have dependencies between iterations.

## 3.3 Translation

Our loop vectorisation is implemented as a translation from the source language to the target language. The translation assumes that these languages use the same set of symbols for variables, and for every variable `x` in that set, they give the corresponding types, i.e., `x` has a type `τ` in the source language if and only if it has the type `τ''` in the target language. Using this assumption, the translation converts expressions of all three types in the source language to themselves except that variables `x''` in those expressions are replaced by `x''`. We put a line on top of expressions in the source language to denote their translations. Thus, `E`, `Z`, and `I` mean the translated `E`, `Z`, and `I`.

The translation of commands is defined inductively by the rules in Fig. 5. The most important rule is the one for translating a for-loop and thereby implementing the main part of our loop vectorisation; the other rules simply apply the translation inductively.

Consider the translation of $C \equiv (\text{for } x^{\text{int}} \text{ in range}(n) \text{ do } C_1)$ in the figure. The translated $\bar{C}$ is a loop \_fixpt\_noacc loop run under the scope of the `extend_index(α, n)` construct. A single iteration of this loop corresponds to the *vectorised parallel execution* of the body $C_1$ of the original for-loop over all the indices in `range(n)`. The translation implements this parallel execution using the `extend_index(α, n)` construct, which changes the input index set $A$ to $A' = \{i \oplus [(\alpha, k)] : i \in A, k \in \{0, ..., n-1\}\}$ and increases the level of parallelism by the factor of $n$ since $|A| \times n = |A'|$; the increased parallelism is then used to run all the $n$ iterations of the loop body $C_1$ simultaneously.
$$C_2 = \left[ \begin{array}{l} \text{extend\_index("vec", 10)} \\ \text{loop\_fixpt\_noacc(10)} \\ \text{shift("vec"); } t^{\text{int}\dagger} := \text{lookup\_index("vec"); } C_0 \end{array} \right].$$

<caption>Fig. 6. Example translation of *HMM* in Fig. 1 to the target language. The command $C_0$ was defined in Fig. 3.</caption>

The parallel execution above is *speculative* in that the execution of the translated loop body $\overline{C_1}$ under an index $(i \leftarrow [(\alpha, k+1))] \in A'$ does not wait until the result of the previous iteration under the index $(i \leftarrow [(\alpha, k))]$ is available—it just proceeds with a predicted result of the previous iteration.

To preserve the semantics of the original for-loop despite the speculative execution, the translation of *C* repeats the parallel execution up to *n* times using the *loop>_fixpt_noacc(n)* construct, and forces *communication* between different threads of the parallel execution (which represent computations under different indices in *A'* ) using the *shift*(*α*) command. Each round of the repetition starts by running *shift*(*α*), which moves values between different threads such that the $(i \leftarrow [(\alpha, k+1))]$-indexed thread for the $(k+1)$-th iteration of the loop body $\overline{C_1}$ in this round receives, at the start of execution, the result of the $(i \leftarrow [(\alpha, k))]$-indexed thread for the *k*-th iteration from the previous round. Moving values this way gradually eliminates the speculative nature of the parallel execution: at the end of the *j*-th round, the threads for the 0, ..., *j*-th iterations of $\overline{C_1}$ are guaranteed to start with correct results for the previous iterations so that they are no longer speculative.

Another important aspect of the translation of *C* is that the *loop>_fixpt_noacc* loop is stopped early if the state is not changed by the previous round of the loop, which is detected by the *fixed-point check* of the *loop>_fixpt_noacc* construct. This early stopping occurs often in practice, and it is crucial for the performance improvement of our loop vectorisation.

Finally, we point out that *only the last round* of the *loop>_fixpt_noacc* loop contributes to the final score map *T'* computed by the loop; the score maps computed by all the other rounds of the loop are simply discarded. Using the score map of the last round only comes from the fact that only the parallel execution of the last round is guaranteed to be the non-speculative parallel execution of all the iterations of the original for-loop in the source language.

**Example 3** ($C_2$ in Fig. 6). Recall the *HMM* example from Fig. 1. Translating *HMM* to the target language gives the command $C_2$ in Fig. 6. In the rest of this subsection, we will go through the key steps of the execution of $C_2$, and explain how $C_2$ uses *speculative execution* and *fixed-point check* to vectorise the loop in *HMM* even though the loop iterations have data dependencies between them.

Let *D* be a random database as defined in Example 1. Also, let $A_{[1]} = \{[]\}$ be the singleton set of the empty index, and $\sigma_0$ be the state where every variable stores the constant-zero function. Suppose that $C_2$ is run with $(D, \sigma_0, A_{[1]})$. Then, the *loop>_fixpt_noacc* loop inside $C_2$ performs two iterations, and $C_2$ produces the state $\sigma'$ and the score map $T' : A_{[1]} \to \mathbb{R}$ given by

$$\sigma'(t) = \lambda i. 9, \quad \sigma'(x) = \sigma'(y) = \lambda i. z_9, \quad \sigma'(v) = \sigma_0(v) \text{ for all other variables } v, \\ T'([]) = \sum_{\ell=0}^{9} (\log p_N(z_\ell; f(z_{\ell-1}), 1.0) + \log p_N(o(\ell); g(z_\ell), 1.0)),$$

where $z_{-1} = 0.0$. At the index [], this outcome matches that of *HMM* in the source language. Specifically, $T'([])$ equals to $r_{HMM}$ from Eq. (5), and for every variable *u*, the value $\sigma'(u)([])$ equals to the value of *u* in the final state of *HMM*. We now examine the execution of $C_2$ in detail.

*First round of the loop.* The execution of $C_2$ begins by extending indices in $A_{[1]}$ with ("vec", *l*) for all *l* ∈ {0, ..., 9}, and updating the current index set to that of the extended indices, i.e., $A_{vec} = \{i_0, ..., i_9\}$ for $i_l = [("vec", l)]$. Then, it runs the body of the loop of $C_2$ with $(D, \sigma_0, A_{vec})$. The first command in the loop body is *shift*("vec"), and it acts as no-op (or *skip*) this time, because all the variables in $\sigma_0$ store constant-zero maps and so moving values between entries does not
change those maps. The next command looks up the value bound to “vec” in each index $i$ in $A_{vec}$, stores the looked-up value at the $i$-related entries of the map of $t$ (i.e., entries at indices $i \to i'$ for all $i'$), and returns the following state $\hat{\sigma}_0$:

$\hat{\sigma}_0 = \sigma_0[t \mapsto m_t]$ where $m_t(i) = \text{if } (i = i_l \to i' \text{ for some } l \in \{0, \dots, 9\} \text{ and } i') \text{ then } l \text{ else } \sigma_0(t)(i)$.

The rest of the loop body executes equivalently to $C_0$ in Example 1, producing the same outcome: a state $\sigma'_0$ and a score map $T'_0: A_{vec} \to \mathbb{R}$ given by

$\sigma'_0 = \sigma_0[t \mapsto m_t, x \mapsto m'_0, y \mapsto m'_0], \quad m'_0(i) = \begin{cases} z_l & \text{if } i = i_l \to i' \text{ for some } l \in \{0, \dots, 9\} \text{ and } i', \\ 0.0 & \text{otherwise,} \end{cases}$
$T'_0(i_l) = \log p_N(z_l; f(0.0), 1.0) + \log p_N(o(l); g(z_l), 1.0) \quad \text{for all } l \in \{0, \dots, 9\}$.

Note that for all $i_l \in A_{vec}$, the $i_l$ part of the execution of the loop body corresponds to the $l$-th iteration of the original for-loop of *HMM*, but this correspondence is not perfect due to the *incorrect speculation* used by the former: the $i_l$ part of the execution speculates that in the original for-loop of *HMM*, the variable $y$ stores $\sigma_0(y)(i_l) = 0.0$ after the $(l-1)$-th iteration, but if $l \ge 1$ and $z_{l-1} \ne 0$, the speculation is wrong because the correct value of $y$ is $z_{l-1}$. As a result of incorrect speculation and imperfect correspondence, $T'(i_l)$ is not equal to the score of the $l$-th iteration of *HMM* in general; the former uses $f(0.0)$ instead of $f(z_{l-1})$.

The incorrect speculation is detected by the *fixed-point check* of the *loop_fixpt_noacc* construct of $C_2$, and it is corrected also by the *re-execution* of the body of the *loop_fixpt_noacc* loop with improved speculation. Concretely, the *fixed-point check* of *loop_fixpt_noacc* tests whether the state $\sigma'_0$ is the same as the initial state $\sigma_0$, and gets a negative answer, which indicates that the speculated value of $y$ at some iteration of the for-loop of *HMM* is wrong. Because of the negative answer, the body of *loop_fixpt_noacc* is executed again but this time speculation is made based on the current state $\sigma'_0$, instead of the old state $\sigma_0$. Using the more recent state for speculation always leads to improvement in terms of correctness.

*Subsequent rounds.* In our example, the speculation made at the $j$-th round of *loop_fixpt_noacc* is guaranteed to be correct for the $k$-th iteration of the original for-loop of *HMM* for all $k \le j$ where correctness means that the speculated value of $y$ at the beginning of the $k$-th iteration equals the actual value of $y$ at that iteration of the for-loop of *HMM*. In the example, we have more than this guaranteed improvement, because all the speculated values of $y$ at the second round of the *loop_fixpt_noacc* loop are correct. The use of the *correct speculation* implies that the state $\sigma'_1$ after the second round equals $\sigma'_0$, and so the *fixed-point check* of *loop_fixpt_noacc* succeeds this time and the *loop_fixpt_noacc* loop terminates after this second round.

The final state $\sigma'$ and the score map $T'$ of $C_2$ are then constructed from $\sigma'_1$ and the score map $T'_1$ computed by the second round, according to the semantics of **extend_index** (“vec”, 10):

$$ \sigma' = \sigma'_{0}[t \mapsto m'_{t}, x \mapsto m', y \mapsto m'], \quad m'_{t} = \lambda i. 9, \quad m' = \lambda i. z_{9}, \\ T' = [[] \mapsto \sum_{l=0}^{9} (\log p_{N}(z_{l}; f(z_{l-1}), 1.0) + \log p_{N}(o(l); g(z_{l}), 1.0))] $$

Note that the final $\sigma'$ and $T'$ of $C_2$ at the index $[]$ are the same as the results of *HMM* in Fig. 1. $\square$

# 4 Semantics of the Target Language and Soundness of Translation

In this section, we present an operational semantics of the target language, formalising the intuitive description of the language in §3.1 and §3.2 (§4.1). We then prove that the translation from the source language to the target language, described in §3.3, preserves the semantics (§4.2).
## 4.1 Semantics of the Target Language

*Definition of states.* The semantics relies on the standard prefix order $\sqsubseteq$ on indices: for all $i, j \in \mathit{Index}$, we have $i \sqsubseteq j$ if and only if $|i| \le |j|$ and $i.k = j.k$ for all $k \in \{1, \dots, |i|\}$. Here $|i|$ means the length of $i$, and $i.k$ means the $k$-th element of $i$. We denote the upward and downward closures of an index $i$ and an index set $L$ using the following notations: $i\uparrow = \{j \in \mathit{Index} : i \sqsubseteq j\}$, $i\downarrow = \{j \in \mathit{Index} : j \sqsubseteq i\}$, $L\uparrow = \bigcup_{j \in L} j\uparrow$, and $L\downarrow = \bigcup_{j \in L} j\downarrow$. Also, we let

$$ \max(L) = \text{if } (i \text{ is the } \sqsubseteq\text{-largest element of } L) \text{ then } i \text{ else undefined.} $$

Thus, $i = \max(L)$ if and only if $i \in L$ and for all $i' \in L$, $i' \sqsubseteq i$. Note that if $L \cap j\downarrow \neq \emptyset$ for some index $j$, then $\max(L \cap j\downarrow)$ always exists, because $L \cap j\downarrow$ is totally ordered by $\sqsubseteq$.

For each type $\tau \in \{\text{int}, \text{real}\}$, let $[[\text{int}]] = \mathbb{Z}$ and $[[\text{real}]] = \mathbb{R}$. A state $\sigma$ in the target language is a function sending variables to finite partial maps from indices to integers or reals such that for every variable $x^{\tau^\dagger}$, (i) the range of the partial map $m_x = \sigma(x)$ is $[[\tau]]$, and (ii) the domain of $m_x$ contains the empty index $[]$. A recommended reading is to view the partial map $m_x$ here as a data structure representing the following total function extend($m_x$):

$$ \text{extend}(m_x) : \mathit{Index} \to [[\tau]], \quad \text{extend}(m_x)(i) = m_x(\max(\text{dom}(m_x) \cap i\downarrow)) \quad \text{for all } i \in \mathit{Index}. \quad (6) $$

Note that for every $i \in \mathit{Index}$, the set $(\text{dom}(m_x) \cap i\downarrow)$ is nonempty because it contains the empty index $[]$ at least, and the set is also totally ordered because $i\downarrow$ is totally ordered. As a result, for all $i$, $\max(\text{dom}(m_x) \cap i\downarrow)$ is a well-defined element in $\text{dom}(m_x)$, which in turn implies that $\text{extend}(m_x)$ is a well-defined total map. Although the semantics of the target language to be presented shortly is defined in terms of finite partial maps stored in variables in states, it depends only on the total functions represented by those partial maps: for all states $\sigma, \sigma'$, if $\text{extend}(\sigma(x)) = \text{extend}(\sigma'(x))$ for all variables $x$, the semantics is unable to distinguish between $\sigma$ and $\sigma'$. When we explain the semantics in the rest of the paper, we will often refer to the total function $\text{extend}(\sigma(x))$ as the function stored in a variable $x$, although $\text{extend}(\sigma(x))$ is not stored in $x$. We write $\text{State}_{\text{tgt}}$ for the set of all the states in the target language.

Our *extend* operation models *tensor broadcasting* [69] as implemented in popular tensor libraries such as PyTorch. To see this, first note that the operation can be applied to any finite partial map $m: \mathit{Index} \to V$ for a set $V$, even when the empty index $[]$ is not in $\text{dom}(m)$. The result $\text{extend}(m)$ is still well-defined in this general case, although it might not be a total map. Next, as an example, consider two integer-type tensors $t_1$ and $t_2$ with shapes $(3)$ and $(2, 3)$, respectively. In our setup, they are represented as partial maps $m_1, m_2: \mathit{Index} \to \mathbb{Z}$ with $\text{dom}(m_1) = \{[(\mathbf{snd}, \ell_2)] : \ell_2 \in \{0, 1, 2\}\}$ and $\text{dom}(m_2) = \{[(\mathbf{snd}, \ell_2); (\mathbf{fst}, \ell_1)] : \ell_2 \in \{0, 1, 2\}, \ell_1 \in \{0, 1\}\}$. Although $t_1$ and $t_2$ have different shapes, they can be added using tensor broadcasting: the addition first broadcasts $t_1$ into a tensor of shape $(2, 3)$, and then performs pointwise addition with $t_2$. In our setup, this broadcasting is modelled by applying *extend* to $m_1$, which yields a partial map defined at all indices in $\text{dom}(m_2)$, with each value corresponding to the broadcasted entries of $t_1$. The pointwise addition is then modelled by computing the elementwise sum of $\text{extend}(m_1)$ and $m_2$ at all indices in $\text{dom}(m_2)$. Lastly, note that $\text{extend}(m_1)$ models not only this specific instance of broadcasting $t_1$ with respect to $t_2$, but also all possible instances of broadcasting $t_1$ with respect to any tensor $t'_2$ of a larger “compatible” shape.

*Semantics of target language.* The semantics uses the following five semantic domains and their structures: *Index*, *State*<sub>tgt</sub>, *FAChain*, *Tensor*<sub>Z</sub>, and *Tensor*<sub>R</sub>. We have already defined the first two in this list. The next is *FAChain*, the domain for finite antichains of indices. A set of indices $L$ is an **antichain** if for all $i, j \in L$ with $i \neq j$, we have neither $i \sqsubseteq j$ nor $j \sqsubseteq i$. The domain *FAChain* consists of all the **finite antichains** of indices. When denoting an element of *FAChain*, we often use the symbol $A$, instead of the usual symbol $L$ for an index set, in order to emphasise that it is
<figure><img src="image_4.png" /><figcaption>Fig. 8. Examples for update and copy. Nodes over braces represent the empty indices, and nodes under braces represent indices $i_ℓ = [("rv", ℓ)]$ for $ℓ \in \{0, 1, 2\}$. Values in solid nodes are stored in partial maps, while values in dotted nodes are induced by extend. Thus, $\sigma(x^{\tau^\dagger}) = [[] \mapsto 1, i_1 \mapsto 3, i_2 \mapsto 4]$, $T_1 = [[] \mapsto 6]$, and $T_2 = [i_0 \mapsto 8, i_2 \mapsto 9]$. Also, $\text{extend}(\sigma(x^{\tau^\dagger}))(i_0) = 1$, and $\text{extend}(T_1)(i_ℓ) = 6$ for $ℓ \in \{0, 1, 2\}$. The functions $\rho_1$ and $\rho_2$ denote the partial injections on Index used in the semantics for shift and extend_index (Fig. 9), respectively, and are depicted graphically in the top right; e.g., $\rho_2$ is defined only on $i_2$ and maps it to [].</figcaption></figure>

A good way to understand these slightly unusual update and copy operations on states is to recall that finite partial maps stored in variables represent total functions obtained by the extend operation, and examine what the operations do on these represented functions. For a function $f: Index \to V$ and a partial function $g: Index \to V$, let $f[g]$ be the outcome of updating $f$ with $g$, that is, the function mapping each index $i \in \text{dom}(g)$ to $g(i)$ and every other index $i$ to its original value $f(i)$. The next lemma characterises the effects of our update and copy operations by means of standard update and copy operations on the represented functions. Its proof is in the extended version [40].

**Lemma 1.** Let $\sigma$ be a state, $x^{\tau^\dagger}$, $y^{\tau^\dagger}$ be variables, $T$ be a tensor in $\text{Tensor}_{\llbracket \tau \rrbracket}$, and $\rho$ be a finite partial injection on Index. Then, we have the following equations:

$$ \begin{aligned} \text{extend}(\sigma[x^{\tau^\dagger}: T](y^{\tau^\dagger})) &= \text{if } (y^{\tau^\dagger} \equiv x^{\tau^\dagger}) \text{ then } \text{extend}(\sigma(x^{\tau^\dagger})[\text{extend}(T)]) \text{ else } \text{extend}(\sigma(y^{\tau^\dagger})), \\ \text{extend}(\sigma\langle\rho\rangle(y^{\tau^\dagger})) &= \text{extend}(\sigma(y^{\tau^\dagger})[\text{extend}(T_{y,\rho})]), \end{aligned} $$

$$ \text{where } T_{y,\rho}(i) = \sigma(y^{\tau^\dagger})(\rho^{-1}(i)) \text{ if } i \in \rho(\text{dom}(\sigma(y^{\tau^\dagger}))), \text{ and it is undefined otherwise.} $$

Here we apply the extend operation in Eq. (6) to the tensor $T$, although $T$ may be undefined at the empty index unlike in the case of the partial maps stored in states. As a result, this application may lead to a partial function of type Index $\to \llbracket \tau \rrbracket$, which we recall below:

$$ \text{extend}(T)(i) = \text{if } (i \in \text{dom}(T)\uparrow) \text{ then } T(\max(\text{dom}(T) \cap i\downarrow)) \text{ else undefined.} $$

Note that the equations in the lemma use `(`[extend($T$)] and `)[`extend($T_{y,\rho}$)], instead of `(`[$T$] and `)[$T_{y,\rho}$]. This means that both operations change more entries than one typically expects. For instance, in the first equation, each $(i, T(i))$ entry of the partial map $T$ results in the changes of all the entries in extend($\sigma(x^{\tau^\dagger})$) whose indices are in $i\uparrow$. We use this aggressive version of change in our operations because it enables us to implement the operations and our target language easily and efficiently using the tensor-broadcasting mechanism of PyTorch.

This aggressive change also helps maintain the following important invariant of our target language. For an index set $L$, we say a state $\sigma$ is an *L-state* if $\text{dom}(\sigma(x^{\tau^\dagger})) \subseteq L\downarrow$ for all variables $x^{\tau^\dagger}$. Intuitively, it means that for every variable $x^{\tau^\dagger}$, if the partial map $\sigma(x^{\tau^\dagger})$ stored in $x^{\tau^\dagger}$ represents a function $f = \text{extend}(\sigma(x^{\tau^\dagger}))$, then the function is determined by what it does on $L\downarrow$, i.e., $f = \text{extend}(f|_{L\downarrow})$. The invariant is that if a command $C$ is executed under a finite antichain $A$ and an input state is an $A$-state, then the output state is also an $A$-state. This invariant supports the view that the command $C$ is run in parallel by the threads with their ids in $A$ and the state consists of the disjoint parts for these threads. Formal lemmas and proofs are detailed in the extended version [40].
$$ T = [i \mapsto D(\llbracket I \rrbracket(\sigma, i)) : i \in A] \\ (x^{\text{real}\dagger} := \text{fetch}(I)), D, \sigma, A \llbracket \text{tgt} \sigma[x^{\text{real}\dagger} : T], T_A^z \rightleftharpoons $$

$$ T = [i \mapsto \llbracket E \rrbracket(\sigma, i) : i \in A] \\ \text{score}(E), D, \sigma, A \llbracket \text{tgt} \sigma, T \rightleftharpoons $$

$$ \frac{T = [i \mapsto \llbracket Z \rrbracket(\sigma, i) : i \in A]}{(x^{\text{int}\dagger} := Z), D, \sigma, A \llbracket \text{tgt} \sigma[x^{\text{int}\dagger} : T], T_A^z} $$

$$ T = [i \mapsto \llbracket E \rrbracket(\sigma, i) : i \in A] \\ (x^{\text{real}\dagger} := E), D, \sigma, A \llbracket \text{tgt} \sigma[x^{\text{real}\dagger} : T], T_A^z \rightleftharpoons $$

$$ \frac{\text{skip}, D, \sigma, A \llbracket \text{tgt} \sigma, T_A^z}{C_1, D, \sigma, A \llbracket \text{tgt} \sigma', T' \quad C_2, D, \sigma', A \llbracket \text{tgt} \sigma'', T''} \\ (C_1; C_2), D, \sigma, A \llbracket \text{tgt} \sigma'', (T' \oplus T'')} $$

$$ A_1 = \{i \in A : \llbracket Z \rrbracket(\sigma, i) = 0\} \quad C_1, D, \sigma, A_1 \llbracket \text{tgt} \sigma', T' $$

$$ A_2 = \{i \in A : \llbracket Z \rrbracket(\sigma, i) \neq 0\} \quad C_2, D, \sigma', A_2 \llbracket \text{tgt} \sigma'', T'' $$

$$ \frac{(\text{ifz} Z C_1 C_2), D, \sigma, A \llbracket \text{tgt} \sigma'', (T' \oplus T'')}{\text{for } x^{\text{int}\dagger} \text{ in range}(n_+) \text{ do } C), D, \sigma, A \llbracket \text{tgt} \sigma, T_k} $$

$$ \frac{\sigma_0 = \sigma \quad C, D, \sigma_k[x^{\text{int}\dagger}: [i \mapsto k : i \in A]], A \llbracket \text{tgt} \sigma_{k+1}, T_{k+1} \quad \text{for all } k \in \{0, \dots, n_+ - 1\}}{(\text{for } x^{\text{int}\dagger} \text{ in range}(n_+) \text{ do } C), D, \sigma, A \llbracket \text{tgt} \sigma_{n_+}, \bigoplus_{k=1}^{n_+} T_k} $$

$$ \frac{\alpha \in \text{dom}(i) \text{ for all } i \in A}{T = [i \mapsto i(\alpha) : i \in A]} \\ (x^{\text{int}\dagger} := \text{lookup\_index}(\alpha)), D, \sigma, A \llbracket \text{tgt} \sigma[x^{\text{int}\dagger} : T], T_A^z \rightleftharpoons $$

$$ \alpha \notin \text{dom}(i) \text{ for all } i \in A \quad A' = \{i \mapsto [(\alpha, k)] : i \in A, k \in \{0, \dots, n_+ - 1\}\} \quad C, D, \sigma, A' \llbracket \text{tgt} \sigma', T' $$

$$ \rho = [i \mapsto [(\alpha, n_+ - 1)] \mapsto i : i \in A] \quad T'' = [i \mapsto \sum_{0 \le k < n_+} T'(i \mapsto [(\alpha, k)]) : i \in A] $$

$$ \frac{(\text{extend\_index}(\alpha, n_+), C), D, \sigma, A \llbracket \text{tgt} \sigma' \langle \rho \rangle, T''}{0 \le m < n_+ - 1 \quad \sigma_0 = \sigma} $$

$$ 0 \le m < n_+ - 1 \quad \sigma_0 = \sigma $$

$$ C, D, \sigma_k, A \llbracket \text{tgt} \sigma_{k+1}, T_{k+1} \quad \text{and } \sigma_k \neq \sigma_{k+1} \text{ for all } k \in \{0, \dots, m-1\} \quad C, D, \sigma_m, A \llbracket \text{tgt} \sigma_m, T_{m+1} $$

$$ \frac{(\text{loop\_fixpt\_noacc}(n_+), C), D, \sigma, A \llbracket \text{tgt} \sigma_m, T_{m+1}}{\sigma_0 = \sigma \quad C, D, \sigma_k, A \llbracket \text{tgt} \sigma_{k+1}, T_{k+1} \quad \text{for all } k \in \{0, \dots, n_+ - 1\} \quad \sigma_k \neq \sigma_{k+1} \quad \text{for all } k \in \{0, \dots, n_+ - 2\}} $$

$$ \frac{\sigma_0 = \sigma \quad C, D, \sigma_k, A \llbracket \text{tgt} \sigma_{k+1}, T_{k+1} \quad \text{for all } k \in \{0, \dots, n_+ - 1\} \quad \sigma_k \neq \sigma_{k+1} \quad \text{for all } k \in \{0, \dots, n_+ - 2\}}{(\text{loop\_fixpt\_noacc}(n_+), C), D, \sigma, A \llbracket \text{tgt} \sigma_{n_+}, T_{n_+}} $$

$$ \rho(i) = \begin{cases} i \mapsto [(\alpha, 0)] & \text{if } \alpha \notin \text{dom}(i) \land i \mapsto [(\alpha, 0)] \in A, \\ j \mapsto [(\alpha, m+1)] & \text{if } i = j \mapsto [(\alpha, m)] \in A \downarrow \land j \mapsto [(\alpha, m+1)] \in A \text{ for some } j \text{ and } m \ge 0, \\ \text{undefined} & \text{otherwise} \end{cases} $$

$$ \frac{\text{shift}(\alpha), D, \sigma, A \llbracket \text{tgt} \sigma \langle \rho \rangle, T_A^z}{\rho(i) = \begin{cases} i \mapsto [(\alpha, 0)] & \text{if } \alpha \notin \text{dom}(i) \land i \mapsto [(\alpha, 0)] \in A, \\ j \mapsto [(\alpha, m+1)] & \text{if } i = j \mapsto [(\alpha, m)] \in A \downarrow \land j \mapsto [(\alpha, m+1)] \in A \text{ for some } j \text{ and } m \ge 0, \\ \text{undefined} & \text{otherwise} \end{cases}} $$

<caption>Fig. 9. Semantics of commands in the target language.</caption>

**Semantics of commands.** The rules for the $\llbracket \text{tgt} \sigma \rrbracket$ relation in Fig. 9 mostly follow the general principle of running each command over multiple input data specified by the input antichain $A$, except for two unusual parts. First, updating $x^{\tau\dagger}$ with operation $\sigma[x^{\tau\dagger} : T]$ modifies the partial map $\sigma(x^{\tau\dagger})$ beyond typical SIMD behaviour: if an entry in $\sigma(x^{\tau\dagger})$ has the index $i \in \text{dom}(T)\uparrow$, the entry gets updated (if $i \in \text{dom}(T)$) or removed (if $i \in \text{dom}(T)\uparrow \setminus \text{dom}(T)$). Second, the rule for the conditional command in the semantics is unusual. It splits the input index set $A$ into $A_1$ and $A_2$, where $A_1$ consists of indices leading to the true branch and $A_2$ those leading to the false branch. The command then runs the true branch under $A_1$, and then the false branch under $A_2$. The disjointedness of $A_1$ and $A_2$, and the fact that $A_1 \cup A_2$ is an antichain prevent interference between the two branches: they do not write to the same entry of a variable, and neither of these branches reads from an entry of a variable written by the other branch. This non-interference means the execution order of the true and false branches is interchangeable without changing the semantics of the command.
We now go through the last five rules of the $\llbracket tgt \rrbracket$ relation in Fig. 9, which describe the behaviour of the new constructs in the target language. The rule for the **lookup_index** command says that the command regards each index $i = [(\alpha_1, m_1); \dots; (\alpha_k, m_k)] \in A$ as a partial function mapping string $\alpha_i$ to integer $m_i$, and it runs only when $\alpha$ is in the domain of this partial function (i.e., $\alpha \in \text{dom}(i)$). In that case, the command looks up the $\alpha$ entry of the partial function $i$ for all $i \in A$, and stores the results of these lookups in the variable $x^{\text{int}^\dagger}$. The next rule is concerned with the **extend_index** command, and says that the command works in two steps. First, it runs the command $C$ in its body under the set $A'$ of extended indices, where $A'$ consists of the indices $i \dashv [(\alpha, k)]$ for all $i \in A$ and $k \in \{0, \dots, n-1\}$. Second, it goes through every variable $x^{\tau^\dagger}$, accesses the partial map $m_x$ stored in $x^{\tau^\dagger}$, and copies, for every $i \in A$, the value $m_x(i \dashv [(\alpha, n-1)])$ to the $i$ entry of $m_x$ while removing the entries of $m_x$ whose indices are in $i \upharpoonright \{i\}$. So, the copy here clears the entries at indices $i \dashv [(\alpha, k)]$ for $k \in \{0, \dots, n-1\}$ if such entries exist in $m_x$. Note that the copy is implemented with the help of the $(\_)(\rho)$ operation (Fig. 8). Also, note that the tensor $T''$ in the result of the **extend_index** command is constructed from the values of the $i \dashv [(\alpha, k)]$ entries in $T'$ for all $i \in A$ and $k \in \{0, \dots, n-1\}$. If some of these entries are not defined, the execution does not proceed. But this failure case does not occur, as indicated in the extended version [40], because the domain of the tensor in the result of execution is always the same as the given input antichain for the execution, and so $\text{dom}(T') = A'$.

The following two rules describe the behaviour of the (**loop_fixpt_noacc**($n_+$) $C$) command. They say that the command repeatedly runs its body $C$ up to $n_+$ times, but unlike the usual loop, this repetition may stop early at some iteration $m+1$, and skip the remaining $n_+-(m+1)$ iterations; this early stopping happens when a fixed point is reached at the iteration $m$ (for the first time) and so the fixed-point check at the following iteration succeeds. Another important point is that the (**loop_fixpt_noacc**($n_+$) $C$) command does not accumulate the score maps computed during iterations, and keeps only the score map at the last iteration (which is $T_{m+1}$ in the case of early stopping and $T_{n_+}$ otherwise). Discarding the score maps from all the non-last iterations comes from our loop vectorisation; it is exploited by our translation for vectorisation. The last rule describes the behaviour of the **shift** command: for every variable $x^{\tau^\dagger}$ in the current state, it looks up the partial map $m_x$ stored in the variable $x^{\tau^\dagger}$, and shifts values in $m_x$ according to $\rho$ in the rule, using the $(\_)(\rho)$ operation (Fig. 8). When $\rho(i) = i'$, the definition of $\rho$ ensures that $i \in A_\downarrow$ and $i' \in A$ for the input antichain $A$ so that the source index of the shift is always in $A_\downarrow$ and the target index is in $A$.

## 4.2 Soundness of Translation

The next theorem states that our parallelising translation preserves the semantics.

**Theorem 1.** Let $C$ be a command in the source language, and $\bar{C}$ be the parallelising translation of $C$ in the target language. Let $D \in \text{RDB}$, $\sigma_0 \in \text{State}_{\text{src}}$, and $\sigma_1 \in \text{State}_{\text{tgt}}$. If $\sigma_1(x^{\tau^\dagger}) = [[\quad] \mapsto \sigma_0(x^\tau)]$ for all variables $x^\tau$ in the source language, there exist $\sigma'_0 \in \text{State}_{\text{src}}$, $\sigma'_1 \in \text{State}_{\text{tgt}}$, and $r \in \mathbb{R}$ such that

$C, D, \sigma_0 \llbracket \sigma'_0, r, \bar{C}, D, \sigma_1, [[\quad]] \llbracket_{tgt} \sigma'_1, [[\quad] \mapsto r]$, and $\sigma'_1(x^{\tau^\dagger}) = [[\quad] \mapsto \sigma'_0(x^\tau)]$ for all $x^\tau$.

Intuitively, Theorem 1 says that if a command $C$ in the source language and its translation $\bar{C}$ in the target language are run in the same states, then $C$ and $\bar{C}$ terminate successfully, and their final states and final scores are essentially the same. The proof of Theorem 1 is provided in the extended version [40].

The key point of the proof stems from the fact that $\bar{C}$ may store in a single target language state information about multiple source language states. Specifically, when $C$ contains a for-loop, $\bar{C}$ executes all the iterations of this loop in parallel. Thus, each of the states that arise during the execution of the translated loop in $\bar{C}$ corresponds to multiple states that arise during the execution of the original loop in $C$. The difficulty is to capture this many-to-one correspondence between states in the source and target languages, and to prove its invariance so as to establish the theorem. We address
the challenge by first embedding states and commands in the source language to those in the target language, then defining a form of simulation relation between states of the target language, which we call a **coupling** relation, and finally proving that modulo embedding, the relation is preserved by the execution of a command in the source language and that of its translation in the target language.

# 5 Implementation

We implemented our vectorisation technique for Pyro [9], a probabilistic programming language built on PyTorch [51]. Our choice of Pyro is motivated by its powerful support for parallel execution, which leverages PyTorch’s efficient tensor operations, tensor broadcasting, and automatic differentiation, and by its rich set of inference engines, including Stochastic Variational Inference (SVI) and variable elimination. We point out that our implementation is not specific to Pyro and can be adapted to other PPLs that support similar features. Beyond standard Pyro/PyTorch primitives, our design incorporates a non-standard runtime to handle variable reading and writing. This runtime reflects the SIMD-style semantics of expressions and assignment commands based on the antichains in our target language. We also define two novel primitives specifically for the target language’s control flow (ifz) and loop (for) constructs. These functionalities are integrated into a library that overrides standard Python variable access and provides new constructors. The key aspects of our implementation are summarised below.

**Representing finite partial maps with PyTorch tensors.** Our runtime relies on PyTorch tensors to represent finite partial maps from *Index* in our target language, such as V-tensors *T* and partial maps *m* stored in variables. This approach, however, needs to address several representation problems. One of these is that PyTorch uses integer sequences as indices, but our target language uses string-integer sequences as indices. To bridge this gap, our library maintains a global partial map from strings to integers, allocating a unique dimension for each string, as shown in Fig. 10. The other problem is that a finite partial map in our target language accepts indices of varying length, while a PyTorch tensor assumes indices of a fixed length. To address this problem, we track the maximum integer associated with each string present in the partial map’s domain. We then construct a PyTorch tensor where the maximum index of each tensor dimension is one larger than the tracked maximum integer of the corresponding string. For instance, in Fig. 10, a partial map *T* has strings *α* and *β* with the maximum integers 9 and 8, respectively. Therefore, it is represented by a PyTorch tensor of shape (11, 10), where the first dimension corresponds to *α* and the second to *β* by the global mapping. Here, the one additional component in each dimension of the PyTorch tensor, accessed by -1 in the figure, is used to represent the case that the corresponding string is absent in an index of the partial map. Each cell of the PyTorch tensor is then assigned the value evaluated by the *extend(T)* function at the corresponding index. For example, in Fig. 10, for all $0 \le i \le 9$ and $0 \le j \le 8$, the cell of the PyTorch tensor at index [*i*, *j*] is assigned the value of *extend(T)([*(α, i); (β, j)*)] = c*<sub>*i*,*j*</sub>*, the cell at index [*i*, -1] is set to the value of *extend(T)([*(α, i)*)] = b*<sub>*i*</sub>*, the cell at index [-1, *j*] is given the value of *extend(T)([*(β, *j*)]) = *a*, and the cell at index [-1, -1] is assigned to the value of *extend(T)([]) = *a*. This strategy allows us to effectively represent any finite partial map in the target language using a single PyTorch tensor.

**Antichains and PyTorch tensor operations.** To preserve the correspondence between finite partial maps and PyTorch tensors, our library overrides the standard variable read and write of Python in a way that respects the antichain. An antichain, being a finite set of indices, is itself represented by a PyTorch tensor, where each tensor cell holds a boolean value indicating whether the corresponding index is in the antichain. Based on this representation, our library reads a variable by constructing a PyTorch tensor consisting solely of elements at indices for which the current antichain evaluates to true. Subsequent operations (e.g., arithmetic) on PyTorch tensors are then performed element-wise, facilitated by PyTorch’s *broadcasting* semantics. The outcomes of these operations are then used to
<figure><img src="image_5.png" /><figcaption>Fig. 10. Correspondence between a PyTorch tensor (left) and a finite partial map *T* in our target language (right). The `strings_to_dim` mapping ensures this correspondence by allocating strings to unique dimensions.</figcaption></figure>

update the PyTorch tensor assigned to the variable, but in a more careful way: it updates not only the element at the current antichain's true indices, but also the elements at the upward closure of those indices, which can be also efficiently handled by broadcasting. These steps are implemented efficiently in PyTorch thanks to its support for broadcasting and to the preservation of the encoding of finite partial maps into PyTorch tensors across all the base operations of the target language.

**Variable elimination.** Discrete random variables pose difficulties for probabilistic inference algorithms that rely on gradient computation, such as SVI and HMC. A common solution, known as *variable elimination*, enumerates all the possible combinations of the values of these discrete random variables, computes a total score for each combination, and sums them up. While variable elimination is costly in general, Pyro efficiently performs it by allocating a dedicated tensor dimension for the enumeration of each discrete random variable, and leveraging PyTorch's broadcasting semantics and tensor-based parallelism to compute total scores. To support these optimised algorithms, our library internally manages metadata (e.g., dimension-variable mappings) for the correct handling of these additional dimensions for enumeration.

# 6 Experimental Evaluation

In this section, we report on the evaluation of probabilistic programs translated using our method, focusing on the following three research questions:

(RQ1) Consistency: Does our approach produce scores and their gradients that are consistent with those produced by existing methods (modulo floating-point errors)?

(RQ2) Time and memory: How does our approach compare to existing techniques in terms of execution time and peak memory usage?

(RQ3) Scalability: Does our approach scale with dataset size, assuming a fixed model?

*Benchmarks.* We used five classes of probabilistic models, with complementary features.

* Hidden Markov models (hmm-{ord1, ord2, neural}) [10, 48]. These models describe piano scores [12] using variants of a hidden Markov model (HMM). The models hmm-{ord1, ord2} are first- and second-order HMMs involving 3 nested loops but no neural networks, whereas hmm-neural is a first-order HMM with 2 nested loops and neural networks. In these models, the outermost loop iterates over 229 (hmm-{ord1, neural}) or 50 (hmm-ord2) sequences of music notes, the first nested loop stands for the HMM iteration over time steps of length 129, and the innermost loop of hmm-{ord1, ord2} iterates over piano keys of length 51. The first and third loops are easily parallelisable since each iteration is conditionally independent.

* Deep Markov model (dmm) [10, 36]. This model is another HMM variant describing the same music dataset. It is a first-order Markov model, consisting of two nested loops over 229 sequences of music notes and 129 time steps, respectively. It differs from hmm-* in two ways: first, dmm uses
continuous hidden states, while hmm-* uses discrete states; second, dmm uses neural networks for both prior and likelihood, while hmm-neural does so only for likelihood.

* *Nested hidden Markov models* (nhmm-{train, stock}). These first-order nested HMMs model the number of arrivals at a train station over 44,640 time steps (hours) [65], and the stock price of 10 companies over 2,560 time steps (trading days) [3], respectively. Both models consist of 4 nested loops, either over 5 years, 12 months, 31 days, 24 hours, or over 10 companies, 10 years, 4 quarters, 64 days. Only the outermost loops can be easily parallelised.

* *Autoregressive model* (arm). This variant of autoregressive model (AR) [61] describes speech data taken from [50]. While standard AR models have a constant order, the order of this model is a latent random variable. This model contains one loop iterating over 2,000 time steps.

* *Temperature controller model* (tcm) [39, 60]. This state-space model describes the evolution of the temperature inside a room, equipped with a stochastic temperature controller. It includes if-then-else statements, required to convey different types of temperature dynamics based on the controller state. It consists of two nested loops over 10 temperature sequences and 200 time steps, respectively, where the first loop can be easily parallelised.

* *Baseline methods*. Pyro features several loop vectorisation methods.

* *Plate constructor* [68]. It is a tensor-based loop constructor that vectorises loops without data dependencies. However, it may lead to incorrect inference results if the variables in a loop have data dependencies across iterations. Plates can be nested but cannot contain if-then-else statements.

* *Vectorised Markov constructor* [67]. It is a "functor"-based loop constructor that vectorises certain types of loops with data dependencies. It supports scalar or vector indices determined by the user-defined dependence order, and includes an efficient automatic variable elimination algorithm for enumeration. However, it cannot vectorise if-then-else statements and nested loops.

* *Discrete HMM distribution* [66]. It is a distribution class in Pyro that expresses a first-order discrete HMM. Given transition and emission matrices, it efficiently computes log-likelihood and posterior marginals in parallel using algorithms such as the forward-backward algorithm. However, it is limited to first-order discrete HMMs without if-then-else statements or nested loops.

* *Hand-coded vectorisation*. This baseline refers to the manual vectorisation of loops using tensor operations. As a result, applying it requires expertise in both the Pyro PPL and tensor-based programming, and often entails non-standard manipulation of the PPL's internal structures.

We compare our method (ours) with five reference methods that combine some of these features:

* seq: It is a sequential implementation of the model, where all loops are executed sequentially.

* plate: It uses only plate constructors to vectorise loops without data dependencies.

* vmarkov: It uses plate constructors and vectorised Markov constructors to vectorise loops.

* discHMM: It uses plate constructors and discrete HMM constructors to vectorise loops.

* hand: It refers to manual vectorisation of loops using tensor operations.

Note that all methods except seq have limitations and cannot be applied to all models. In particular, some benchmark models include if-then-else statements, nested loops, or higher-order dependencies that these methods cannot handle. We consider plate, vmarkov, or discHMM inapplicable to a model if no loop in the model can be vectorised by a plate constructor, a vectorised Markov constructor, or a discrete HMM constructor, respectively. Also, hand is considered inapplicable to a model when a loop in the model contains an if-then-else statement. This is because vectorisation using tensor operations in such cases is highly non-trivial. See the extended version [40] for further details on the applicability. We also note that the use of plate or vmarkov requires care since they are only sound for loops without data dependencies or when the dependence order is correctly specified by the user. See the extended version [40] for the code snippets of the model hmm-neural using our method and the baselines.
<table><caption>Table 1. Average relative errors of objective function and gradient values computed by plate, vmarkov, discHMM, and ours with respect to seq in SVI or MAP inference. “-” indicates the method is not applicable. Note that plate is not applicable to arm as it contains only a single loop with data dependencies.</caption><thead><tr><th rowspan="2"></th><th colspan="2">plate</th><th colspan="2">vmarkov</th><th colspan="2">discHMM</th><th colspan="2">ours</th></tr><tr><th>obj</th><th>grad</th><th>obj</th><th>grad</th><th>obj</th><th>grad</th><th>obj</th><th>grad</th></tr></thead><tbody><tr><td>hmm-ord1</td><td>0</td><td>3.5e-7</td><td>0</td><td>3.5e-7</td><td>8.7e-8</td><td>2.0e-6</td><td>0</td><td>3.4e-7</td></tr><tr><td>hmm-neural</td><td>2.0e-9</td><td>4.5e-6</td><td>5.0e-9</td><td>4.0e-6</td><td>4.0e-9</td><td>4.1e-6</td><td>5.0e-9</td><td>4.0e-6</td></tr><tr><td>hmm-ord2</td><td>0</td><td>2.7e-7</td><td>0</td><td>2.7e-7</td><td>-</td><td>-</td><td>0</td><td>2.6e-7</td></tr><tr><td>dmm</td><td>7.7e-9</td><td>3.1e-7</td><td>2.5e-8</td><td>3.1e-7</td><td>-</td><td>-</td><td>2.4e-8</td><td>3.1e-7</td></tr><tr><td>nhmm-train</td><td>0</td><td>8.1e-6</td><td>-</td><td>-</td><td>-</td><td>-</td><td>0</td><td>8.1e-6</td></tr><tr><td>nhmm-stock</td><td>0</td><td>5.4e-6</td><td>-</td><td>-</td><td>-</td><td>-</td><td>0</td><td>5.5e-6</td></tr><tr><td>arm</td><td>-</td><td>-</td><td>0</td><td>3.0e-9</td><td>-</td><td>-</td><td>0</td><td>3.0e-9</td></tr><tr><td>tcm</td><td>-</td><td>-</td><td>-</td><td>-</td><td>-</td><td>-</td><td>0</td><td>0</td></tr></tbody></table>

*Experimental setup.* We run three types of inference engines, Maximum a Posteriori (MAP) estimation, Stochastic Variational Inference (SVI), and Markov Chain Monte Carlo (MCMC) sampling. MAP estimation is performed for the models hmm-* and nhmm-* . SVI is run using RNN-based amortised variational distributions for dmm or independent normal variational distributions for arm and tcm. The inference engines for MAP estimation and SVI require solving an optimisation problem, where an objective function is the posterior probability density or the evidence lower bound. We use the Adam optimiser [35] to solve this problem, i.e., to train learnable parameters. MCMC is used for models without neural networks, i.e., hmm-{ord1, ord2}, nhmm-* , arm, and tcm. Specifically, we employ the Hamiltonian Monte Carlo (HMC) sampler [18, 19] with fixed number of leapfrog steps and step size to sample from the posterior distribution of continuous variables, while a discrete variable in arm is handled via Gibbs sampling [26] based on Metropolis-Hastings [28, 42] with a random walk proposal. All timings are measured using an NVIDIA 2080Ti GPU.

[RQ1] Consistency. To validate that our method yields correct outputs, we compared each of the vectorisation methods except hand against seq using SVI or MAP inference, in terms of (i) the objective function value, computed as the sum of score values at each score command, and (ii) the gradient of the objective function with respect to learnable parameters computed by automatic differentiation. We then compared these results across the different methods to assess their consistency with seq.

More precisely, we measured the relative errors between these values in $L_2$-norm. We computed the objective function or gradient values $v \in \mathbb{R}^d$ using seq, and $\hat{v} \in \mathbb{R}^d$ using one of the other vectorisation methods, over the first training step with the same random seed; and then computed the relative error $||v - \hat{v}||_2/||v||_2$. To alleviate the high computational cost of the seq method, we used only 10 of the outermost sequences of the datasets used in the hmm-* and dmm benchmarks, while the entire datasets were used for the other benchmarks.

Table 1 presents these relative errors averaged over 100 random seeds for each model and each vectorisation method. It demonstrates that the outputs from ours are the same as, or very close to the outputs from seq for all models considered. We believe that the non-zero errors come from the non-associativity of floating-point addition/multiplication. For objective function values, we manually verified that they are caused by different orderings of floating-point operations: hmm-neural and dmm involve matrix multiplications, and PyTorch computes the addition/multiplication operations there in different orders in ours and seq, due to different tensor dimensions in the computation. For gradient values, we conjecture that a similar reordering of operations is introduced by PyTorch during backpropagation, due to different tensor dimensions or non-contiguous memory layouts. We also note that similar discrepancies caused by these factors appeared across all other vectorisation
<figure><img src="image_6.png" /><figcaption>Fig. 11. Score time (s), total time (s), and GPU memory usage (GB) in each training step of SVI/MAP with 95% confidence intervals. “×” below the x-axis indicates inapplicability. Large values are truncated.</figcaption></figure>

methods. Overall, these observations suggest that our method yields consistent outputs, and the remaining errors are well within the bounds of expected floating-point behaviour.

**[RQ2] Time and Memory.** We compared the wall-clock time and GPU memory usage of our method with other baselines using SVI, MAP inference, and MCMC sampling. For SVI and MAP, we measured, for each training step, the time to compute score (score time) and the time to run the step entirely (total time). The total time includes components other than score computation: e.g., sampling from variational distributions, performing variable elimination, and computing gradients. For MCMC, we measured the time to execute each sampling step. We also measured the maximum GPU memory usage over a training step for SVI and MAP, and over a sampling step for MCMC. For SVI and MAP, every reported number is obtained by averaging over 5 runs, where each run consists of 10 training steps (seq) or 1,000 training steps (the others), and the first 1 out of 10 steps (seq) or 10 out of 1,000 steps (the others) are excluded from the averages, as considered burn-in steps. Concretely, we averaged over 45 data points for seq and 4,950 data points for the other methods. For MCMC, each reported number is the average over 8 runs, each with a different number of sampling steps performed within one hour (i.e., 8 × the average steps per hour in total). In this evaluation, the full datasets are used for all benchmarks. The main results are shown in Fig. 11 and Table 2, and the full results are in the extended version [40].

*Results on SVI/MAP time.* Compared to seq, ours shows significant speedups in both score time (36.2–31399.3×) and total time (54.8–14302.9×). These speedups are mainly enabled by two factors. First, in a score computation, ours drastically reduces the iteration counts required for score
computation (2 to 10.1) compared to seq's extensive iterations (2,000 to 1,506,591), achieved by the loop vectorisation with early termination. Second, ours consistently employs tensors, while seq uses a list of scalars. Many operations in Pyro, including variable elimination and sampling from variational distributions, are more efficient over tensors than separate scalars.

Compared to plate, ours again shows significant speedups in score time (4.5-145.8×) and total time (1.8-78.4×). These gains are also due to plate's high iteration count (129 to 8,928) and ours's efficient tensor operations. Furthermore, unlike ours, plate is not applicable to the models arm and tcm, as all the loops in these models have data dependencies or if-then-else statements.

Compared to vmarkov, ours is faster in both score time (2.0-6.0×) and total time (1.1-4.6×). These speedups stem from two main sources. First, ours exhibits a fewer number of iterations: for a loop with dependence order K, ours performs K + 1 iterations, while vmarkov requires 2K + 1 iterations. Second, vmarkov introduces additional overheads due to the use of "functors" (functional tensors). A conversion from PyTorch tensors to Pyro functors, followed by operations over functors, is much slower than direct operations over tensors. In addition, vmarkov is not applicable to the models nhmm-* and tcm, as they contain nested loops with dependencies or if-then-else statements.

Compared to discHMM, ours is slower in score time (0.3-0.4×) and total time (0.2-0.4×) for the models hmm-{ord1, neural}. However, discHMM only supports first-order HMMs without nesting and if-then-else statements, and thus cannot be applied to the other six models.

Compared to hand, ours is slower in score time (0.1-0.3×) and total time (0.2-0.9×) for all models except tcm where hand could not be applied due to the presence of if-then-else statements.

*Results on SVI/MAP memory usage.* Compared to seq, ours shows a decrease of memory usage for some models and an increase for others. For models nhmm-stock, arm, and tcm, ours significantly reduces peak GPU memory usage, requiring only 3-15% of the memory used by seq. For models hmm-{neural, ord2} and nhmm-train, ours requires 119-144% of the memory used by seq. Two effects play against each other here. Our approach uses antichains representing indices considered in a given state. In our current implementation, these antichains are represented as boolean tensors, which may be sparse, hence causing memory inefficiencies. However, even then, ours may benefit from more efficient memory allocation schemes implemented by GPUs, e.g., with allocation in large fixed-sized chunks. In contrast, seq allocates larger numbers of smaller chunks, which may increase memory overhead, especially when the maximum vector dimension used in the innermost loop is small. These factors explain that ours can sometimes achieve lower memory usage under specific conditions, although it seems to have a higher memory requirement in theory.

Compared to plate, ours reduces memory usage to 70% on nhmm-stock benchmark, due to the trade-off between usage of antichains and memory allocation schemes discussed above. On the other models, ours uses more memory, ranging from 119% to 284% of the memory used by plate.

Compared to vmarkov, ours uses less memory in all models except hmm-ord1. As we remarked above, ours stores antichains which require additional memory. However, vmarkov also induces an important memory overhead, due to additional execution of iterations with scalar indices and the implementation of functors. The results demonstrate that the overall overhead due to functor in vmarkov is higher than that generated by antichains in ours.

Compared to discHMM, ours uses more memory for the two models that the former can handle. However, as mentioned earlier, discHMM is specialised to a narrow class of models, so it could not be applied to the other six models.

Compared to hand, ours uses more memory for all models except tcm where hand could not be applied. Specifically, ours uses 518% and 200% of the memory used by hand on arm and hmm-ord1. For the others, ours uses 102-109% of the usage of hand, showing comparable memory usage.
<table><caption>Table 2. Speedup and relative GPU memory usage per MCMC sampling step, compared between ours and the others. Each entry shows the range over applicable benchmarks. See the extended version [40] for details.</caption><thead><tr><th></th><th>seq</th><th>plate</th><th>vmarkov</th><th>discHMM</th><th>hand</th></tr></thead><tbody><tr><td>MCMC speedup (other → ours)</td><td>67-1210x</td><td>1.83-65.63x</td><td>1.07-2.56x</td><td>0.28x</td><td>0.26-0.80x</td></tr><tr><td>MCMC memory usage (ours/other)</td><td>3-120%</td><td>69-253%</td><td>24-109%</td><td>562%</td><td>102-1290%</td></tr></tbody></table>

*Results on MCMC.* As shown in Table 2, the results are consistent with the observations from the SVI and MAP inference experiments across all benchmarks, except for hmm-neural and dmm benchmarks that we do not consider in MCMC experiments. Overall, ours exhibits superior sampling speed (i.e., lower time per sampling step) and broader applicability compared to the vectorisation baselines, often while maintaining comparable or lower memory consumption. For instance, compared to vmarkov, ours performs 1.07-2.56x more sampling steps, while using 24-109% of the GPU memory consumed by vmarkov, across all benchmarks except nhmm-* and tcm, for which vmarkov is inapplicable. The applicability of each method to the benchmarks is consistent with that observed in the SVI and MAP inference experiments, except for the case of seq on hmm-ord1, where it fails to complete a single step within one hour.

**[RQ3] Scalability.** To assess the scalability of our method, we considered the autoregressive model arm', where N denotes the number of samples and K the order of the model: $y_i \sim N(y_{i-1} + \dots + y_{i-K}, 1)$ for $i = 0, \dots, N-1$. Here $y_i$ is the i-th sample for $i = 0, \dots, N-1$, and the initial values are $y_{-K} = \dots = y_{-1} = 0$. We express the model in Pyro using a loop of length N to generate the $y_i$'s, which takes $K+1$ iterations to reach a fixed point when using our method.

We measured the time to compute score (score time) under two setups: use a fixed $K \in \{1, 50\}$ and vary $N$; and use a fixed $N \in \{1000, 4000\}$ and vary $K$. All samples were generated from the standard normal distribution, and the time for generating the samples was excluded from the measurement. For each $(N, K)$ pair, we averaged the time over 99 runs out of 100, excluding the first burn-in step. We compared the time taken by the vectorised loop produced by our method (ours) against the sequential loop (seq). The results are shown in Fig. 12.

When $K$ is fixed to a value significantly smaller than $N$ (i.e., 1 and 50), ours exhibits a notably slower increase rate with respect to $N$ compared to seq. The improvement is more significant when $K$ is small. This demonstrates that, when the dependence order of the loop is low, our vectorisation method can improve performance more substantially by leveraging the GPU parallelism.

In contrast, when $N$ is fixed (i.e., 1000 and 4000) and $K$ is close to $N$, the runtime of seq remains almost constant, while that of ours is quadratic in $K$. Indeed, when $K$ and $N$ are close to each other, both the number of iterations and the size of the vectors processed within each iteration grow proportionally to $K$. These findings suggest that our method may not be optimal for higher-order loops. Nevertheless, in many practical applications of autoregressive models, such as speech processing [8, 61], heart rate variability analysis [11], and image dynamic texture modelling [31], the value of $K$ is typically much smaller than $N$. To mitigate the performance degradation when $K$ approaches $N$, an adaptive strategy could be employed to switch between sequential and vectorised execution based on the loop dependence order (e.g., $K$), which our method can automatically detect.

## 7 Related Work

*Model translations in PPLs.* Several prior approaches have translated models from non-tensor-based PPLs into tensor-based languages, such as compiling Stan [17] models to Pyro or NumPyro models for improved inference performance with provable correctness [6], or converting high-level discrete probabilistic programs into NumPy [70] programs with optimised tensor operations (e.g., opt_einsum [59]) for exact inference over discrete random variables [49]. These approaches do
<figure><img src="image_7.png" /><figcaption>Fig. 12. Score time (s) of ours and seq for the model arm' with different N and K. Each point denotes the mean over 99 runs. 95% confidence intervals are omitted as their relative width to the mean is below 2%.</figcaption></figure>

not alter the structures of the original models as aggressively as our method does for loops; instead, they mostly rely on generic optimisation of the target language for performance improvements.

*Vectorisation primitives in PPLs.* Most tensor-based PPLs have introduced convenient features to vectorise independent operations in probabilistic models. A key example is the plate primitive in Pyro and NumPyro, which originates from graphical models [15] and is used in PPLs to vectorise loops, particularly when the random variables inside the loop are conditionally independent or when the loop has no data dependencies. Similarly, JAX provides vmap [13] and TensorFlow offers vectorized_map [2], both of which aid in vectorising function evaluation over independent tensor dimensions. These primitives also help vectorise MCMC sampling, improving inference efficiency across multiple chains [37]. While these primitives are effective for simple loops or sequences of independent operations, our work extends this capability to loops with more complex dependencies.

Other works in PPL aim at extending vectorisation primitives to handle more complex dependencies, similarly to our approach. For instance, TensorFlow Probability (TFP) [23] provides MarkovChain distributions in order to vectorise loops with short-range dependencies in models such as random walks or autoregressive models. Similarly, Pyro provides a family of hidden Markov model (HMM) distributions (e.g., DiscreteHMM, GaussianHMM) to vectorise key computations in HMMs, such as smoothing, filtering algorithms, and score computation [57]. However, these approaches usually require the restructuring of a model so as to make explicit use of specific vectorisation primitives, or have limitations related to the kind of dependencies they support (e.g., supporting only predefined transition distributions). More closely related to our work, Funsor [47] introduces vectorized_markov to vectorise loops with dependencies, but requires users to manually specify dependency lengths, and is not generally applicable to all dependency structures (e.g., nested structures or conditional branches). Our work extends these ideas by introducing a more general, automatic vectorisation method that ensures correctness through fixed-point check and eliminates the need for any additional user input or extensive model change. Moreover, it can handle complex dependencies, which may be nested or have dependencies determined at runtime.

*Optimisation of inference algorithms.* Our work is related to approaches that optimise parameter learning or posterior inference algorithms by exploiting properties of PPL models. A static analysis for the smoothness properties of PPL models has been developed to propose an SVI algorithm that exploits these properties to reduce variance in gradient estimation [38]. Other work focuses on conditional independence properties in PPL models and proposes a type system for capturing these properties, along with an optimised variable elimination algorithm that leverages them [27]. Meanwhile, methods have been introduced to automatically derive correct and efficient gradient estimators for a given SVI objective when it is expressed as a program, by transforming it into gradient estimation code [7]. Our method does not directly optimise learning or inference algorithms; instead, it optimises the execution of probabilistic programs by vectorising loops with dependencies, which in turn can accelerate various inference algorithms relying on efficient program execution.
*Automatic vectorisation in general-purpose languages.* Our work is also related to studies in general-purpose programming languages that aim at automating vectorisation. For example, dynamic data dependency analysis techniques have been proposed to infer SIMD parallelisation potential by identifying all possible dependency-preserving reorderings, together with further analysis of memory access patterns [30]. Additionally, methods have been developed to automatically batch operations in dynamic neural network libraries [45], such as DyNet [44], and to handle data-dependent control flow and recursion in programs using a pointer per batch element to manage variables and program counters [55]. PyTorch's framework for automatic batching [14] leverages mask propagation and control flow rewriting to simplify minibatching in dynamic models. Recently, Autovesk [62] automates vectorisation by transforming scalar instructions into vectorised instructions over a graph of operations, using heuristics to minimise the overall instruction numbers. Although these methods rely on static or dynamic analysis of data dependencies to exploit vectorisation opportunities, our work avoids the need for complex analysis by an automated fixed-point check that ensures correctness of vectorisation.

There are also approaches that use thread-level speculation (TLS) to parallelise general sequential programs. R-LRPD [21], for instance, speculatively executes loop iterations in parallel and records read/write bit arrays to re-execute dependent iterations. GPU-TLS [72] and STMLite [41] both focus on re-execution of mis-speculated iterations, and dynamic analysis to ensure correctness. Other methods improve performance by integrating TLS with transactional memory (TM) [5] or adopting in-place speculation [46]. However, these methods typically discard partial results on mis-speculation during re-execution, which can cause a large number of re-executions and slow convergence, particularly in loops with many inter-iteration dependencies, as observed in most of the models in our experiments. In contrast, our method can extract useful partial information even from such mis-speculated iterations through the *shift* construct, achieving the faster convergence in those cases, as demonstrated by our experimental results.

# 8 Conclusion and Limitations

We have presented an automatic sound method for the vectorisation of loops over a fixed range, as often found in PPLs. It is based on a speculative execution of the body of for-loops, using a non-standard semantics, and on a fixed-point check to ensure correctness of the final result. Its implementation upon Pyro proceeds by translation into a lower level PPL with primitives for vectorisation. We proved this correctness of the method and demonstrated its effectiveness in accelerating inference for several families of complex probabilistic models including sequence models and hierarchical models, and for various probabilistic inference engines.

While our work leads to performance improvements in our experiments, it also has limitations. First, its fixed-point check introduces a performance overhead that might not always be compensated for when the iteration count is not substantially reduced by our method—for instance, due to long data dependencies, as discussed in §6. Although such cases are not expected to be dominant in practice, a potential avenue for improvement lies in exploring a hybrid approach that dynamically decides whether to apply vectorising optimisation or fall back to a standard runtime based on program characteristics. Second, our work does not consider the vectorisation of sample generation, which is central to many inference algorithms (e.g., importance sampling, Metropolis-Hastings, and Sequential Monte Carlo [22]). Our method instead focuses on the vectorisation of density computations, which is orthogonal to sample generation and is central to other inference algorithms (e.g., SVI and HMC). Addressing this limitation would yield a more comprehensive vectorisation framework and further accelerate a broader class of inference algorithms.
## Acknowledgements

We thank the anonymous reviewers for their insightful comments and suggestions that helped us improve this paper. This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korean Government (MSIT) (No. RS-2023-00279680), by the Inria Associated Team AISAPPL (Abstract Interpretation-based Static Analysis for Probabilistic Programming Languages), by the French ANR Grant EMASS, and by the French PEPR SecurEval programe.

## References

[1] Oriol Abril-Pla, Virgile Andreani, Colin Carroll, Larry Dong, Christopher J Fonnesbeck, Maxim Kochurov, Ravin Kumar, Junpeng Lao, Christian C Luhmann, Osvaldo A Martin, et al. 2023. PyMC: a modern, and comprehensive probabilistic programming framework in Python. *PeerJ Computer Science* 9 (2023), e1516. doi:10.7717/peerj-cs.1516

[2] Ashish Agarwal. 2019. Static automatic batching in TensorFlow. In *International Conference on Machine Learning (ICML)*. 92–101. https://proceedings.mlr.press/v97/agarwal19a.html

[3] Ran Aroussi. 2020. yfinance: Download market data from Yahoo! Finance's API. https://github.com/ranaroussi/yfinance
Accessed: 2025-07-10.

[4] Sudipto Banerjee, Bradley P Carlin, and Alan E Gelfand. 2003. *Hierarchical Modeling and Analysis for Spatial Data*. Chapman and Hall/CRC. doi:10.1201/9780203487808

[5] Joao Barreto, Aleksandar Dragojevic, Paulo Ferreira, Ricardo Filipe, and Rachid Guerraoui. 2012. Unifying thread-level speculation and transactional memory. In *International Middleware Conference (Middleware)*. 187–207. doi:10.1007/978-3-642-35170-9_10

[6] Guillaume Baudart, Javier Burroni, Martin Hirzel, Louis Mandel, and Avraham Shinnar. 2021. Compiling Stan to generative probabilistic languages and extension to deep probabilistic programming. In *ACM Conference on Programming Language Design and Implementation (PLDI)*. 497–510. doi:10.1145/3453483.3454058

[7] McCoy R. Becker, Alexander K. Lew, Xiaoyan Wang, Matin Ghavami, Mathieu Huot, Martin C. Rinard, and Vikash K. Mansinghka. 2024. Probabilistic Programming with Programmable Variational Inference. *Proc. ACM Program. Lang.* 8, PLDI, Article 233 (June 2024), 25 pages. doi:10.1145/3656463

[8] Maria A Berezina, Daniel Rudoy, and Patrick J Wolfe. 2010. Autoregressive modeling of voiced speech. In *International Conference on Acoustics, Speech and Signal Processing (ICASSP)*. 5042–5045. doi:10.1109/ICASSP.2010.5495058

[9] Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul A. Szerlip, Paul Horsfall, and Noah D. Goodman. 2019. Pyro: Deep Universal Probabilistic Programming. *Journal of Machine Learning Research* 20, 28 (2019), 1–6. http://jmlr.org/papers/v20/18-403.html

[10] Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul A. Szerlip, Paul Horsfall, and Noah D. Goodman. 2019. Pyro: Deep Universal Probabilistic Programming with Python and PyTorch. https://github.com/pyro-ppl/pyro
Accessed: 2025-07-10.

[11] Anita Boardman, Fernando Soares Schlindwein, and Ana Paula Rocha. 2002. A Study on The Optimum Order of Autoregressive Models for Heart Rate Variability. *Physiological Measurement* 23, 2 (2002), 325. doi:10.1088/0967-3334/23/2/308

[12] Nicolas Boulanger-Lewandowski, Yoshua Bengio, and Pascal Vincent. 2012. Modeling temporal dependencies in high-dimensional sequences: application to polyphonic music generation and transcription. In *International Conference on Machine Learning (ICML)*. 1159–1166. https://icml.cc/2012/papers/590.pdf

[13] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018. *JAX: composable transformations of Python+NumPy programs*. http://github.com/jax-ml/jax
Accessed: 2025-07-10.

[14] James Bradbury and Chunli Fu. 2018. Automatic batching as a compiler pass in PyTorch. In *Workshop on Systems for ML at NeurIPS 2018*. http://learningsys.org/nips18/assets/papers/107CameraReadySubmissionMatchbox_LearningSys_Abstract_%20(2).pdf

[15] Wray L Buntine. 1994. Operations for learning with graphical models. *Journal of Artificial Intelligence Research* 2 (1994), 159–225. doi:10.1613/jair.62

[16] Ricardo Cannizzaro, Michael Groom, Jonathan Routley, Robert Osazuwa Ness, and Lars Kunze. 2025. COBRA-PPM: A Causal Bayesian Reasoning Architecture Using Probabilistic Programming for Robot Manipulation Under Uncertainty. *arXiv:2403.14488* (2025). doi:10.48550/arXiv.2403.14488

[17] Bob Carpenter, Andrew Gelman, Matthew Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. 2017. Stan: A Probabilistic Programming Language. *Journal of Statistical Software* 76, 1 (2017), 1–32. doi:10.18637/jss.v076.i01

[18] Wenlin Chen, Mingtian Zhang, Brooks Paige, José Miguel Hernández-Lobato, and David Barber. 2024. Diffusive Gibbs Sampling. In *International Conference on Machine Learning (ICML)* (Proceedings of Machine Learning Research, Vol. 235).
PMLR, 7731-7747. https://proceedings.mlr.press/v235/chen24be.html

[19] Wenlin Chen, Mingtian Zhang, Brooks Paige, José Miguel Hernández-Lobato, and David Barber. 2024. *Diffusive Gibbs Sampling*. https://github.com/Wenlin-Chen/DiGS Accessed: 2025-10-20.

[20] Peter D Congdon. 2019. *Bayesian Hierarchical Models: With Applications Using R*. Chapman and Hall/CRC. doi:10.1201/9780429113352

[21] F. Dang, Hao Yu, and L. Rauchwerger. 2002. The R-LRPD test: speculative parallelization of partially parallel loops. In *International Parallel and Distributed Processing Symposium (IPDPS)*. doi:10.1109/IPDPS.2002.1015493

[22] Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. 2006. Sequential Monte Carlo Samplers. *Journal of the Royal Statistical Society Series B: Statistical Methodology* 68, 3 (05 2006), 411–436. doi:10.1111/j.1467-9868.2006.00553.x

[23] Joshua V Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. 2017. Tensorflow distributions. *arXiv:1711.10604* (2017). doi:10.48550/arXiv.1711.10604

[24] Simon Duane, A.D. Kennedy, Brian J. Pendleton, and Duncan Roweth. 1987. Hybrid Monte Carlo. *Physics Letters B* 195, 2 (1987), 216–222. doi:10.1016/0370-2693(87)91197-X

[25] Hong Ge, Kai Xu, and Zoubin Ghahramani. 2018. Turing: a language for flexible probabilistic inference. In *International Conference on Artificial Intelligence and Statistics (AISTATS)*. 1682–1690. http://proceedings.mlr.press/v84/ge18b.html

[26] Stuart Geman and Donald Geman. 1984. Stochastic Relaxation, Gibbs Distributions, and the Bayesian Restora-
tion of Images. *IEEE Transactions on Pattern Analysis and Machine Intelligence* PAMI-6, 6 (1984), 721–741.
doi:10.1109/TPAMI.1984.4767596

[27] Maria I. Gorinova, Andrew D. Gordon, Charles Sutton, and Matthijs Vákár. 2021. Conditional Independence by Typing. *ACM Trans. Program. Lang. Syst.* 44, 1, Article 4 (Dec. 2021), 54 pages. doi:10.1145/3490421

[28] W Keith Hastings. 1970. Monte Carlo sampling methods using Markov chains and their applications. *Biometrika* 57, 1 (04 1970), 97–109. doi:10.1093/biomet/57.1.97

[29] Matthew D. Hoffman, David M. Blei, Chong Wang, and John Paisley. 2013. Stochastic variational inference. *J. Mach. Learn. Res.* 14, 1 (2013), 1303–1347. https://jmlr.org/papers/v14/hoffman13a.html

[30] Justin Holewinski, Ragavendar Ramamurthi, Mahesh Ravishankar, Naznin Fauzia, Louis-Noël Pouchet, Atanas Rountev, and P Sadayappan. 2012. Dynamic trace-based analysis of vectorization potential of applications. In *ACM Conference on Programming Language Design and Implementation (PLDI)*. 371–382. doi:10.1145/2254064.2254108

[31] Midori Hyndman, Allan D Jepson, and David J Fleet. 2007. Higher-order Autoregressive Models for Dynamic Textures. In *British Machine Vision Conference (BMVC)*. doi:10.5244/C.21.76

[32] Krishna Murthy Jatavallabhula, Miles Macklin, Dieter Fox, Animesh Garg, and Fabio Ramos. 2023. Bayesian Object Models for Robotic Interaction with Differentiable Probabilistic Programming. In *Conference on Robot Learning (CoRL)*. 1563–1574. https://proceedings.mlr.press/v205/jatavallabhula23a.html

[33] Marc Kéry and Kenneth F Kellner. 2024. *Applied Statistical Modelling for Ecologists: A Practical Guide to Bayesian and Likelihood Inference Using R, JAGS, NIMBLE, Stan and TMB*. Elsevier. doi:10.1016/C2022-0-02766-5

[34] Jack Kiefer and Jacob Wolfowitz. 1952. Stochastic estimation of the maximum of a regression function. *The Annals of Mathematical Statistics* (1952), 462–466. doi:10.1214/aoms/1177729392

[35] Diederik P. Kingma and Jimmy Ba. 2015. Adam: A Method for Stochastic Optimization. In *International Conference on Learning Representations (ICLR)*. doi:10.48550/arXiv.1412.6980

[36] Rahul Krishnan, Uri Shalit, and David Sontag. 2017. Structured Inference Networks for Nonlinear State Space Models. In *AAAI Conference on Artificial Intelligence (AAAI)*. doi:10.1609/aaai.v31i1.10779

[37] Junpeng Lao, Christopher Suter, Ian Langmore, Cyril Chimisov, Ashish Saxena, Pavel Sountsov, Dave Moore, Rif A Saurous, Matthew D Hoffman, and Joshua V Dillon. 2020. tfp.mcmc: Modern Markov chain Monte Carlo tools built for modern hardware. *arXiv:2002.01184* (2020). doi:10.48550/arXiv.2002.01184

[38] Wonyeol Lee, Xavier Rival, and Hongseok Yang. 2023. Smoothness Analysis for Probabilistic Programs with Application to Optimised Variational Inference. *Proc. ACM Program. Lang.* 7, POPL, Article 12 (Jan. 2023), 32 pages. doi:10.1145/3571205

[39] Wonyeol Lee, Hangyeol Yu, and Hongseok Yang. 2018. Reparameterization Gradient for Non-differentiable Models. In *Annual Conference on Neural Information Processing Systems (NeurIPS)*. https://dl.acm.org/doi/10.5555/3327345.3327459

[40] Sangho Lim, Hyoungjin Lim, Wonyeol Lee, Xavier Rival, and Hongseok Yang. 2025. Optimising Density Computations in Probabilistic Programs via Automatic Loop Vectorisation. *arXiv:2511.11070* (2025). doi:10.48550/arXiv.2511.11070

[41] Mojtaba Mehrara, Jeff Hao, Po-Chun Hsu, and Scott Mahlke. 2009. Parallelizing sequential applications on commodity hardware using a low-cost software transactional memory. In *ACM Conference on Programming Language Design and Implementation (PLDI)*. 166–176. doi:10.1145/1542476.1542495

[42] Nicholas Metropolis, Arianna W. Rosenbluth, Marshall N. Rosenbluth, Augusta H. Teller, and Edward Teller. 1953. Equation of State Calculations by Fast Computing Machines. *The Journal of Chemical Physics* 21, 6 (06 1953), 1087–1092. doi:10.1063/1.1699114
[43] Radford M Neal et al. 2011. MCMC using Hamiltonian dynamics. *Handbook of Markov Chain Monte Carlo* 2, 11 (2011), 2. doi:10.1201/b10905

[44] Graham Neubig, Chris Dyer, Yoav Goldberg, Austin Matthews, Waleed Ammar, Antonios Anastasopoulos, Miguel Ballesteros, David Chiang, Daniel Clothiaux, Trevor Cohn, Kevin Duh, Manaal Faruqui, Cynthia Gan, Dan Garrette, Yangfeng Ji, Lingpeng Kong, Adhiguna Kuncoro, Gaurav Kumar, Chaitanya Malaviya, Paul Michel, Yusuke Oda, Matthew Richardson, Naomi Saphra, Swabha Swayamdipta, and Pengcheng Yin. 2017. DyNet: The Dynamic Neural Network Toolkit. arXiv:1701.03980 (2017). doi:10.48550/arXiv:1701.03980

[45] Graham Neubig, Yoav Goldberg, and Chris Dyer. 2017. On-the-fly operation batching in dynamic computation graphs. In *Annual Conference on Neural Information Processing Systems (NeurIPS)*. https://dl.acm.org/doi/10.5555/3294996.3295153

[46] Cosmin E Oancea, Alan Mycroft, and Tim Harris. 2009. A lightweight in-place implementation for software thread-level speculation. In *Annual Symposium on Parallelism in Algorithms and Architectures (SPAA)*. 223–232. doi:10.1145/1583991.1584050

[47] Fritz Obermeyer, Eli Bingham, Martin Jankowiak, Du Phan, and Jonathan Chen. 2020. Functional Tensors for Probabilistic Programming. arXiv:1910.10775 (2020). doi:10.48550/arXiv:1910.10775

[48] Fritz Obermeyer, Eli Bingham, Martin Jankowiak, Neeraj Pradhan, Justin Chiu, Alexander Rush, and Noah Goodman. 2019. Tensor Variable Elimination for Plated Factor Graphs. In *International Conference on Machine Learning (ICML)*, Vol. 97. 4871–4880. https://proceedings.mlr.press/v97/obermeyer19a.html

[49] Jingwen Pan and Amir Shaikhha. 2023. Compiling discrete probabilistic programs for vectorized exact inference. In *International Conference on Compiler Construction (CC)*. 13–24. doi:10.1145/3578360.3580258

[50] Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. 2015. Librispeech: an Asr Corpus based on Public Domain Audio Books. In *International Conference on Acoustics, Speech and Signal Processing (ICASSP)*. 5206–5210. doi:10.1109/ICASSP.2015.7178964

[51] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In *Annual Conference on Neural Information Processing Systems (NeurIPS)*. 8024–8035. https://dl.acm.org/doi/10.5555/3454287.3455008

[52] Du Phan, Neeraj Pradhan, and Martin Jankowiak. 2019. Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro. arXiv:1912.11554 (2019). doi:10.48550/arXiv:1912.11554

[53] Matti Pärssinen, Mikko Kotila, and Ilkka Sillanpää. 2021. Probabilistic Programming Method for Time-Series Forecasting of COVID-19 Cases Based on Empirical Data. *American Journal of Epidemiology and Infectious Disease* 9 (2021), 18–23. doi:10.12691/ajeid.9.1-4

[54] Song S Qian, Mark R DuFour, and Ibrahim Alameddine. 2022. Bayesian Applications in Environmental and Ecological Studies with *R* and *Stan*. Chapman and Hall/CRC. doi:10.1201/9781351018784

[55] Alexey Radul, Brian Patton, Dougal Maclaurin, Matthew Hoffman, and Rif A. Saurous. 2020. Automatically batching control-intensive programs for modern accelerators. In *Machine Learning and Systems (MLSys)*, Vol. 2. 390–399. https://proceedings.mlsys.org/paper_files/paper/2020/file/39945d578f616735572174bf5e8f155d-Paper.pdf

[56] Herbert Robbins and Sutton Monro. 1951. A stochastic approximation method. *The Annals of Mathematical Statistics* (1951), 400–407. doi:10.1214/aoms/1177729586

[57] Simo Särkkä and Ángel F García-Fernández. 2020. Temporal parallelization of Bayesian smoothers. *IEEE Trans. Automat. Control* 66, 1 (2020), 299–306. doi:10.1109/TAC.2020.2976316

[58] Viktor Senderov, Jan Kudlicka, Daniel Lundén, Viktor Palmkvist, Mariana P Braga, Emma Granqvist, David Broman, and Fredrik Ronquist. 2023. TreePPL: a universal probabilistic programming language for phylogenetics. *bioRxiv* (2023), 2023-10. doi:10.1101/2023.10.10.561673

[59] Daniel G. A. Smith and Johnnie Gray. 2018. opt_einsum - A Python package for optimizing contraction order for einsum-like expressions. *Journal of Open Source Software* 3, 26 (2018), 753. doi:10.21105/joss.00753

[60] Sadegh Esmaeil Zadeh Soudjani, Rupak Majumdar, and Tigran Nagapetyan. 2017. Multilevel Monte Carlo Method for Statistical Model Checking of Hybrid Systems. In *International Conference on Quantitative Evaluation of Systems (QEST)*. 351–367. doi:10.1007/978-3-319-66335-7_24

[61] Gintautas Tamulevičius and Jonas Kaukenas. 2017. High-order Autoregressive Modeling of Individual Speaker's Qualities. In *IEEE Workshop on Advances in Information, Electronic and Electrical Engineering (AIEEE)*. doi:10.1109/AIEEE.2017.8270551

[62] Hayfa Tayeb, Ludovic Paillat, and Bérenger Bramas. 2023. Autovesk: Automatic Vectorized Code Generation from Unstructured Static Kernels Using Graph Transformations. *ACM Transactions on Architecture and Code Optimization* 21 (2023), 1–25. doi:10.1145/3631709

[63] David Tolpin, Jan-Willem van de Meent, Hongseok Yang, and Frank Wood. 2016. Design and Implementation of Probabilistic Programming Language Anglican. In *Symposium on the Implementation and Application of Functional*
Programming Languages (IFL). Article 6. doi:10.1145/3064899.3064910

[64] Dustin Tran, Matthew D. Hoffman, Dave Moore, Christopher Suter, Srinivas Vasudevan, Alexey Radul, Matthew Johnson, and Rif A. Saurous. 2018. Simple, Distributed, and Accelerated Probabilistic Programming. In *Annual Conference on Neural Information Processing Systems (NeurIPS)*. https://dl.acm.org/doi/10.5555/3327757.3327859

[65] Bay Area Rapid Transit. 2024. Ridership Reports. https://www.bart.gov/about/reports/ridership Accessed: 2025-07-10.

[66] Uber Technologies, Inc. 2025. Pyro Documentation: Distributions (DiscreteHMM). https://docs.pyro.ai/en/dev/distributions.html#pyro.distributions.DiscreteHMM Accessed: 2025-07-10.

[67] Uber Technologies, Inc. 2025. Pyro Documentation: Funsor-Based Pyro (vectorized_markov). https://docs.pyro.ai/en/dev/contrib.funsor.html#pyro.contrib.funsorvectorized_markov Accessed: 2025-07-10.

[68] Uber Technologies, Inc. 2025. Pyro Documentation: Primitives (plate). https://docs.pyro.ai/en/dev/primitives.html#pyro.primitives.plate Accessed: 2025-07-10.

[69] Stéfan van der Walt, S. Chris Colbert, and Gaël Varoquaux. 2011. The NumPy Array: A Structure for Efficient Numerical Computation. *Computing in Science & Engineering* 13, 2 (2011), 22–30. doi:10.1109/MCSE.2011.37

[70] Stefan van der Walt, S. Chris Colbert, and Gael Varoquaux. 2011. The NumPy Array: A Structure for Efficient Numerical Computation. *Computing in Science & Engineering* 13, 2 (2011), 22–30. doi:10.1109/MCSE.2011.37

[71] Frank Wood, Andrew Warrington, Saeid Naderiparizi, Christian Weilbach, Vaden Masrani, William Harvey, Adam Ścibior, Boyan Beronov, John Grefenstette, Duncan Campbell, and S. Ali Nasseri. 2022. Planning as Inference in Epidemiological Dynamics Models. *Frontiers in Artificial Intelligence* 4 (2022). doi:10.3389/frai.2021.550603

[72] Chenggang Zhang, Guodong Han, and Cho-Li Wang. 2013. GPU-TLS: An Efficient Runtime for Speculative Loop Parallelization on GPUs. In *International Symposium on Cluster, Cloud, and Grid Computing (CCGrid)*. 120–127. doi:10.1109/CCGrid.2013.34

[73] Walter Zucchini and Iain L MacDonald. 2009. *Hidden Markov Models for Time Series: An Introduction Using R*. Chapman and Hall/CRC. doi:10.1201/9781420010893

Received 2025-07-08; accepted 2025-11-06