Proof obligations for jump threading and block merging

The jump threading and block merging transformations of a program’s control flow graph do not change the ordering of the program’s non-control statement execution at all executions of the program. Both transformations are correct since a terminator does not produce any observable effects, and, therefore, the edges which are removed during the transformations are not traceable.

The idea of block merging is to remove an edge between a block and a successor which has only one predecessor, thus, removing one jump from the path on each such edge. Jump threading duplicates a block for a particular input edge when the result of the conditional branch at the end of the block is known for the block; the duplicated block is then changed in a way that instead of the conditional branch there is an unconditional jump to the appropriate successor of the block. Consequently, the path through the duplicated block differs from the path through the source only in that the former performs an unconditional jump instead of evaluating the conditional branch.

The two transformations mutually reinforce each other since jump threading produces blocks with only one input which are then absorbed by the block merging, whereas LLVM’s JumpThreading pass merges a block into its predecessor in order to thread the condition through the predecessors of that predecessor [1]. In the sections below, there is a simulation argument for each transformation, followed by a description of an OCaml program which checks a threaded and merged graph against the original graph.

Setting

Let us look at a control flow graph that has an entry block $\mathit{entry}$. The body of each block $b$ consists of $|b|$ straight-line instructions followed by a terminator which is either an unconditional jump or a two-way branch on a register, or a return. A state $(b, i, \rho, m)$ represents the current block $b$, the index $i$ of the next instruction along with the register environment $\rho$ and the memory $m$. The condition $i = |b|$ says that the next step is the terminator. In one step, either one instruction or one terminator is executed and gives rise to a finite trace $w$ of observable events that is empty unless there are some external calls or volatile accesses.

Correctness follows CompCert: a transformation is expected to preserve the property that whenever the source program does not run into error, all observable behaviors of the target program are also observable behaviors of the source program, no matter whether the behavior is terminating or non-terminating. Leroy proves the inclusion via forward simulation combined with the determinism of the target semantics. The forward simulation ensures that each observable behavior of the source has a corresponding behavior of the target and the deterministic target has exactly one behavior to give, so that the behavior of the source must match the behavior of the target. Since both rewrites do not introduce any choice in the target, all claims below will use the forward direction because the relation between states is easier to define.

The forward simulation is a relation $\mathcal{R}$ between states of the source and target programs which relates initial states and is preserved by each step of the source. The most general definition of the relation also allows the source to make an extra step which has no counterpart in the target program and the measure $\mu$ on the source states strictly decreases. Since $\mu$ takes values in a well-founded order, the source program makes only finitely many such steps before the target program has to make one step. CompCert's proof of branch tunneling is equivalent to the diagram below for each step of the LTL language, with the measure coming from the node count.

Divergence is introduced into the equation via the second part of this diagram, as there is the possibility that even with the use of a measure, an infinite number of steps performed by the source which does not make any progress can be attached to a single target state that does not make any progress either. Leroy makes use of this problem to validate the measure since in its absence, the simulation will certify a target that terminates, while the source will diverge [3]. The jump threading procedure needs only the first part of the diagram, in that all steps of the copy correspond to only one step of the original block, while the block merging requires the second, because the jump does not have any counterpart in the target.

Jump threading

The case that motivates the transformation appears in Figure 2 where block $A$ branches on $c$ and both arms rejoin at $J$. That block merges $u$ and $v$ through a φ and performs a load before it branches on $c$ a second time, and since nothing on the way from $A$ to that second branch changes $c$ the edge $(B, J)$ is certain to reach $T$ while the edge $(C, J)$ is certain to reach $F$. Mueller and Whalley describe this situation as a path on which the result of a conditional branch is already known, and they avoid the branch by replicating code along that path [5]. Threading the edge $(p, b)$ creates a copy $b_p$ whose terminator is a jump to the successor selected on that edge, redirects $p$ to $b_p$ and leaves every other predecessor of $b$ pointing at the original. Threading both edges into $J$ therefore leaves it with no predecessors so that it can be deleted, and every run then executes one conditional branch on $c$ where it previously executed two.

