Branch Prediction as a Reinforcement Learning Problem: Why, How and Case Studies

A number of tabular RL methods exist; most popular ones include TD-learning [15], SARSA [14], Q-Learning [17] and double Q-Learning [6]. Here we focus on the $Q$-Learning algorithm that provides specific convergence guarantees [17]³³3We have experimented with most of the tabular methods and found their branch prediction performance to be similar..

$Q$-Learning stores the $Q$-values $Q(s,a)$ for every state and action pair in a fixed-sized table. Given a state $s$ from the environment, $Q$-Learning predicts the action greedily using the policy $\operatorname{\pi_{\footnotesize\text{greedy}}}\left(s\right)$. The $Q$-Learning update rule works as follows: (a) choose an $a\in\mathcal{A}$ based on the current state using the $\varepsilon$-greedy policy; (b) observe reward $r$, next state $s^{\prime}$ after action $a$; (c) update $Q(s,a)\mathrel{{+}{=}}\alpha[r+\gamma\max_{a^{\prime}\in{\mathcal{A}}}Q(s^{\prime},a^{\prime})-Q(s,a)]$ for a positive learning rate parameter $\alpha$.

In function approximation methods, the policy is parametrized using a vector parameter $\bm{\theta}$ and is denoted by $\pi_{\bm{\theta}}(a|s)$. The goal is to optimize the expected cumulative reward by maximizing

where $d^{\pi}(s)$ is the stationary distribution⁴⁴4In general, a branch predictor might never converge to a stationary distribution due to the partial-observability of the state space. As an example, consider data-dependent branches whereby branch outcomes are correlated with data values but the actual data values are not visible to the branch predictor. of the Markov Chain defined by $\pi_{\bm{\theta}}$ on a given MDP. Using gradient ascent, $\bm{\theta}$ can be updated towards the direction of the gradient $\nabla J(\bm{\theta})$ to find an (local) optimal $\bm{\theta}$ that maximizes the average future cumulative reward, i.e., $\bm{\theta}\leftarrow\bm{\theta}+\alpha\nabla J(\bm{\theta})$.

As a representative example, the REINFORCE Policy-Gradient method [15, 18] works as follows:

• •

  Collect reward $r$ on state $s\in\mathcal{S}$ with action $a\in\mathcal{A}$.

• •

  Update policy by maximizing $J(\bm{\theta})$ via an optimization strategy (i.e., online gradient ascent).

Given a discrete action space (for example of size two), the policy is usually parametrized by a softmax function $\pi_{\bm{\theta}}(a|s):={\textrm{e}}^{h(s,a,\bm{\theta})}/\sum_{a\in{\mathcal{A}}}{\textrm{e}}^{h(s,a,\bm{\theta})}$ for an arbitrary differentiable function $h(s,a,\bm{\theta})$.

The second step of the Policy-Gradient method requires a gradient computation. In Section IV-B, a special case will be considered where $h$ is a linear function on $\bm{\theta}$, i.e., $h(s,a,\bm{\theta}):=\bm{\theta}^{\top}\bm{x}(s,a)$, where $\bm{x}(s,a)$ is a vector representation of the state-action space and $\bm{\theta}^{\top}\bm{x}$ denotes vector dot-product. In this specific case, the policy gradient theorem [15, Chapter 13] gives us a closed form solution for the gradient, i.e., $\nabla J(\bm{\theta})=\bm{x}(s,a)-\bm{\mu}$, where $\bm{\mu}:=\sum_{a\in\mathcal{A}}\pi_{\bm{\theta}}(a|s)\bm{x}(s,a)$.

In this section, we show that branch prediction is a task that can be seamlessly formulated as an RL problem as it abides by similar theoretical principles. In particular, through the lenses of RL, branch prediction is viewed as a pure online optimization task where the high-level goal is to minimize the rate of branch mispredictions. In RL parlance: “Maximization of the future cumulative reward equals to minimization of future cumulative branch mispredictions.”