The semantic obligation splits into two parts that are easy to conflate, and the difference matters because a sound analysis can still justify an unsound thread. The first part is that some fact $\varphi_{(p,b)}$ holds on every execution of the edge $(p, b)$, which is an ordinary soundness property of whatever analysis produced it, and in the example that fact is the outcome of the branch in $A$ carried along the path through $B$. The second part is that the fact survives the body of $b$ and still decides the branch when control reaches the terminator. It is the entailment below where $\mathrm{wp}$ is the weakest precondition of the straight line body, and a failure of either part lets the copy jump to a successor that the source would not have taken.

In a language with mutable registers the body can kill the fact by redefining $c$ or one of its operands. SSA form rules that out because a name cannot be redefined, so the danger moves to a condition computed inside $b$ from a load whose memory the fact does not describe. The weakest precondition of an assignment substitutes the assigned expression for the register in the postcondition, so taking it across the whole body pushes the branch condition backwards until it mentions only values that are available on entry to $b$. The same entailment therefore covers the other common case in which $c$ is computed in $b$ from φ nodes whose operands on the edge from $p$ are constants, because the precondition then reduces to a statement about those incoming operands that evaluation alone can decide.

The simulation relation for the copy strengthens equality of states with the weakest precondition of the part of the body that has yet to execute. In it $\theta$ maps each value defined in $b$ to its counterpart in $b_p$ and every block other than the copies is related to itself by equality. Both sides share the index $i$ and the memory $m$ because the copy holds the same instructions in the same order, so only its terminator and the names of its values can differ.

The entailment establishes that strengthening on entry to $b_p$ and each instruction of the copy preserves it because a state satisfying the weakest precondition of a suffix steps to a state satisfying the weakest precondition of the rest of that suffix, and at $i = |b|$ the strengthening says exactly that the source branch goes to the successor to which the target jumps. Every source step is therefore matched by exactly one target step as in the lock step diagram on the left of Figure 1, and threading needs no measure at all. The invariant cannot be avoided because the source state at the start of $b$ does not determine which target block is executing, since that depends on the edge by which control arrived and the fact carried in from $p$ is the only link between the two. Equality of states alone would hold just as well for a copy whose jump pointed at the wrong successor, and the proof would then break at the terminator where the two programs first disagree.

SSA repair

With mutable registers the copy can reuse every name and so $\theta$ is the identity and the relation above is plain equality of environments. SSA forbids a second definition of a name and so every value defined in $b$ receives a fresh name in $b_p$, after which a use that was dominated by its single definition in $b$ can be reached from both definitions while being dominated by neither.

The load in Figure 2 is the smallest instance because $y$ was defined in $J$ and used in $R$, and after threading the two copies define $y_1$ and $y_2$ and $R$ needs a new φ to choose between them. The φ belongs at $R$ because that block lies in the dominance frontier of both copies, since each copy dominates one predecessor of $R$ without dominating $R$ itself. Restoring the dominance property in general places such φ nodes at the iterated dominance frontier of the blocks that now define the value, which is the placement problem Cytron et al. solve during construction, and LLVM's JumpThreading hands exactly this problem to an SSAUpdater after it clones a block to thread an edge [1, 6].

The φ nodes of $b$ itself have to be handled at both ends of the threaded edge because that edge now enters the copy instead of the original. In the copy each φ collapses to its operand from $p$, which is why $x$ disappears from panel (b) of Figure 2 and $T$ multiplies $u$ directly. In the original each φ loses its operand for $p$ because $p$ is no longer a predecessor, and once every incoming edge has been threaded as it is for $J$ those φ nodes have no operands left and disappear together with the block. Every successor $s$ of $b$ acquires $b_p$ as a new predecessor, so each φ in $s$ needs an operand for the edge $(b_p, s)$ naming the copy of whatever value it received from $b$, and forgetting that operand leaves a φ that no longer matches the predecessor list of its block.

The relation inherits all of this as a renaming that depends on the path. Along the path through $J_B$ the source value $y$ corresponds to $y_1$ in the target until control reaches $R$ where it corresponds to the new φ, while along the path through $J_C$ the same source value corresponds to $y_2$. No single renaming serves both paths, so $\theta$ has to be indexed by the edge on which the copied region was entered and extended at each join where a φ was placed. None of that bookkeeping changes the shape of the argument because each instruction still preserves the relation once $\theta$ is fixed for the current path, although it accounts for most of the length of any mechanised version.

Undefined values

Everything so far assumed that evaluating $c$ twice in the same state gives the same answer, and that assumption fails in an IR where a value may be resolved independently at each use. Threading sits on the harmless side of that failure because forcing the second branch to agree with the first removes executions and never adds one. In Figure 2 a source run could take $B$ at $A$ and then $F$ at $J$ if $c$ resolved differently at the two branches, while the threaded target has no such run because $J_B$ always continues to $T$. If the source may resolve each branch independently then the behaviours of the target form a subset of the behaviours of the source and refinement still holds.

LLVM declares branching on an undefined value to be undefined behaviour and provides freeze to turn undef or poison into an arbitrary but fixed value that every use of that freeze is guaranteed to observe identically [7]. JumpThreading folds a branch on freeze undef to whichever successor it prefers only when the freeze has exactly one use, because choosing a successor also chooses the frozen value and any other use would have to observe that same choice [1]. The unsound direction appears as soon as the copy is also used to propagate equalities, which is what GVN does when it replaces $w$ by $y$ inside the true arm of a branch on $t = y$ where $t$ and $w$ both equal $x + 1$. Lee et al. observe that if $y$ is poison while $w$ is not then the replacement passes poison to a use that received a defined value in the source. That replacement is sound only if branching on poison is undefined behaviour, because the source run then has no defined behaviour for the target to preserve, while loop unswitching assumed the opposite since it hoists a branch out of a loop and evaluates its condition on runs that might never have reached it. Lee et al. resolve the conflict by making branching on poison undefined and freezing the condition that unswitching hoists [8].

Loops and reducibility

Under the conditions above threading always preserves behaviour but it can still destroy the loop structure that later passes depend on. Threading one of the two entry edges of a loop header $H$ on which its branch is decided produces a copy $H'$ that jumps straight into the body. The cycle through $H$, $B$ and $L$ in Figure 3 then has two entries at $H$ from $P_2$ and at $B$ from $H'$, and since neither entry dominates the other the region is irreducible in the sense of Hecht and Ullman [9]. Any pass that finds loops by looking for a single header that dominates the whole cycle then stops seeing it as a loop at all because an irreducible cycle is not a natural loop. Threading the back edge instead gives the latch its own copy $H_L$ of the header so that the original header runs once on entry and the loop becomes the cycle through $B$, $L$ and $H_L$ with $B$ as its only entry, which amounts to peeling the header and leaves the graph reducible.

LLVM approximates loop headers by the targets of back edges and by default refuses to thread across any of them, which rules out both rewrites in Figure 3 including the one through the back edge that would have kept the graph reducible, so by default the pass gives up the rewrite that preserves reducibility along with the one that destroys it. The comment on findLoopHeaders concedes that this also forbids profitable cases such as threading a jump from the preheader or a back edge directly to an exit block, because its simple analysis cannot track those properties and keep them current as the graph mutates [1].

Block merging

Merging applies to an edge $(a, b)$ when the conditions below hold, and it replaces $a$ by a block $a'$ whose body is $\mathrm{body}(a) \cdot \mathrm{body}(b)$ and whose terminator is that of $b$ as drawn in panel (a) of Figure 4. The jump that ended $a$ disappears along with $b$ itself, and every predecessor of $a$ now enters $a'$ instead, so the only edges that move are those that used to leave $b$ and now leave $a'$.