The branch predictor corresponds to the agent and the processor’s execution environment (e.g., program counter, registers, memory/cache content, etc) corresponds to the environment. The environment is private and communicates its state to the branch-predictor agent so that the latter can decide which action $a$ from the action space $\mathcal{A}:=\{\text{T}{},\text{NT}{}\}$⁵⁵5T stands for taken, NT for not-taken. will be followed. Depending on the correctness of the chosen action, the environment responds back with a reward $R\in{\mathbb{R}}$ that guides the predictor’s update (e.g., confidence counter increase or reset on a, respectively, correct or incorrect prediction). At that point, the next state of the environment is also disclosed. Generally, the state content can consist of any information available to the BP; e.g., the branch PC address, the local and/or global branch history, etc. For instance, if the global branch history is part of the state definition, its instance after a prediction (next state) will include the predicted branch outcome (note that this is the common case of speculative history update in BP).

In analogy with RL methods, branch predictors can be categorized into two separate classes: tabular and function-approximation models. Tabular predictors model their actual predictions with a set of finite-state machine (FSMs) represented by up/down saturating counters stored in storage-constrained tables. Their goal is to couple branches with the correct FSM by matching patterns in the execution history. Examples of such predictors include the Bimodal predictor [13], two-level branch predictors, such as gshare [10], and, most recently, the TAGE-based [12] predictors.

Meanwhile, existing branch predictors that fall under the function-approximation models are mainly variants of the Perceptron [7] predictor, which is based on the homonymous online learning algorithm [1]. Perceptron learns from previous experiences by training a single-layer neural network, which stands in contrast to the memoization approach of tabular models. O-GEHL [11], Hashed Perceptron [16], and, most recently, Multiperspective Perceptron [8], are examples of this class of predictors.

In Figure 1, we classify common branch predictors into the two categories of tabular and function-approximation models. We respectively map them to their RL-based counterparts of either $Q$-Learning or Policy-Gradient origin. In essence, tabular predictors that are based on $N$-bit saturating counters can be formulated with a $Q$-Learning based scheme, while function-approximation models are a good fit for Policy-Gradient methods. We explain this mapping for $Q$-Learning based predictors in the next section; Policy-Gradient based schemes are analyzed later in Section IV.

We now give a detailed example of how an existing tabular branch predictor can be cast in RL terms. For illustrative purposes, we focus on relatively simple predictors – Bimodal and gshare (Section IV-B discusses a more complex Perceptron predictor). We opt not to include TAGE in our exploration as its design principles originate from the method of partial string matching used in data compression [4]. However, a similar RL-based optimization of TAGE could be feasible; we leave this case for future work.

Both Bimodal and gshare can be viewed as $Q$-Learning instances; for Bimodal the state space is the set of branch PC addresses, while for gshare it is the set of branch PC addresses exclusive-ORed with (a part of) the global branch-history register (GHR). In this simple form, the RL-based versions of these predictors do not perform any exploration (see section IV-D), that is, $\varepsilon=0$. According to the description of section II-A, their update rule will be reduced to $Q(s,a)=(1-\alpha)Q(s,a)+\alpha r$, since future rewards are not considered, and thus $\gamma=0$.

We have implemented in the CBP-5 [3] framework the aforementioned $Q$-Learning based variant of gshare, which we call G-QLAg (Global-History $Q$-Learning Agent). G-QLAg consists of a $N$-entry table, each containing two $Q$-values: $Q$_T and $Q$_NT. $Q$-values are signed floating-point numbers ranging from $-1$ to $1$. Experimentally, we found that using 6-bit $Q$-values is a good trade-off in our setup. Like gshare, G-QLAg is indexed with a hash of the branch PC and GHR modulo the size of the prediction table. For the matching entry, when $Q\textsubscript{\text{T}}>Q\textsubscript{\text{NT}}$, the prediction is T. Similarly, when $Q\textsubscript{\text{T}}<Q\textsubscript{\text{NT}}$, the prediction is NT. In the case where $Q\textsubscript{\text{T}}=Q\textsubscript{\text{NT}}$, the predicted direction (T, NT) is picked at random with equal probability. $Q$-values are initialized to $0$. At update time, the reward $r$ is set to $1$ or $-1$ if the prediction was correct or incorrect, respectively. The $Q$-value that corresponds to the prediction ($\max(Q\textsubscript{\text{T}},Q\textsubscript{\text{NT}})$) is updated according to the formula in the previous paragraph with learning rate $\alpha=0.2$.