The successor condition means that $a$ ends in a terminator whose only target is $b$, so that terminator can be dropped between the two bodies. The predecessor condition ensures that no edge still lands at the start of $b$ once it has gone, and the entry condition protects the implicit edge from the caller. The last condition excludes a block that jumps to itself, which under the other three can only be an unreachable self loop because its sole predecessor is itself and it is not the entry, so there is nothing to concatenate. A predecessor set computed from terminators also misses any edge that is not written in a terminator, which is why LLVM additionally refuses to merge a block whose address is taken and used, since an indirect branch through that address is an edge into $b$ that no predecessor lists [1]. SSA adds a renaming obligation of its own because every φ in $b$ has a single operand from $a$ and must be replaced by it, and every successor of $b$ that names $b$ as an incoming block in a φ must name $a$ instead because the edge it describes now leaves the merged block. Every copy has one predecessor and one successor after threading and so the example merges $B$, $J_B$ and $T$ into a single block and does the same on the other arm. That leaves panel (b) of Figure 4 where $A$ branches once and each arm runs straight into $R$, whose φ nodes now name $B$ and $C$.

The relation maps the prefix of $a'$ to $a$ and the suffix to $b$ with identity everywhere else, and $\mu$ is one at the jump $(a, |a|)$ and zero at every other state. Registers and memory are equal on both sides because the merged block runs exactly the instructions of $a$ and then those of $b$ in their original order and only drops the jump between them, which changes neither.

The jump at the end of $a$ and the first instruction of $b$ are both related to the target state at index $|a|$ of $a'$, so the source step between them is matched by no target step at all. Its trace is empty and the measure falls from one to zero, which is the diagram on the right of Figure 1, while every other source step is matched one for one by the target step at the corresponding index and the proof reduces to a case split on whether the source sits at $(a, |a|)$. The measure is trivially bounded in this case because the predecessor condition allows at most one deleted jump between any two target steps.

Dropping the predecessor condition gives tunnelling, which redirects every jump into a block consisting only of a branch straight to that block's eventual target, and there the measure stops being trivial. The source may cross an arbitrarily long chain of such blocks for a single target jump, so $\mu$ has to bound the remaining length of the chain, and a chain that closes into a cycle has no finite length. CompCert's tunnelling pass on LTL avoids computing that length directly by building a union find structure in which every block of the form branch s is placed in the class of s and then rewriting each successor to the representative of its class, so a cycle of branches ends up as a single branch from its representative to itself and the silent divergence survives the rewrite [4]. A rewrite that shortcut the cycle would instead let the target reach code that the source never executes while it spins forever, and the checker later in this post rejects precisely such a target. The same pass also turns a conditional branch whose two successors share a representative into an unconditional branch, which is the degenerate form of threading in which the outcome of the branch no longer matters.

Checking the result

Both arguments have the same shape once they are stated over individual instructions, which suggests checking each transformed graph against its source in the style of the verified validator of Tristan and Leroy or the bounded checks that Alive2 applies to individual runs of LLVM optimisations [10, 11]. A checker of that kind catches a bug in the transformation on the graph where it occurs, so an incorrect rewrite is rejected before it can reach the generated code.

The code below does this for a register transfer language in which each node holds one instruction and names its successors, so a block is only a maximal chain of nodes and block merging becomes the deletion of Goto nodes. The transformation supplies revmap, which sends each target node to the source node it was copied from. The checker requires every target instruction to equal the source instruction at its image with two exceptions that correspond to the two proofs. A target successor may skip a chain of source Goto nodes bounded by the number of source nodes so that a chain which revisits a node is rejected, and a source Cond may become a target Goto when a must analysis over the target graph knows the value of the condition register at that node:

type node = int
type reg = int

type instr =
  | Op of reg * string * reg list * node
  | Cond of reg * node * node
  | Goto of node
  | Return of reg

type graph = { entry : node; code : (node, instr) Hashtbl.t }

module Facts = Map.Make (Int)

let successors = function
  | Op (_, _, _, s) | Goto s -> [ s ]
  | Cond (_, t, f) -> [ t; f ]
  | Return _ -> []

let transfer i facts s =
  match i with
  | Op (d, _, _, _) -> Facts.remove d facts
  | Cond (c, t, f) when t <> f -> Facts.add c (s = t) facts
  | Cond _ | Goto _ | Return _ -> facts

let meet =
  Facts.merge (fun _ x y ->
      match (x, y) with Some x, Some y when x = y -> Some x | _ -> None)

let known_branches g =
  let facts = Hashtbl.create 16 and work = Queue.create () in
  Hashtbl.replace facts g.entry Facts.empty;
  Queue.add g.entry work;
  while not (Queue.is_empty work) do
    let n = Queue.pop work in
    Option.iter
      (fun i ->
        let here = Hashtbl.find facts n in
        List.iter
          (fun s ->
            let out = transfer i here s in
            let next, changed =
              match Hashtbl.find_opt facts s with
              | None -> (out, true)
              | Some old ->
                  let m = meet old out in
                  (m, not (Facts.equal Bool.equal m old))
            in
            if changed then begin
              Hashtbl.replace facts s next;
              Queue.add s work
            end)
          (successors i))
      (Hashtbl.find_opt g.code n)
  done;
  facts

let rec reaches g fuel n target =
  n = target
  || fuel > 0
     &&
     match Hashtbl.find_opt g.code n with
     | Some (Goto m) -> reaches g (fuel - 1) m target
     | _ -> false