Figure 2 compares the mispredictions per kilo-instructions (MPKI) of gshare and G-QLAg by varying the global branch-history length. We consider the same storage budget (64KB) for both predictors and configure their number of entries accordingly. Our implementation of G-QLAg uses 12 bits (2 $Q$-values) per entry, resulting in significantly fewer entries than gshare, which requires only 2 bits per entry.

As the figure shows, the two predictors perform similarly across the range of histories. As expected, for both predictors, MPKI reduces as history length is increased. However, at longer histories, the capacity disadvantage of G-QLAg becomes more pronounced, resulting in higher MPKI than gshare.

As the study shows, it is indeed straightforward to cast existing branch predictors as RL, and can be done without considering RL aspects such as $\gamma$ and $\varepsilon=0$ for G-QLAg. However, as we show in section IV-A, considering these and other aspects of RL opens new venues for BP policy optimization.

Following the Policy-Gradient framework explained in Section II-B, we next discuss how RL can facilitate the development of future BPs. We identify four key ingredients that a branch predictor needs to specify: (1) policy, (2) state representation, (3) loss function optimized by the predictor and (4) optimization strategy.

Non-linear models, including multi-layer neural networks, are fundamentally more expressive than linear ones and have the capacity to capture non-linear correlations. The limitation of non-linear models, however, is that they do not land themselves to a straight-forward microarchitectural implementation. For instance, the recent BranchNet predictor [19] attempts to overcome implementation limitations with a specialized software-hardware co-design: a plurality of convolutional neural networks (CNNs) are trained at compile-time (for up to 18 hours depending on the application) for a few hard-to-predict (H2P) branches. BranchNet’s inference engine is implemented in hardware for predicting only those branches, alongside TAGE. While the design is shown to be effective, it is complex. We believe that finding the right balance between model simplicity and its capacity will be a major focus in future ML-inspired microarchitectural research.

Table I summarizes these design elements for a selection of well-known branch predictors, as well as PolGAg, a Policy-Gradient based branch predictor introduced in the next section.

In this section, we study the Perceptron BP through the perspective of Policy Gradient, with the aim of showcasing how the RL construction can be used to explore the BP design-space.

To that end, we design a Perceptron-like branch predictor, which we call Policy Gradient Agent BP (PolGAg). The Policy Gradient Agent is based on the REINFORCE algorithm [15], a well-known Policy-Gradient method. Specifically, we adjust the REINFORCE algorithm for the context of the Perceptron BP. The resulting design is described with Algorithm 1. We derive Algorithm 1 from a holistic optimization view of the problem, i.e., maximizing the objective function of Equation 2 and making specific decisions on the definition of the policy $\pi_{\bm{\theta}}(a|s)$.

PolGAg represents a specific design point with a policy, state space, loss function and optimization strategy (Section IV-A). It suffices to define the policy and the state space $\mathcal{S}$, since the REINFORCE method already specifies the loss function (Equation 2) and the optimization strategy (online gradient ascent). For the policy, we use a linear softmax policy function with $h(s,a,\bm{\theta})=\bm{\theta}^{\top}\bm{x}(s,a)$. Recall that $\bm{x}(s,a)$ is a vector representation of the state-action space. The state space $\mathcal{S}$ is defined as the set of all branch addresses (PC) and the content of GHR with length $l$. GHR is viewed as a vector $\bm{q}\in{\{\pm 1\}}^{l}$. For the purpose of the analysis, assume that there are $p$ unique branch addresses during the whole program execution. Also, let $\text{OHE}(PC)\in\{0,1\}^{p}$ be the one-hot encoding⁷⁷7In the one-hot-encoding representation, each coordinate corresponds to a unique address. $\text{OHE}(PC)$ has $1$ in the coordinate that corresponds to PC and $0$, otherwise. (OHE) representation of the branch address. Moreover, define $\bm{q}(\text{T}):=[1,\bm{q}]$ and $\bm{q}(\text{NT})=-\bm{q}(\text{T})$, where the constant $1$ corresponds to a bias term. Finally, define $\bm{x}(s,a)$ as the $(p(l+1))$-vector $\bm{x}(s,a)=\text{OHE}(PC)\otimes\bm{q}$, where $\otimes$ denotes the Kronecker product. $\bm{x}(s,a)$ is defined for every state and action. Using the above policy and state representation, the REINFORCE method translates to a specific update step of our PolGAg Algorithm. Continuing from Section II-B, a direct calculation shows that $\nabla J(\bm{\theta})=\bm{x}(s,a)-\bm{\mu}=2\pi(\bar{a}|s)\bm{x}(s,a)$, where $\bar{a}$ flips the action from T to NT or vice-versa. Therefore, Step 5 of Algorithm 1 becomes $\bm{\theta}\leftarrow\bm{\theta}+2\alpha r\pi(\bar{a}|s)\bm{x}(s,a)$.