let validate src tgt revmap =
  let facts = known_branches tgt and fuel = Hashtbl.length src.code in
  let follows s s' =
    match Hashtbl.find_opt revmap s' with
    | Some n -> reaches src fuel s n
    | None -> false
  in
  let known n' c = Option.bind (Hashtbl.find_opt facts n') (Facts.find_opt c) in
  let matches n' i' =
    match (Option.bind (Hashtbl.find_opt revmap n') (Hashtbl.find_opt src.code), i') with
    | Some (Op (d, o, a, s)), Op (d', o', a', s') ->
        d = d' && o = o' && a = a' && follows s s'
    | Some (Cond (c, t, f)), Cond (c', t', f') ->
        c = c' && follows t t' && follows f f'
    | Some (Cond (c, t, f)), Goto s' -> (
        match known n' c with
        | Some true -> follows t s'
        | Some false -> follows f s'
        | None -> false)
    | Some (Goto s), Goto s' -> follows s s'
    | Some (Return r), Return r' -> r = r'
    | _ -> false
  in
  follows src.entry tgt.entry
  && Hashtbl.fold (fun n' i' ok -> ok && matches n' i') tgt.code true

The fuel in reaches is the computational content of the measure, since a target step that lands at s' is matched by the source step to s followed by the silent Goto steps from s to the image of s' and the relation holds throughout with the remaining chain length as $\mu$. That length is finite because a walk longer than the node count would have to repeat a node, and a chain of Goto nodes that repeats a node is a cycle from which the source never escapes.

Computing the facts on the target graph looks backwards at first but is what the forward relation needs. The relation already requires equal registers and the fact is an inductive invariant of target states at the node where the copy's jump replaced the source branch, so equal registers carry it across to the source branch being matched, and meet keeps a fact only when every incoming path agrees on it, which is what makes the analysis a must analysis. Registers are compared by name because this language allows multiple assignments and so copies share names and $\theta$ is the identity, which is also why duplication on CompCert's RTL needs no repair at all when RTL is not in SSA form.

Written as RTL with the φ replaced by assignments to a shared register and the jump from $B$ to $J$ kept as a separate Goto node, the example from Figure 2 is accepted after threading and merging. The checker rejects the same target with one copy's jump pointed at the wrong arm because follows cannot reach the image of that arm from the successor the known fact selects. It rejects the pair in which the load writes the condition register in both programs before the second branch because transfer drops the fact at the load. A target that skips a two node Goto cycle to reach a Return that the source never executes is rejected as well, since the walk in reaches goes round the cycle until its fuel runs out. The checker is not itself verified here, although a proof in the style of Tristan and Leroy would consist mostly of the two lemmas already stated, one for each exception, together with the soundness of the must analysis.

Cost

Each thread copies $|b|$ instructions to save one dynamic branch on the threaded path, so the code size is paid once per copy while the saving depends on how often the threaded edge actually runs, which a size threshold alone does not measure. LLVM caps the size of a block it will duplicate at six instructions by default and three for functions marked minsize, with jump-threading-threshold overriding both [1]. That cap bounds a single copy and says nothing about how many copies are needed for the paths that replication has to separate.

When $d$ diamonds branching on independent conditions $c_1, \ldots, c_d$ are followed by a second test of each condition in turn, every later test needs its own fact to be known at the program point just after the last diamond. That point therefore has to be copied once for each of the $2^d$ combinations of outcomes before every second test can be removed, and the growth is exponential in the number of correlated tests even when every block involved is small. Bodík, Gupta and Soffa detect interprocedural paths along which a branch outcome is known with a demand driven analysis and eliminate the branch by restructuring that replicates code along those paths [12].

References

1. The LLVM jump threading pass, including its merging of single predecessor blocks, duplication threshold, loop header approximation, SSA update and folding of branches on frozen values, is implemented in llvm/lib/Transforms/Scalar/JumpThreading.cpp.
2. Xavier Leroy describes the CompCert verified compiler and the semantic preservation property it establishes for safe source programs in Formal Verification of a Realistic Compiler, published in Communications of the ACM 52(7) in 2009.
3. Xavier Leroy gives the simulation diagrams, the measure that rules out infinite stuttering and the passage from forward simulation to backward simulation for deterministic targets in A Formally Verified Compiler Back-end, published in the Journal of Automated Reasoning 43(4) in 2009.
4. The branch tunnelling pass of CompCert is defined in backend/Tunneling.v and proved correct with a measured forward simulation in backend/Tunnelingproof.v.
5. Frank Mueller and David B. Whalley replicate code along paths on which the result of a conditional branch is known in Avoiding Conditional Branches by Code Replication, presented at PLDI 1995.
6. Ron Cytron, Jeanne Ferrante, Barry K. Rosen, Mark N. Wegman and F. Kenneth Zadeck place φ nodes at iterated dominance frontiers in Efficiently Computing Static Single Assignment Form and the Control Dependence Graph, published in ACM Transactions on Programming Languages and Systems 13(4) in 1991.
7. The LLVM Language Reference Manual states that branching on an undefined value is undefined behaviour and that all uses of the value returned by one freeze observe the same value.
8. Juneyoung Lee, Yoonseung Kim, Youngju Song, Chung-Kil Hur, Sanjoy Das, David Majnemer, John Regehr and Nuno P. Lopes introduce freeze and make branching on poison undefined behaviour in Taming Undefined Behavior in LLVM, presented at PLDI 2017 and published in ACM SIGPLAN Notices 52(6).
9. Matthew S. Hecht and Jeffrey D. Ullman characterise reducible flow graphs in Flow Graph Reducibility, presented at STOC 1972.
10. Jean-Baptiste Tristan and Xavier Leroy prove correct a validator that checks the output of an unverified implementation of lazy code motion in Verified Validation of Lazy Code Motion, presented at PLDI 2009 and published in ACM SIGPLAN Notices 44(6).
11. Nuno P. Lopes, Juneyoung Lee, Chung-Kil Hur, Zhengyang Liu and John Regehr describe bounded translation validation of LLVM optimisations in Alive2: Bounded Translation Validation for LLVM, presented at PLDI 2021.
12. Rastislav Bodík, Rajiv Gupta and Mary Lou Soffa eliminate conditional branches along interprocedural paths on which their outcome is known in Interprocedural Conditional Branch Elimination, presented at PLDI 1997 and published in ACM SIGPLAN Notices 32(5).