In Figure 3, we compare the MPKI of Perceptron [7] and PolGAg predictors. To avoid aliasing effects, we consider both predictors with unbounded storage, i.e., each branch has a dedicated set of weights. Firstly, for both predictors, each weight entry is stored using 8-bits. For PolGAg, we use 8-bit floats (1-bit sign, 5-bit exponent, 2-bit Mantissa) and for Perceptron 8-bit integers. As such, PolGAg has more precision by employing floats, but also a higher computation complexity. However, as this is a limit study, we ignore the higher computational requirements of PolGAg.

As the figure shows, PolGAg has a much higher prediction accuracy than the Perceptron. We attribute the difference to two factors: (a) the use of floats for PolGAg and (b) the use of logistic loss and policy gradient update versus Perceptron’s hinge loss update. Unfortunately, it is not easy to quantify the effect of each of these factors separately, since the choice of the loss function and update rule is coupled. In terms of using floats vs ints, the update rule of PolGAg involves multiplication with a prediction probability, which makes it unfriendly to an integer representation. Instead, we assess a variant of PolGAg that uses float32, rather than float8, for each weight. We observe that the use of higher precision further improves PolGAg’s accuracy, reducing MPKI by a full point compared to PolGAg-float8, which leads us to conclude that numerical precision does play a role in predictor’s efficacy.

To summarize, we showed that the Perceptron predictor can be viewed as a Policy-Gradient method by appropriate selection of policy, state space, loss function and optimization strategy. While the resulting predictor, PolGAg, was shown to outperform the baseline Perceptron, the former may be implementation-unfriendly due to its use of floats. Rather than espousing one design or the other, our purpose is to highlight the fact that viewing branch predictors through the RL prism opens new optimization opportunities.

We briefly show two more branch predictors that can be viewed in the Policy-Gradient framework:

RL methods allow the agent to explore unknown states, potentially by making a decision likely to be disadvantageous in the short term, in order to achieve a higher cumulative long-term reward. We posit that such an approach could benefit existing or new BP schemes. In the RL-based BP paradigm, the exploration process would effectively require the agent (predictor) to consider information from wrong execution paths. But how to explore a wrong execution path without greatly affecting processor performance?

We observe that on a misprediction, the BP continues to make predictions (and, accordingly, advance processor’s control flow) up to the point of the pipeline flush. In today’s BP designs, whatever was learned by the BP on the mispredicted path is discarded. However, the path observed under a misprediction may contain information valuable in later phases of the execution. In RL terms, there exists an opportunity to improve prediction accuracy through exploration and learning under a misprediction. In the case of $Q$-Learning, for example, to fully apply its update rule (Sec. II-A), it would be necessary to consider in the predictor’s update the first branch address on the wrong path. In this way, the agent would be able to perform exploration of the control-flow path according to the definition of the $\varepsilon$-greedy update step. To the best of our knowledge, the only proposed branch predictor that considers a similar option is the prophet/critic hybrid scheme [5]. We believe that the RL formulation opens the door to more systematic studies of exploration.