[llvm] [BOLT][AArch64] Expand cmpbr when reversing would overflow (PR #202998)

Alexandros Lamprineas via llvm-commits llvm-commits at lists.llvm.org
Tue Aug 4 05:44:59 PDT 2026


https://github.com/labrinea updated https://github.com/llvm/llvm-project/pull/202998

>From 754ab44a4e7fa5304fbada343cc8c172c8579f7b Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Thu, 4 Jun 2026 10:57:50 +0100
Subject: [PATCH 1/6] [BOLT][AArch64] Transform cmpbr ~> cmp + br when
 inversion not possible (#185731)

When reordering blocks we may have to invert branches. Sometimes this
isn't possible for compare-and-branch instructions because the immediate
value would overflow/underflow after the adjustment. In such cases I am
splitting the instruction into a compare followed by a branch. For this
to be legal we should be sure that the condition flags are not being
clobbered. Liveness analysis may help here.
---
 bolt/include/bolt/Core/BinaryFunction.h       |   4 +-
 bolt/include/bolt/Core/MCPlusBuilder.h        |  11 +-
 bolt/include/bolt/Passes/LongJmp.h            |   4 +-
 bolt/include/bolt/Utils/CommandLineOpts.h     |   4 +
 bolt/lib/Core/BinaryFunction.cpp              |   7 +-
 bolt/lib/Passes/BinaryPasses.cpp              |  16 +-
 bolt/lib/Passes/LongJmp.cpp                   |  27 +++-
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   |  90 +++++++++--
 bolt/lib/Target/AArch64/CMakeLists.txt        |   6 +-
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  |   6 +-
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |   6 +-
 bolt/lib/Utils/CommandLineOpts.cpp            |   6 +
 .../AArch64/compare-and-branch-inversion.S    | 118 ++++++++++-----
 bolt/unittests/Core/CMakeLists.txt            |   1 +
 bolt/unittests/Core/MCPlusBuilder.cpp         | 140 ++++++++++++++++--
 15 files changed, 365 insertions(+), 81 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 84fbd5661fd0a..51651e41497da 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -65,6 +65,8 @@ class DWARFUnit;
 
 namespace bolt {
 
+class DataflowInfoManager;
+
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
 /// Types of macro-fusion alignment corrections.
@@ -2463,7 +2465,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches();
+  void fixBranches(DataflowInfoManager *DIM = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 84b44a9ab5483..c98b56d14523c 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -52,6 +52,7 @@ namespace bolt {
 class BinaryBasicBlock;
 class BinaryContext;
 class BinaryFunction;
+class DataflowInfoManager;
 
 /// Different types of indirect branches encountered during disassembly.
 enum class IndirectBranchType : char {
@@ -475,7 +476,8 @@ class MCPlusBuilder {
   }
 
   /// Check whether this conditional branch can be reversed
-  virtual bool isReversibleBranch(const MCInst &Inst) const {
+  virtual bool isReversibleBranch(const MCInst &Inst,
+                                  DataflowInfoManager *DIM = nullptr) const {
     assert(!isUnsupportedInstruction(Inst) && isConditionalBranch(Inst) &&
            "Instruction is not known conditional branch");
 
@@ -2141,8 +2143,11 @@ class MCPlusBuilder {
   }
 
   /// Reverses the branch condition in Inst and update its taken target to TBB.
-  virtual void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                                      MCContext *Ctx) const {
+  /// Assumes that the branch is reversible.
+  virtual void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const {
     llvm_unreachable("not implemented");
   }
 
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index 4633d30104d43..ccdebf9e1ed0a 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -14,6 +14,8 @@
 namespace llvm {
 namespace bolt {
 
+class DataflowInfoManager;
+
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
 /// pull this pass inside BOLT because here we can do a better job at stub
@@ -74,7 +76,7 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(BinaryFunction &BF);
+  void relaxLocalBranches(BinaryFunction &BF, DataflowInfoManager *DIM);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/include/bolt/Utils/CommandLineOpts.h b/bolt/include/bolt/Utils/CommandLineOpts.h
index e11b18d3489cf..88b56217f9512 100644
--- a/bolt/include/bolt/Utils/CommandLineOpts.h
+++ b/bolt/include/bolt/Utils/CommandLineOpts.h
@@ -132,6 +132,10 @@ extern llvm::cl::opt<bool> UpdateDebugSections;
 // dbgs() for output within DEBUG().
 extern llvm::cl::opt<unsigned> Verbosity;
 
+// Option to control whether liveness analysis should be used by
+// FixupBranches and LongJmpPass. Needed for branch inversion on AArch64.
+extern llvm::cl::opt<bool> LivenessAnalysis;
+
 /// Return true if we should process all functions in the binary.
 bool processAllFunctions();
 
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 200e286d8e80e..ddfd163be84ec 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3650,7 +3650,7 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches() {
+void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3709,7 +3709,7 @@ void BinaryFunction::fixBranches() {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        if (!MIB->isReversibleBranch(*CondBranch)) {
+        if (!MIB->isReversibleBranch(*CondBranch, DIM)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3719,7 +3719,8 @@ void BinaryFunction::fixBranches() {
         std::swap(TSuccessor, FSuccessor);
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
-        MIB->reverseBranchCondition(*CondBranch, TSuccessor->getLabel(), Ctx);
+        MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
+                                    Ctx, DIM);
         return true;
       };
 
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index adf2bbae52d11..55b88233dcbde 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -11,8 +11,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/BinaryPasses.h"
+#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/FunctionLayout.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Passes/ReorderAlgorithm.h"
 #include "bolt/Passes/ReorderFunctions.h"
 #include "bolt/Utils/CommandLineOpts.h"
@@ -545,12 +547,22 @@ bool ReorderBasicBlocks::modifyFunctionLayout(BinaryFunction &BF,
 }
 
 Error FixupBranches::runOnFunctions(BinaryContext &BC) {
+  std::unique_ptr<BinaryFunctionCallGraph> CG;
+  std::unique_ptr<RegAnalysis> RA;
+  std::unique_ptr<DataflowInfoManager> DIM;
+
+  if (opts::LivenessAnalysis) {
+    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
+    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
+  }
   for (auto &It : BC.getBinaryFunctions()) {
     BinaryFunction &Function = It.second;
     if (!BC.shouldEmit(Function) || !Function.isSimple())
       continue;
 
-    Function.fixBranches();
+    if (opts::LivenessAnalysis)
+      DIM = std::make_unique<DataflowInfoManager>(Function, RA.get(), nullptr);
+    Function.fixBranches(DIM.get());
   }
   return Error::success();
 }
@@ -961,7 +973,7 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryFunction &BF) {
       uint64_t Count = 0;
       if (CondSucc != BB) {
         // Patch the new target address into the conditional branch.
-        MIB->reverseBranchCondition(*CondBranch, CalleeSymbol, Ctx);
+        MIB->reverseBranchCondition(PredBB, *CondBranch, CalleeSymbol, Ctx);
         // Since we reversed the condition on the branch we need to change
         // the target for the unconditional branch or add a unconditional
         // branch to the old target.  This has to be done manually since
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index b771e6a8b120a..03d5aa4bf3e2c 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -11,7 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/LongJmp.h"
+#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/Support/MathExtras.h"
 
@@ -662,7 +664,8 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
   return Error::success();
 }
 
-void LongJmpPass::relaxLocalBranches(BinaryFunction &BF) {
+void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
+                                     DataflowInfoManager *DIM) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -832,7 +835,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF) {
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool IsReversibleBranch = MIB->isReversibleBranch(Inst);
+      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, DIM);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -850,7 +853,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF) {
       if (ShouldReverseBranch && IsReversibleBranch) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
-        MIB->reverseBranchCondition(Inst, NextBB->getLabel(), BC.Ctx.get());
+        MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(), BC.Ctx.get(),
+                                    DIM);
       } else {
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
@@ -935,12 +939,23 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
+  std::unique_ptr<BinaryFunctionCallGraph> CG;
+  std::unique_ptr<RegAnalysis> RA;
+  std::unique_ptr<DataflowInfoManager> DIM;
+
+  if (opts::LivenessAnalysis) {
+    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
+    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
+  }
+
   if (opts::CompactCodeModel) {
     BC.outs()
         << "BOLT-INFO: relaxing branches for compact code model (<128MB)\n";
 
     ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF);
+      if (opts::LivenessAnalysis)
+        DIM = std::make_unique<DataflowInfoManager>(BF, RA.get(), nullptr);
+      relaxLocalBranches(BF, DIM.get());
     };
 
     ParallelUtilities::PredicateTy SkipPredicate =
@@ -965,12 +980,14 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     tentativeLayout(BC, Sorted);
     updateStubGroups();
     for (BinaryFunction *Func : Sorted) {
+      if (opts::LivenessAnalysis)
+        DIM = std::make_unique<DataflowInfoManager>(*Func, RA.get(), nullptr);
       if (auto E = relax(*Func, Modified))
         return Error(std::move(E));
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches();
+        Func->fixBranches(DIM.get());
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index b4a9dff9d25b2..ecb7bac9d7796 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -22,6 +22,7 @@
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/MCInstUtils.h"
 #include "bolt/Core/MCPlusBuilder.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDisassembler/MCDisassembler.h"
@@ -2045,6 +2046,25 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     exit(1);
   }
 
+  unsigned getInvertedCC(unsigned Opcode) const {
+    // clang-format off
+    switch (Opcode) {
+    default:
+      llvm_unreachable("Failed to invert condition code");
+      return Opcode;
+    // Compare register with immediate and branch.
+    case AArch64::CBGTWri:  return AArch64CC::LE;
+    case AArch64::CBGTXri:  return AArch64CC::LE;
+    case AArch64::CBLTWri:  return AArch64CC::GE;
+    case AArch64::CBLTXri:  return AArch64CC::GE;
+    case AArch64::CBHIWri:  return AArch64CC::LS;
+    case AArch64::CBHIXri:  return AArch64CC::LS;
+    case AArch64::CBLOWri:  return AArch64CC::HS;
+    case AArch64::CBLOXri:  return AArch64CC::HS;
+    }
+    // clang-format on
+  }
+
   unsigned getInvertedBranchOpcode(unsigned Opcode) const {
     // clang-format off
     switch (Opcode) {
@@ -2171,38 +2191,78 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  bool isReversibleBranch(const MCInst &Inst) const override {
+  bool isReversibleBranch(const MCInst &Inst,
+                          DataflowInfoManager *DIM = nullptr) const override {
     if (isCompAndBranch(Inst)) {
+      bool MayClobberFlags =
+          DIM ? DIM->getLivenessAnalysis().getLiveIn(Inst).test(getFlagsReg())
+              : true;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
-      if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0)
+      if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0 &&
+          MayClobberFlags)
         return false;
-      if (needsImmInc(InvertedOpcode) && Inst.getOperand(1).getImm() == 63)
+      if (needsImmInc(InvertedOpcode) && Inst.getOperand(1).getImm() == 63 &&
+          MayClobberFlags)
         return false;
     }
     return MCPlusBuilder::isReversibleBranch(Inst);
   }
 
-  void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                              MCContext *Ctx) const override {
-    if (!isReversibleBranch(Inst)) {
-      errs() << "BOLT-ERROR: Cannot reverse branch " << Inst << "\n";
-      exit(1);
-    }
+  void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const override {
+    assert(isReversibleBranch(Inst, DIM) && "Irreversible branch");
 
     if (isTB(Inst) || isCB(Inst) || isCompAndBranch(Inst)) {
+      bool ImmediateOutOfBounds = false;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
-      Inst.setOpcode(InvertedOpcode);
-      assert(Inst.getOpcode() != 0 && "Invalid branch instruction");
+      assert(InvertedOpcode != 0 && "Invalid branch instruction");
       // The FEAT_CMPBR compare-and-branch instructions cannot encode all
       // the possible condition codes, therefore we either have to adjust
       // the immediate value by +-1, or to swap the register operands
       // when reversing the branch condition.
       if (needsRegSwap(InvertedOpcode))
         std::swap(Inst.getOperand(0), Inst.getOperand(1));
-      else if (needsImmDec(InvertedOpcode))
-        Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() - 1);
-      else if (needsImmInc(InvertedOpcode))
-        Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() + 1);
+      else if (needsImmDec(InvertedOpcode)) {
+        if (Inst.getOperand(1).getImm() == 0)
+          ImmediateOutOfBounds = true;
+        else
+          Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() - 1);
+      } else if (needsImmInc(InvertedOpcode)) {
+        if (Inst.getOperand(1).getImm() == 63)
+          ImmediateOutOfBounds = true;
+        else
+          Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() + 1);
+      }
+      if (ImmediateOutOfBounds) {
+        auto is32BitVariant = [](unsigned Opcode) {
+          switch (Opcode) {
+          default:
+            return false;
+          case AArch64::CBGTWri:
+          case AArch64::CBLTWri:
+          case AArch64::CBHIWri:
+          case AArch64::CBLOWri:
+            return true;
+          }
+        };
+        InstructionListType Code;
+        MCInstBuilder Cmp =
+            is32BitVariant(InvertedOpcode)
+                ? MCInstBuilder(AArch64::SUBSWri).addReg(AArch64::WZR)
+                : MCInstBuilder(AArch64::SUBSXri).addReg(AArch64::XZR);
+        Cmp.addReg(Inst.getOperand(0).getReg())
+            .addImm(Inst.getOperand(1).getImm())
+            .addImm(0);
+        Code.emplace_back(std::move(Cmp));
+        Code.emplace_back(MCInstBuilder(AArch64::Bcc)
+                              .addImm(getInvertedCC(Inst.getOpcode()))
+                              .addExpr(MCSymbolRefExpr::create(TBB, *Ctx)));
+        Parent->replaceInstruction(Parent->findInstruction(&Inst), Code);
+        return;
+      }
+      Inst.setOpcode(InvertedOpcode);
     } else if (Inst.getOpcode() == AArch64::Bcc) {
       Inst.getOperand(0).setImm(AArch64CC::getInvertedCondCode(
           static_cast<AArch64CC::CondCode>(Inst.getOperand(0).getImm())));
diff --git a/bolt/lib/Target/AArch64/CMakeLists.txt b/bolt/lib/Target/AArch64/CMakeLists.txt
index 1e171748aece6..e28ed0bd66ba1 100644
--- a/bolt/lib/Target/AArch64/CMakeLists.txt
+++ b/bolt/lib/Target/AArch64/CMakeLists.txt
@@ -29,7 +29,11 @@ add_llvm_library(LLVMBOLTTargetAArch64
   AArch64CommonTableGen
   )
 
-target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE LLVMBOLTCore LLVMBOLTUtils)
+target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE
+  LLVMBOLTCore
+  LLVMBOLTPasses
+  LLVMBOLTUtils
+  )
 
 include_directories(
   ${LLVM_MAIN_SRC_DIR}/lib/Target/AArch64
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index d1a0572277874..b954a392c4f26 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -162,8 +162,10 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                              MCContext *Ctx) const override {
+  void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const override {
     auto Opcode = getInvertedBranchOpcode(Inst.getOpcode());
     Inst.setOpcode(Opcode);
     replaceBranchTarget(Inst, TBB, Ctx);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 11a297f514530..88dcf994e5dc7 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2811,8 +2811,10 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(CC));
   }
 
-  void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                              MCContext *Ctx) const override {
+  void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const override {
     unsigned InvCC = getInvertedCondCode(getCondCode(Inst));
     assert(InvCC != X86::COND_INVALID && "invalid branch instruction");
     Inst.getOperand(Info->get(Inst.getOpcode()).NumOperands - 1).setImm(InvCC);
diff --git a/bolt/lib/Utils/CommandLineOpts.cpp b/bolt/lib/Utils/CommandLineOpts.cpp
index 20b24c3b4acc5..82fde5788397a 100644
--- a/bolt/lib/Utils/CommandLineOpts.cpp
+++ b/bolt/lib/Utils/CommandLineOpts.cpp
@@ -365,6 +365,12 @@ cl::opt<unsigned>
               cl::init(0), cl::ZeroOrMore, cl::cat(BoltCategory),
               cl::sub(cl::SubCommand::getAll()));
 
+cl::opt<bool> LivenessAnalysis(
+    "liveness-analysis",
+    cl::desc("use liveness analysis in FixupBranches and LongJmpPass"
+             "(needed for branch inversion on AArch64)"),
+    cl::init(false), cl::cat(BoltCategory));
+
 bool processAllFunctions() {
   if (opts::AggregateOnly)
     return false;
diff --git a/bolt/test/AArch64/compare-and-branch-inversion.S b/bolt/test/AArch64/compare-and-branch-inversion.S
index 28167416c31cb..0ea084e904101 100644
--- a/bolt/test/AArch64/compare-and-branch-inversion.S
+++ b/bolt/test/AArch64/compare-and-branch-inversion.S
@@ -1,18 +1,24 @@
 # This test checks that branch inversion works when reordering blocks which
 # contain short range conditional branches. Handles edge cases, like when
 # the immediate value is the upper or lower allowed value in which case the
-# transformation bails.
+# transformation bails. If liveness analysis proves that the condition flags
+# are dead we can replace the branch with cmp + b.cc
 
 # REQUIRES: system-linux, asserts
 
 # RUN: %clang %cflags -march=armv9-a+cmpbr -Wl,-q %s -o %t
 # RUN: link_fdata --no-lbr %s %t %t.fdata
 # RUN: llvm-strip --strip-unneeded %t
+#
 # RUN: llvm-bolt -v=1 %t -o %t.bolt --data %t.fdata --reorder-blocks=ext-tsp --compact-code-model \
-# RUN:   | FileCheck %s --check-prefix=BOLT-INFO
-# RUN: llvm-objdump -d %t.bolt | FileCheck %s
+# RUN:   | FileCheck %s --check-prefix=BOLT-INFO-NO-LIVENESS
+# RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=NO-LIVENESS
+#
+# RUN: llvm-bolt -v=1 %t -o %t.bolt --data %t.fdata --reorder-blocks=ext-tsp --compact-code-model \
+# RUN:   --liveness-analysis | FileCheck %s --check-prefix=BOLT-INFO-LIVENESS
+# RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=LIVENESS
 
-# CHECK: Disassembly of section .text:
+# COMMON: Disassembly of section .text:
 
   .globl immediate_increment
   .type immediate_increment, %function
@@ -29,12 +35,12 @@ immediate_increment:
     mov x0, #2
     ret
 
-# CHECK: <immediate_increment>:
-# CHECK-NEXT:            {{.*}} cblt x0, #0x1, 0x[[ADDR0:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR0]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <immediate_increment>:
+# COMMON-NEXT:            {{.*}} cblt x0, #0x1, 0x[[ADDR0:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR0]]: {{.*}} mov x0, #0x1 // =1
+# COMMON-NEXT:            {{.*}} ret
 
   .globl immediate_decrement
   .type immediate_decrement, %function
@@ -51,12 +57,12 @@ immediate_decrement:
     mov x0, #2
     ret
 
-# CHECK: <immediate_decrement>:
-# CHECK-NEXT:            {{.*}} cbhi x0, #0x0, 0x[[ADDR1:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR1]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <immediate_decrement>:
+# COMMON-NEXT:            {{.*}} cbhi x0, #0x0, 0x[[ADDR1:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR1]]: {{.*}} mov x0, #0x1 // =1
+# COMMON-NEXT:            {{.*}} ret
 
   .globl register_swap
   .type register_swap, %function
@@ -73,37 +79,77 @@ register_swap:
     mov x0, #2
     ret
 
-# CHECK: <register_swap>:
-# CHECK-NEXT:            {{.*}} cbgt x1, x0, 0x[[ADDR2:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR2]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <register_swap>:
+# COMMON-NEXT:            {{.*}} cbgt x1, x0, 0x[[ADDR2:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR2]]: {{.*}} mov x0, #0x1 // =1
+# COMMON-NEXT:            {{.*}} ret
 
-  .globl irreversible
-  .type irreversible, %function
-irreversible:
+  .globl immediate_overflow
+  .type immediate_overflow, %function
+immediate_overflow:
 .entry3:
-# FDATA: 1 irreversible #.entry3# 10
+# FDATA: 1 immediate_overflow #.entry3# 10
     cbgt x0, #63, .exit3
 .cold3:
-# FDATA: 1 irreversible #.cold3# 1
+# FDATA: 1 immediate_overflow #.cold3# 1
     mov x0, #1
     ret
 .exit3:
-# FDATA: 1 irreversible #.exit3# 10
+# FDATA: 1 immediate_overflow #.exit3# 10
+    mov x0, #2
+    ret
+
+# BOLT-INFO-NO-LIVENESS: unable to swap successors in immediate_overflow
+#
+# Without liveness the blocks get reordered, but since the branch is
+# irreversible an additional unconditional branch is emitted.
+# This codegen is suboptimal yet correct.
+#
+# NO-LIVENESS: <immediate_overflow>:
+# NO-LIVENESS-NEXT:            {{.*}} cbgt x0, #0x3f, 0x[[ADDR3:[0-9a-f]+]] <{{.*}}>
+# NO-LIVENESS-NEXT:            {{.*}} b               0x[[ADDR4:[0-9a-f]+]] <{{.*}}>
+# NO-LIVENESS-NEXT: [[ADDR3]]: {{.*}} mov  x0, #0x2 // =2
+# NO-LIVENESS-NEXT:            {{.*}} ret
+# NO-LIVENESS-NEXT: [[ADDR4]]: {{.*}} mov x0, #0x1 // =1
+# NO-LIVENESS-NEXT:            {{.*}} ret
+
+# LIVENESS: <immediate_overflow>:
+# LIVENESS-NEXT:            {{.*}} cmp  x0, #0x3f
+# LIVENESS-NEXT:            {{.*}} b.le 0x[[ADDR5:[0-9a-f]+]] <{{.*}}>
+# LIVENESS-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# LIVENESS-NEXT:            {{.*}} ret
+# LIVENESS-NEXT: [[ADDR5]]: {{.*}} mov x0, #0x1 // =1
+# LIVENESS-NEXT:            {{.*}} ret
+
+  .globl irreversible
+  .type irreversible, %function
+irreversible:
+.entry4:
+# FDATA: 1 irreversible #.entry4# 10
+    cmp x0, #63
+    cbgt x0, #63, .exit4
+.cold4:
+# FDATA: 1 irreversible #.cold4# 1
+    csel x0, x1, x2, le
+    ret
+.exit4:
+# FDATA: 1 irreversible #.exit4# 10
     mov x0, #2
     ret
 
-# BOLT-INFO: unable to swap successors in irreversible
+# BOLT-INFO-NO-LIVENESS: unable to swap successors in irreversible
+# BOLT-INFO-LIVENESS:    unable to swap successors in irreversible
 
-# CHECK: <irreversible>:
-# CHECK-NEXT:            {{.*}} cbgt x0, #0x3f, 0x[[ADDR3:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} b               0x[[ADDR4:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT: [[ADDR3]]: {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR4]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <irreversible>:
+# COMMON-NEXT:            {{.*}} cmp  x0, #0x3f
+# COMMON-NEXT:            {{.*}} cbgt x0, #0x3f, 0x[[ADDR6:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} b               0x[[ADDR7:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT: [[ADDR6]]: {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR7]]: {{.*}} csel x0, x1, x2, le
+# COMMON-NEXT:            {{.*}} ret
 
 ## Force relocation mode.
 .reloc 0, R_AARCH64_NONE
diff --git a/bolt/unittests/Core/CMakeLists.txt b/bolt/unittests/Core/CMakeLists.txt
index 297dec7449202..b755afc8a66da 100644
--- a/bolt/unittests/Core/CMakeLists.txt
+++ b/bolt/unittests/Core/CMakeLists.txt
@@ -24,6 +24,7 @@ target_link_libraries(CoreTests
   PRIVATE
   LLVMBOLTCore
   LLVMBOLTRewrite
+  LLVMBOLTPasses
   LLVMBOLTProfile
   LLVMBOLTUtils
   )
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index e67460fe2a6a6..5ea242d65e533 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -18,6 +18,9 @@
 
 #include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/BinaryFunctionCallGraph.h"
+#include "bolt/Passes/BinaryPasses.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
@@ -233,8 +236,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                            .addExpr(MCSymbolRefExpr::create(
                                TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(NeedsImmInc));
-  BC->MIB->reverseBranchCondition(NeedsImmInc, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, NeedsImmInc,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(NeedsImmInc.getOpcode(), AArch64::CBLTXri);
   ASSERT_EQ(NeedsImmInc.getOperand(1).getImm(), 1);
 
@@ -247,8 +250,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                            .addExpr(MCSymbolRefExpr::create(
                                TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(NeedsImmDec));
-  BC->MIB->reverseBranchCondition(NeedsImmDec, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, NeedsImmDec,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(NeedsImmDec.getOpcode(), AArch64::CBHIXri);
   ASSERT_EQ(NeedsImmDec.getOperand(1).getImm(), 0);
 
@@ -261,8 +264,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                    .addExpr(MCSymbolRefExpr::create(
                                        TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompRegNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(CompRegNeedsRegSwap, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompRegNeedsRegSwap,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(CompRegNeedsRegSwap.getOpcode(), AArch64::CBGTXrr);
   ASSERT_EQ(CompRegNeedsRegSwap.getOperand(0).getReg(), AArch64::X1);
   ASSERT_EQ(CompRegNeedsRegSwap.getOperand(1).getReg(), AArch64::X0);
@@ -276,8 +279,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                     .addExpr(MCSymbolRefExpr::create(
                                         TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompByteNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(CompByteNeedsRegSwap, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompByteNeedsRegSwap,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(CompByteNeedsRegSwap.getOpcode(), AArch64::CBBHSWrr);
   ASSERT_EQ(CompByteNeedsRegSwap.getOperand(0).getReg(), AArch64::W1);
   ASSERT_EQ(CompByteNeedsRegSwap.getOperand(1).getReg(), AArch64::W0);
@@ -291,8 +294,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                     .addExpr(MCSymbolRefExpr::create(
                                         TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompHalfNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(CompHalfNeedsRegSwap, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompHalfNeedsRegSwap,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(CompHalfNeedsRegSwap.getOpcode(), AArch64::CBHHIWrr);
   ASSERT_EQ(CompHalfNeedsRegSwap.getOperand(0).getReg(), AArch64::W1);
   ASSERT_EQ(CompHalfNeedsRegSwap.getOperand(1).getReg(), AArch64::W0);
@@ -318,6 +321,123 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
   ASSERT_FALSE(BC->MIB->isReversibleBranch(Overflows));
 }
 
+TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Underflows) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  BinaryBasicBlock *EntryBB = BF->addBasicBlock();
+  BinaryBasicBlock *FallThroughBB = BF->addBasicBlock();
+  BinaryBasicBlock *TargetBB = BF->addBasicBlock();
+  BF->addEntryPoint(*EntryBB);
+  EntryBB->addSuccessor(TargetBB);
+  EntryBB->addSuccessor(FallThroughBB);
+
+  // Inversion requires expansion, immediate value underflows.
+  // cblt x0, #0, target ~> cmp x0, #0
+  //                        b.ge target
+  auto I =
+      EntryBB->addInstruction(MCInstBuilder(AArch64::CBLTXri)
+                                  .addReg(AArch64::X0)
+                                  .addImm(0)
+                                  .addExpr(MCSymbolRefExpr::create(
+                                      TargetBB->getLabel(), *BC->Ctx.get())));
+  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
+  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
+  DataflowInfoManager DIM(*BF, &RA, nullptr);
+
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
+                                  BC->Ctx.get(), &DIM);
+  I = EntryBB->begin();
+  ASSERT_EQ(I->getOpcode(), AArch64::SUBSXri);
+  ASSERT_EQ(I->getOperand(0).getReg(), AArch64::XZR);
+  ASSERT_EQ(I->getOperand(1).getReg(), AArch64::X0);
+  ASSERT_EQ(I->getOperand(2).getImm(), 0);
+  ASSERT_EQ(I->getOperand(3).getImm(), 0);
+  I++;
+  ASSERT_EQ(I->getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(I->getOperand(0).getImm(), AArch64CC::GE);
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  BinaryBasicBlock *EntryBB = BF->addBasicBlock();
+  BinaryBasicBlock *FallThroughBB = BF->addBasicBlock();
+  BinaryBasicBlock *TargetBB = BF->addBasicBlock();
+  BF->addEntryPoint(*EntryBB);
+  EntryBB->addSuccessor(TargetBB);
+  EntryBB->addSuccessor(FallThroughBB);
+
+  // Inversion requires expansion, immediate value overflows.
+  // cbhi w0, #63, target ~> cmp w0, #63
+  //                         b.ls target
+  auto I =
+      EntryBB->addInstruction(MCInstBuilder(AArch64::CBHIWri)
+                                  .addReg(AArch64::W0)
+                                  .addImm(63)
+                                  .addExpr(MCSymbolRefExpr::create(
+                                      TargetBB->getLabel(), *BC->Ctx.get())));
+  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
+  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
+  DataflowInfoManager DIM(*BF, &RA, nullptr);
+
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
+                                  BC->Ctx.get(), &DIM);
+  I = EntryBB->begin();
+  ASSERT_EQ(I->getOpcode(), AArch64::SUBSWri);
+  ASSERT_EQ(I->getOperand(0).getReg(), AArch64::WZR);
+  ASSERT_EQ(I->getOperand(1).getReg(), AArch64::W0);
+  ASSERT_EQ(I->getOperand(2).getImm(), 63);
+  ASSERT_EQ(I->getOperand(3).getImm(), 0);
+  I++;
+  ASSERT_EQ(I->getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(I->getOperand(0).getImm(), AArch64CC::LS);
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_IsReversibleBranch_LiveCondFlags) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  BinaryBasicBlock *EntryBB = BF->addBasicBlock();
+  BinaryBasicBlock *FallThroughBB = BF->addBasicBlock();
+  BinaryBasicBlock *TargetBB = BF->addBasicBlock();
+  BF->addEntryPoint(*EntryBB);
+  EntryBB->addSuccessor(TargetBB);
+  EntryBB->addSuccessor(FallThroughBB);
+
+  // cmp x0, #63
+  EntryBB->addInstruction(MCInstBuilder(AArch64::SUBSXri)
+                              .addReg(AArch64::XZR)
+                              .addReg(AArch64::X0)
+                              .addImm(63)
+                              .addImm(0));
+  // cbgt x0, #63, target
+  auto I =
+      EntryBB->addInstruction(MCInstBuilder(AArch64::CBGTXri)
+                                  .addReg(AArch64::X0)
+                                  .addImm(63)
+                                  .addExpr(MCSymbolRefExpr::create(
+                                      TargetBB->getLabel(), *BC->Ctx.get())));
+  // csel x0, x1, x2, le
+  FallThroughBB->addInstruction(MCInstBuilder(AArch64::CSELXr)
+                                    .addReg(AArch64::X0)
+                                    .addReg(AArch64::X1)
+                                    .addReg(AArch64::X2)
+                                    .addImm(13));
+
+  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
+  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
+  DataflowInfoManager DIM(*BF, &RA, nullptr);
+
+  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &DIM));
+}
+
 TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {
   if (GetParam() != Triple::aarch64)
     GTEST_SKIP();

>From 75d8603c42ec4339a70bb8d0ec5205ab61f3efd4 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Thu, 2 Jul 2026 12:33:22 +0100
Subject: [PATCH 2/6] Run Liveness Analysis upfront and cache it before any CFG
 changes.

Teach the branch-fixing paths to use cached branch liveness information
when deciding whether this expansion is legal. The liveness snapshot is
built before branch relaxation/fixup and is safe for the current users:
they only insert trampolines/stubs between existing CFG edges or invert
branches without changing program semantics.

Since expansion can grow the source basic block, update the local
address accounting used by branch relaxation.
---
 bolt/include/bolt/Core/BinaryFunction.h       |   4 +-
 bolt/include/bolt/Core/MCPlusBuilder.h        |  23 +++-
 bolt/include/bolt/Passes/LongJmp.h            |   4 +-
 bolt/lib/Core/BinaryFunction.cpp              |   6 +-
 bolt/lib/Passes/BinaryPasses.cpp              |  32 ++---
 bolt/lib/Passes/LongJmp.cpp                   | 109 +++++++++++-------
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   |  34 ++++--
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  |   7 +-
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |   7 +-
 bolt/unittests/Core/MCPlusBuilder.cpp         |  16 ++-
 10 files changed, 148 insertions(+), 94 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 51651e41497da..40e7da06666fc 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -65,7 +65,7 @@ class DWARFUnit;
 
 namespace bolt {
 
-class DataflowInfoManager;
+struct BranchLivenessInfo;
 
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
@@ -2465,7 +2465,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches(DataflowInfoManager *DIM = nullptr);
+  void fixBranches(const BranchLivenessInfo *BLI = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index c98b56d14523c..c2859561910d5 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -18,6 +18,7 @@
 #include "bolt/Core/Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/MC/MCAsmBackend.h"
@@ -74,6 +75,15 @@ enum BTIKind {
   JC /// Accepting both.
 };
 
+struct BranchLivenessInfo {
+  DenseMap<const MCInst *, bool> FlagsLiveIn;
+
+  bool mustPreserveFlags(const MCInst &Inst) const {
+    auto It = FlagsLiveIn.find(&Inst);
+    return It == FlagsLiveIn.end() || It->second;
+  }
+};
+
 class MCPlusBuilder {
 public:
   using AllocatorIdTy = uint16_t;
@@ -475,9 +485,16 @@ class MCPlusBuilder {
     return false;
   }
 
+  /// Return liveness info required for branch transformations.
+  virtual BranchLivenessInfo
+  createBranchLivenessInfo(BinaryFunction &BF, DataflowInfoManager &DIM) const {
+    return BranchLivenessInfo();
+  }
+
   /// Check whether this conditional branch can be reversed
-  virtual bool isReversibleBranch(const MCInst &Inst,
-                                  DataflowInfoManager *DIM = nullptr) const {
+  virtual bool
+  isReversibleBranch(const MCInst &Inst,
+                     const BranchLivenessInfo *BLI = nullptr) const {
     assert(!isUnsupportedInstruction(Inst) && isConditionalBranch(Inst) &&
            "Instruction is not known conditional branch");
 
@@ -2147,7 +2164,7 @@ class MCPlusBuilder {
   virtual void
   reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
                          const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const {
+                         const BranchLivenessInfo *BLI = nullptr) const {
     llvm_unreachable("not implemented");
   }
 
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index ccdebf9e1ed0a..90e10943c0f70 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -14,7 +14,7 @@
 namespace llvm {
 namespace bolt {
 
-class DataflowInfoManager;
+struct BranchLivenessInfo;
 
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
@@ -76,7 +76,7 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(BinaryFunction &BF, DataflowInfoManager *DIM);
+  void relaxLocalBranches(BinaryFunction &BF, const BranchLivenessInfo *BLI);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index ddfd163be84ec..4fa0290af8b3b 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3650,7 +3650,7 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
+void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3709,7 +3709,7 @@ void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        if (!MIB->isReversibleBranch(*CondBranch, DIM)) {
+        if (!MIB->isReversibleBranch(*CondBranch, BLI)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3720,7 +3720,7 @@ void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
         MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
-                                    Ctx, DIM);
+                                    Ctx, BLI);
         return true;
       };
 
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 55b88233dcbde..95ebbd8258159 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -547,23 +547,25 @@ bool ReorderBasicBlocks::modifyFunctionLayout(BinaryFunction &BF,
 }
 
 Error FixupBranches::runOnFunctions(BinaryContext &BC) {
-  std::unique_ptr<BinaryFunctionCallGraph> CG;
-  std::unique_ptr<RegAnalysis> RA;
-  std::unique_ptr<DataflowInfoManager> DIM;
+  auto forEachFunction = [&](auto &&Apply) {
+    for (auto &It : BC.getBinaryFunctions()) {
+      BinaryFunction &Function = It.second;
+      if (!BC.shouldEmit(Function) || !Function.isSimple())
+        continue;
+      Apply(Function);
+    }
+  };
 
   if (opts::LivenessAnalysis) {
-    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
-    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
-  }
-  for (auto &It : BC.getBinaryFunctions()) {
-    BinaryFunction &Function = It.second;
-    if (!BC.shouldEmit(Function) || !Function.isSimple())
-      continue;
-
-    if (opts::LivenessAnalysis)
-      DIM = std::make_unique<DataflowInfoManager>(Function, RA.get(), nullptr);
-    Function.fixBranches(DIM.get());
-  }
+    BinaryFunctionCallGraph CG = buildCallGraph(BC);
+    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
+    forEachFunction([&](BinaryFunction &BF) {
+      DataflowInfoManager DIM(BF, &RA, nullptr);
+      BranchLivenessInfo Info = BC.MIB->createBranchLivenessInfo(BF, DIM);
+      BF.fixBranches(&Info);
+    });
+  } else
+    forEachFunction([&](BinaryFunction &BF) { BF.fixBranches(nullptr); });
   return Error::success();
 }
 
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 03d5aa4bf3e2c..871791e6b7988 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -665,7 +665,7 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
 }
 
 void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
-                                     DataflowInfoManager *DIM) {
+                                     const BranchLivenessInfo *BLI) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -711,14 +711,18 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
     DenseMap<const BinaryBasicBlock *, BinaryBasicBlock *> FragmentTrampolines;
 
     // Create a trampoline code after \p BB or at the end of the fragment if BB
-    // is nullptr. If \p UpdateOffsets is true, update FragmentSize and offsets
-    // for basic blocks affected by the insertion of the trampoline.
+    // is nullptr. \p Offset is the fragment size delta caused by the insertion,
+    // including any growth of \p BB before the trampoline.
     auto addTrampolineAfter = [&](BinaryBasicBlock *BB,
                                   BinaryBasicBlock *TargetBB, uint64_t Count,
-                                  bool UpdateOffsets = true) {
+                                  uint64_t Offset) {
       FunctionTrampolines.emplace_back(BB ? BB : FF.back(),
                                        BF.createBasicBlock());
       BinaryBasicBlock *TrampolineBB = FunctionTrampolines.back().second.get();
+      const uint64_t BBGrowth = Offset ? Offset - TrampolineSize : 0;
+
+      if (BB)
+        BB->setOutputEndAddress(BB->getOutputEndAddress() + BBGrowth);
 
       MCInst Inst;
       {
@@ -734,13 +738,23 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       TrampolineBB->setOutputEndAddress(TrampolineAddress + TrampolineSize);
       TrampolineBB->setFragmentNum(FF.getFragmentNum());
 
+      // Shift the fragment-local output address range for blocks at or after
+      // the old end address.
+      auto adjustBasicBlockAddress = [](BinaryBasicBlock *BB, uint64_t Address,
+                                        uint64_t Offset) {
+        if (BB->getOutputStartAddress() < Address)
+          return;
+        BB->setOutputStartAddress(BB->getOutputStartAddress() + Offset);
+        BB->setOutputEndAddress(BB->getOutputEndAddress() + Offset);
+      };
+
       if (!FragmentTrampolines.lookup(TargetBB))
         FragmentTrampolines[TargetBB] = TrampolineBB;
 
-      if (!UpdateOffsets)
+      if (!Offset)
         return TrampolineBB;
 
-      FragmentSize += TrampolineSize;
+      FragmentSize += Offset;
 
       // If the trampoline was added at the end of the fragment, offsets of
       // other fragments should stay intact.
@@ -748,13 +762,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         return TrampolineBB;
 
       // Update offsets for blocks after BB.
-      for (BinaryBasicBlock *IBB : FF) {
-        if (IBB->getOutputStartAddress() >= TrampolineAddress) {
-          IBB->setOutputStartAddress(IBB->getOutputStartAddress() +
-                                     TrampolineSize);
-          IBB->setOutputEndAddress(IBB->getOutputEndAddress() + TrampolineSize);
-        }
-      }
+      for (BinaryBasicBlock *IBB : FF)
+        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
 
       // Update offsets for trampolines in this fragment that are placed after
       // the new trampoline. Note that trampoline blocks are not part of the
@@ -766,11 +775,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
           continue;
         if (IBB == TrampolineBB)
           continue;
-        if (IBB->getOutputStartAddress() >= TrampolineAddress) {
-          IBB->setOutputStartAddress(IBB->getOutputStartAddress() +
-                                     TrampolineSize);
-          IBB->setOutputEndAddress(IBB->getOutputEndAddress() + TrampolineSize);
-        }
+        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
       }
 
       return TrampolineBB;
@@ -791,7 +796,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
 
       BinaryBasicBlock *TrampolineBB =
-          addTrampolineAfter(BB, TargetBB, BI.Count, /*UpdateOffsets*/ false);
+          addTrampolineAfter(BB, TargetBB, BI.Count, /*Offset=*/0);
       BB->replaceSuccessor(TargetBB, TrampolineBB, BI.Count);
     }
 
@@ -824,7 +829,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // case we will need further relaxation.
       const int64_t OffsetToEnd = FragmentSize - InstAddress;
       if (Count == 0 && isBranchOffsetInRange(Inst, OffsetToEnd)) {
-        TrampolineBB = addTrampolineAfter(nullptr, TargetBB, Count);
+        TrampolineBB =
+            addTrampolineAfter(nullptr, TargetBB, Count, TrampolineSize);
         BB->replaceSuccessor(TargetBB, TrampolineBB, Count);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
@@ -835,7 +841,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, DIM);
+      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, BLI);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -843,19 +849,26 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       if (ShouldReverseBranch && !IsReversibleBranch) {
         const uint64_t NextCount = BB->getBranchInfo(*NextBB).Count;
         BinaryBasicBlock *FallThrough =
-            addTrampolineAfter(BB, NextBB, NextCount);
+            addTrampolineAfter(BB, NextBB, NextCount, TrampolineSize);
         BB->replaceSuccessor(NextBB, FallThrough, NextCount);
       }
 
-      // Create a trampoline basic block for the taken target of the branch.
-      TrampolineBB = addTrampolineAfter(BB, TargetBB, Count);
-
       if (ShouldReverseBranch && IsReversibleBranch) {
+        const uint64_t OldBBSize = BB->estimateSize();
         BB->swapConditionalSuccessors();
-        auto L = BC.scopeLock();
-        MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(), BC.Ctx.get(),
-                                    DIM);
+        {
+          auto L = BC.scopeLock();
+          MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(),
+                                      BC.Ctx.get(), BLI);
+        }
+        const uint64_t NewBBSize = BB->estimateSize();
+
+        // Create a trampoline basic block for the original taken target.
+        TrampolineBB = addTrampolineAfter(
+            BB, TargetBB, Count, TrampolineSize + (NewBBSize - OldBBSize));
       } else {
+        // Create a trampoline basic block for the taken target of the branch.
+        TrampolineBB = addTrampolineAfter(BB, TargetBB, Count, TrampolineSize);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
       }
@@ -870,7 +883,10 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       for (auto BBI = FF.begin(); BBI != FF.end(); ++BBI) {
         BinaryBasicBlock *BB = *BBI;
         uint64_t NextInstOffset = BB->getOutputStartAddress();
-        for (MCInst &Inst : *BB) {
+        // Branch reversal may replace the current instruction with a sequence.
+        // Use an index so the next instruction is reloaded after the mutation.
+        for (size_t I = 0; I < BB->size(); ++I) {
+          MCInst &Inst = *(BB->begin() + I);
           const size_t InstAddress = NextInstOffset;
           if (!MIB->isPseudo(Inst))
             NextInstOffset += 4;
@@ -939,30 +955,37 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
-  std::unique_ptr<BinaryFunctionCallGraph> CG;
-  std::unique_ptr<RegAnalysis> RA;
-  std::unique_ptr<DataflowInfoManager> DIM;
-
+  DenseMap<BinaryFunction *, BranchLivenessInfo> BranchLiveness;
   if (opts::LivenessAnalysis) {
-    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
-    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
+    BinaryFunctionCallGraph CG = buildCallGraph(BC);
+    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
+    for (auto &It : BC.getBinaryFunctions()) {
+      BinaryFunction &BF = It.second;
+      if (!BC.shouldEmit(BF) || !BF.isSimple())
+        continue;
+      DataflowInfoManager DIM(BF, &RA, nullptr);
+      BranchLiveness[&BF] = BC.MIB->createBranchLivenessInfo(BF, DIM);
+    }
   }
+  auto getBranchLiveness =
+      [&](BinaryFunction &BF) -> const BranchLivenessInfo * {
+    auto It = BranchLiveness.find(&BF);
+    return It == BranchLiveness.end() ? nullptr : &It->second;
+  };
 
   if (opts::CompactCodeModel) {
     BC.outs()
         << "BOLT-INFO: relaxing branches for compact code model (<128MB)\n";
 
-    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      if (opts::LivenessAnalysis)
-        DIM = std::make_unique<DataflowInfoManager>(BF, RA.get(), nullptr);
-      relaxLocalBranches(BF, DIM.get());
-    };
-
     ParallelUtilities::PredicateTy SkipPredicate =
         [&](const BinaryFunction &BF) {
           return !BC.shouldEmit(BF) || !BF.isSimple();
         };
 
+    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+      relaxLocalBranches(BF, getBranchLiveness(BF));
+    };
+
     ParallelUtilities::runOnEachFunction(
         BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
         SkipPredicate, "RelaxLocalBranches");
@@ -980,14 +1003,12 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     tentativeLayout(BC, Sorted);
     updateStubGroups();
     for (BinaryFunction *Func : Sorted) {
-      if (opts::LivenessAnalysis)
-        DIM = std::make_unique<DataflowInfoManager>(*Func, RA.get(), nullptr);
       if (auto E = relax(*Func, Modified))
         return Error(std::move(E));
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches(DIM.get());
+        Func->fixBranches(getBranchLiveness(*Func));
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index ecb7bac9d7796..8025f75b789e1 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2191,28 +2191,38 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  bool isReversibleBranch(const MCInst &Inst,
-                          DataflowInfoManager *DIM = nullptr) const override {
+  BranchLivenessInfo
+  createBranchLivenessInfo(BinaryFunction &BF,
+                           DataflowInfoManager &DIM) const override {
+    BranchLivenessInfo Info;
+    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+    for (BinaryBasicBlock &BB : BF)
+      for (MCInst &Inst : BB)
+        if (isCompAndBranch(Inst))
+          Info.FlagsLiveIn[&Inst] = LA.getLiveIn(Inst).test(getFlagsReg());
+    return Info;
+  }
+
+  bool
+  isReversibleBranch(const MCInst &Inst,
+                     const BranchLivenessInfo *BLI = nullptr) const override {
     if (isCompAndBranch(Inst)) {
-      bool MayClobberFlags =
-          DIM ? DIM->getLivenessAnalysis().getLiveIn(Inst).test(getFlagsReg())
-              : true;
+      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
       if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0 &&
-          MayClobberFlags)
+          MustPreserveFlags)
         return false;
       if (needsImmInc(InvertedOpcode) && Inst.getOperand(1).getImm() == 63 &&
-          MayClobberFlags)
+          MustPreserveFlags)
         return false;
     }
     return MCPlusBuilder::isReversibleBranch(Inst);
   }
 
-  void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const override {
-    assert(isReversibleBranch(Inst, DIM) && "Irreversible branch");
+  void reverseBranchCondition(
+      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
+      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
+    assert(isReversibleBranch(Inst, BLI) && "Irreversible branch");
 
     if (isTB(Inst) || isCB(Inst) || isCompAndBranch(Inst)) {
       bool ImmediateOutOfBounds = false;
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index b954a392c4f26..34e8392f59743 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -162,10 +162,9 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const override {
+  void reverseBranchCondition(
+      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
+      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
     auto Opcode = getInvertedBranchOpcode(Inst.getOpcode());
     Inst.setOpcode(Opcode);
     replaceBranchTarget(Inst, TBB, Ctx);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 88dcf994e5dc7..98b9eb8fdd4cb 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2811,10 +2811,9 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(CC));
   }
 
-  void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const override {
+  void reverseBranchCondition(
+      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
+      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
     unsigned InvCC = getInvertedCondCode(getCondCode(Inst));
     assert(InvCC != X86::COND_INVALID && "invalid branch instruction");
     Inst.getOperand(Info->get(Inst.getOpcode()).NumOperands - 1).setImm(InvCC);
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 5ea242d65e533..5fdaf47094770 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -345,10 +345,12 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Underflows) {
   BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
   RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
   DataflowInfoManager DIM(*BF, &RA, nullptr);
+  BranchLivenessInfo BranchLiveness =
+      BC->MIB->createBranchLivenessInfo(*BF, DIM);
 
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &DIM);
+                                  BC->Ctx.get(), &BranchLiveness);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSXri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::XZR);
@@ -384,10 +386,12 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
   BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
   RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
   DataflowInfoManager DIM(*BF, &RA, nullptr);
+  BranchLivenessInfo BranchLiveness =
+      BC->MIB->createBranchLivenessInfo(*BF, DIM);
 
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &DIM);
+                                  BC->Ctx.get(), &BranchLiveness);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSWri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::WZR);
@@ -434,8 +438,10 @@ TEST_P(MCPlusBuilderTester, AArch64_IsReversibleBranch_LiveCondFlags) {
   BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
   RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
   DataflowInfoManager DIM(*BF, &RA, nullptr);
+  BranchLivenessInfo BranchLiveness =
+      BC->MIB->createBranchLivenessInfo(*BF, DIM);
 
-  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &DIM));
+  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
 }
 
 TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {

>From e1791b36729b17ba4c039621446d65c0d8bb613e Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Mon, 20 Jul 2026 19:40:36 +0100
Subject: [PATCH 3/6] * update offset calculation in addTrampolineAfter * use
 set instead of map for branch liveness * defer liveness analysis until we
 encounter compare-and-branch inst * add minor comment

---
 bolt/include/bolt/Core/MCPlusBuilder.h        |  8 ++---
 bolt/lib/Passes/LongJmp.cpp                   | 33 ++++++++++---------
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 13 ++++++--
 3 files changed, 31 insertions(+), 23 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index c2859561910d5..d0e76032c7222 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -76,11 +76,10 @@ enum BTIKind {
 };
 
 struct BranchLivenessInfo {
-  DenseMap<const MCInst *, bool> FlagsLiveIn;
+  DenseSet<const MCInst *> FlagsDead;
 
   bool mustPreserveFlags(const MCInst &Inst) const {
-    auto It = FlagsLiveIn.find(&Inst);
-    return It == FlagsLiveIn.end() || It->second;
+    return !FlagsDead.count(&Inst);
   }
 };
 
@@ -2160,7 +2159,8 @@ class MCPlusBuilder {
   }
 
   /// Reverses the branch condition in Inst and update its taken target to TBB.
-  /// Assumes that the branch is reversible.
+  /// Assumes that the branch is reversible. It may replace Inst with a longer
+  /// instruction sequence on some targets.
   virtual void
   reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
                          const MCSymbol *TBB, MCContext *Ctx,
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 871791e6b7988..0e8903697612d 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -711,18 +711,20 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
     DenseMap<const BinaryBasicBlock *, BinaryBasicBlock *> FragmentTrampolines;
 
     // Create a trampoline code after \p BB or at the end of the fragment if BB
-    // is nullptr. \p Offset is the fragment size delta caused by the insertion,
-    // including any growth of \p BB before the trampoline.
+    // is nullptr. \p Offset reflects the size delta of BB caused by splitting
+    // unconditional branches, or replacing a branch with a longer instruction
+    // sequence. It is used to update the output addresses of basic blocks
+    // following the trampoline.
     auto addTrampolineAfter = [&](BinaryBasicBlock *BB,
                                   BinaryBasicBlock *TargetBB, uint64_t Count,
-                                  uint64_t Offset) {
+                                  uint64_t Offset = 0) {
       FunctionTrampolines.emplace_back(BB ? BB : FF.back(),
                                        BF.createBasicBlock());
       BinaryBasicBlock *TrampolineBB = FunctionTrampolines.back().second.get();
-      const uint64_t BBGrowth = Offset ? Offset - TrampolineSize : 0;
-
-      if (BB)
-        BB->setOutputEndAddress(BB->getOutputEndAddress() + BBGrowth);
+      const uint64_t OldBBEnd = BB ? BB->getOutputEndAddress() : 0;
+      if (BB && Offset)
+        BB->setOutputEndAddress(OldBBEnd + Offset);
+      Offset += TrampolineSize;
 
       MCInst Inst;
       {
@@ -763,7 +765,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
 
       // Update offsets for blocks after BB.
       for (BinaryBasicBlock *IBB : FF)
-        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
+        adjustBasicBlockAddress(IBB, OldBBEnd, Offset);
 
       // Update offsets for trampolines in this fragment that are placed after
       // the new trampoline. Note that trampoline blocks are not part of the
@@ -775,7 +777,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
           continue;
         if (IBB == TrampolineBB)
           continue;
-        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
+        adjustBasicBlockAddress(IBB, OldBBEnd, Offset);
       }
 
       return TrampolineBB;
@@ -796,7 +798,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
 
       BinaryBasicBlock *TrampolineBB =
-          addTrampolineAfter(BB, TargetBB, BI.Count, /*Offset=*/0);
+          addTrampolineAfter(BB, TargetBB, BI.Count, /*Offset=*/-4);
       BB->replaceSuccessor(TargetBB, TrampolineBB, BI.Count);
     }
 
@@ -829,8 +831,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // case we will need further relaxation.
       const int64_t OffsetToEnd = FragmentSize - InstAddress;
       if (Count == 0 && isBranchOffsetInRange(Inst, OffsetToEnd)) {
-        TrampolineBB =
-            addTrampolineAfter(nullptr, TargetBB, Count, TrampolineSize);
+        TrampolineBB = addTrampolineAfter(nullptr, TargetBB, Count);
         BB->replaceSuccessor(TargetBB, TrampolineBB, Count);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
@@ -849,7 +850,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       if (ShouldReverseBranch && !IsReversibleBranch) {
         const uint64_t NextCount = BB->getBranchInfo(*NextBB).Count;
         BinaryBasicBlock *FallThrough =
-            addTrampolineAfter(BB, NextBB, NextCount, TrampolineSize);
+            addTrampolineAfter(BB, NextBB, NextCount);
         BB->replaceSuccessor(NextBB, FallThrough, NextCount);
       }
 
@@ -864,11 +865,11 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         const uint64_t NewBBSize = BB->estimateSize();
 
         // Create a trampoline basic block for the original taken target.
-        TrampolineBB = addTrampolineAfter(
-            BB, TargetBB, Count, TrampolineSize + (NewBBSize - OldBBSize));
+        TrampolineBB =
+            addTrampolineAfter(BB, TargetBB, Count, NewBBSize - OldBBSize);
       } else {
         // Create a trampoline basic block for the taken target of the branch.
-        TrampolineBB = addTrampolineAfter(BB, TargetBB, Count, TrampolineSize);
+        TrampolineBB = addTrampolineAfter(BB, TargetBB, Count);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
       }
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 8025f75b789e1..3503cd7df4750 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2194,12 +2194,19 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
   BranchLivenessInfo
   createBranchLivenessInfo(BinaryFunction &BF,
                            DataflowInfoManager &DIM) const override {
-    BranchLivenessInfo Info;
-    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+    SmallVector<MCInst *> CompAndBranchInsts;
     for (BinaryBasicBlock &BB : BF)
       for (MCInst &Inst : BB)
         if (isCompAndBranch(Inst))
-          Info.FlagsLiveIn[&Inst] = LA.getLiveIn(Inst).test(getFlagsReg());
+          CompAndBranchInsts.push_back(&Inst);
+    if (CompAndBranchInsts.empty())
+      return {};
+
+    BranchLivenessInfo Info;
+    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+    for (MCInst *Inst : CompAndBranchInsts)
+      if (!LA.getLiveIn(*Inst).test(getFlagsReg()))
+        Info.FlagsDead.insert(Inst);
     return Info;
   }
 

>From 152117b61c0e5c6aea369e6be022a90d024f7603 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Mon, 3 Aug 2026 17:24:15 +0100
Subject: [PATCH 4/6] changes:

- decoupled dependency between passes and target library
- renamed the cmdline option
- removed extraneous setOutputEndAddress after refactoring addTrampolineAfter
- removed call graph construction and only invoked register analysis when the function has cmpbr instructions
---
 bolt/include/bolt/Core/BinaryFunction.h       |  5 +-
 bolt/include/bolt/Core/MCPlusBuilder.h        | 28 ++-------
 .../include/bolt/Passes/BranchLivenessUtils.h | 31 ++++++++++
 bolt/include/bolt/Passes/LongJmp.h            |  9 ++-
 bolt/include/bolt/Utils/CommandLineOpts.h     |  2 +-
 bolt/lib/Core/BinaryFunction.cpp              |  9 ++-
 bolt/lib/Passes/BinaryPasses.cpp              | 44 +++++++-------
 bolt/lib/Passes/BranchLivenessUtils.cpp       | 57 +++++++++++++++++++
 bolt/lib/Passes/CMakeLists.txt                |  1 +
 bolt/lib/Passes/LongJmp.cpp                   | 45 ++++++++-------
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 35 +++---------
 bolt/lib/Target/AArch64/CMakeLists.txt        |  6 +-
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  |  6 +-
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |  6 +-
 bolt/lib/Utils/CommandLineOpts.cpp            |  6 +-
 .../AArch64/compare-and-branch-inversion.S    |  2 +-
 bolt/unittests/Core/CMakeLists.txt            |  1 -
 bolt/unittests/Core/MCPlusBuilder.cpp         | 31 ++--------
 18 files changed, 182 insertions(+), 142 deletions(-)
 create mode 100644 bolt/include/bolt/Passes/BranchLivenessUtils.h
 create mode 100644 bolt/lib/Passes/BranchLivenessUtils.cpp

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 8dc98630d8c39..6128af7f5a941 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -36,6 +36,7 @@
 #include "bolt/Core/MCPlus.h"
 #include "bolt/Utils/NameResolver.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
@@ -66,8 +67,6 @@ class DWARFUnit;
 
 namespace bolt {
 
-struct BranchLivenessInfo;
-
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
 /// Types of macro-fusion alignment corrections.
@@ -2506,7 +2505,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches(const BranchLivenessInfo *BLI = nullptr);
+  void fixBranches(const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 45e8f82220f9e..0879fa5f74b55 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -18,7 +18,6 @@
 #include "bolt/Core/Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/MC/MCAsmBackend.h"
@@ -53,7 +52,6 @@ namespace bolt {
 class BinaryBasicBlock;
 class BinaryContext;
 class BinaryFunction;
-class DataflowInfoManager;
 
 /// Different types of indirect branches encountered during disassembly.
 enum class IndirectBranchType : char {
@@ -75,14 +73,6 @@ enum BTIKind {
   JC /// Accepting both.
 };
 
-struct BranchLivenessInfo {
-  DenseSet<const MCInst *> FlagsDead;
-
-  bool mustPreserveFlags(const MCInst &Inst) const {
-    return !FlagsDead.count(&Inst);
-  }
-};
-
 class MCPlusBuilder {
 public:
   using AllocatorIdTy = uint16_t;
@@ -484,16 +474,9 @@ class MCPlusBuilder {
     return false;
   }
 
-  /// Return liveness info required for branch transformations.
-  virtual BranchLivenessInfo
-  createBranchLivenessInfo(BinaryFunction &BF, DataflowInfoManager &DIM) const {
-    return BranchLivenessInfo();
-  }
-
   /// Check whether this conditional branch can be reversed
-  virtual bool
-  isReversibleBranch(const MCInst &Inst,
-                     const BranchLivenessInfo *BLI = nullptr) const {
+  virtual bool isReversibleBranch(const MCInst &Inst,
+                                  bool MustPreserveFlags = true) const {
     assert(!isUnsupportedInstruction(Inst) && isConditionalBranch(Inst) &&
            "Instruction is not known conditional branch");
 
@@ -2171,10 +2154,9 @@ class MCPlusBuilder {
   /// Reverses the branch condition in Inst and update its taken target to TBB.
   /// Assumes that the branch is reversible. It may replace Inst with a longer
   /// instruction sequence on some targets.
-  virtual void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         const BranchLivenessInfo *BLI = nullptr) const {
+  virtual void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                                      const MCSymbol *TBB, MCContext *Ctx,
+                                      bool MustPreserveFlags = true) const {
     llvm_unreachable("not implemented");
   }
 
diff --git a/bolt/include/bolt/Passes/BranchLivenessUtils.h b/bolt/include/bolt/Passes/BranchLivenessUtils.h
new file mode 100644
index 0000000000000..aca9df6ef0b5c
--- /dev/null
+++ b/bolt/include/bolt/Passes/BranchLivenessUtils.h
@@ -0,0 +1,31 @@
+//===- bolt/Passes/BranchLivenessUtils.h ------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_PASSES_BRANCHLIVENESSUTILS_H
+#define BOLT_PASSES_BRANCHLIVENESSUTILS_H
+
+#include "llvm/ADT/DenseSet.h"
+
+namespace llvm {
+class MCInst;
+
+namespace bolt {
+class BinaryFunction;
+class RegAnalysis;
+
+/// Return true if \p BF has short-range branches.
+bool hasShortRangeBranch(BinaryFunction &BF);
+
+/// Return the branch instructions where the target flags register is dead.
+DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
+                                                 RegAnalysis &RA);
+
+} // namespace bolt
+} // namespace llvm
+
+#endif
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index 90e10943c0f70..9aedfcf44b731 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -10,11 +10,12 @@
 #define BOLT_PASSES_LONGJMP_H
 
 #include "bolt/Passes/BinaryPasses.h"
+#include "llvm/ADT/DenseSet.h"
 
 namespace llvm {
-namespace bolt {
+class MCInst;
 
-struct BranchLivenessInfo;
+namespace bolt {
 
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
@@ -76,7 +77,9 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(BinaryFunction &BF, const BranchLivenessInfo *BLI);
+  void relaxLocalBranches(
+      BinaryFunction &BF,
+      const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/include/bolt/Utils/CommandLineOpts.h b/bolt/include/bolt/Utils/CommandLineOpts.h
index 88b56217f9512..51ca7702749ce 100644
--- a/bolt/include/bolt/Utils/CommandLineOpts.h
+++ b/bolt/include/bolt/Utils/CommandLineOpts.h
@@ -134,7 +134,7 @@ extern llvm::cl::opt<unsigned> Verbosity;
 
 // Option to control whether liveness analysis should be used by
 // FixupBranches and LongJmpPass. Needed for branch inversion on AArch64.
-extern llvm::cl::opt<bool> LivenessAnalysis;
+extern llvm::cl::opt<bool> FixBranchesWithLiveness;
 
 /// Return true if we should process all functions in the binary.
 bool processAllFunctions();
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index f60f0716f4b64..be5f7f9d4cd0e 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3676,7 +3676,8 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
+void BinaryFunction::fixBranches(
+    const DenseSet<const MCInst *> *DeadFlagBranches) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3735,7 +3736,9 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        if (!MIB->isReversibleBranch(*CondBranch, BLI)) {
+        bool MustPreserveFlags =
+            !DeadFlagBranches || !DeadFlagBranches->count(CondBranch);
+        if (!MIB->isReversibleBranch(*CondBranch, MustPreserveFlags)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3746,7 +3749,7 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
         MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
-                                    Ctx, BLI);
+                                    Ctx, MustPreserveFlags);
         return true;
       };
 
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 0a00458777f8f..2197c7342be6c 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -11,18 +11,21 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/BinaryPasses.h"
-#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/FunctionLayout.h"
 #include "bolt/Core/ParallelUtilities.h"
-#include "bolt/Passes/DataflowInfoManager.h"
+#include "bolt/Passes/BranchLivenessUtils.h"
+#include "bolt/Passes/RegAnalysis.h"
 #include "bolt/Passes/ReorderAlgorithm.h"
 #include "bolt/Passes/ReorderFunctions.h"
 #include "bolt/Utils/CommandLineOpts.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include <atomic>
 #include <cmath>
 #include <mutex>
 #include <numeric>
+#include <optional>
 #include <vector>
 
 #define DEBUG_TYPE "bolt-opts"
@@ -555,25 +558,26 @@ bool ReorderBasicBlocks::modifyFunctionLayout(BinaryFunction &BF,
 }
 
 Error FixupBranches::runOnFunctions(BinaryContext &BC) {
-  auto forEachFunction = [&](auto &&Apply) {
-    for (auto &It : BC.getBinaryFunctions()) {
-      BinaryFunction &Function = It.second;
-      if (!BC.shouldEmit(Function) || !Function.isSimple())
-        continue;
-      Apply(Function);
-    }
-  };
+  const bool ShouldRunRegisterAnalysis =
+      opts::FixBranchesWithLiveness &&
+      llvm::any_of(BC.getBinaryFunctions(), [&](auto &It) {
+        BinaryFunction &BF = It.second;
+        return BC.shouldEmit(BF) && BF.isSimple() && hasShortRangeBranch(BF);
+      });
 
-  if (opts::LivenessAnalysis) {
-    BinaryFunctionCallGraph CG = buildCallGraph(BC);
-    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
-    forEachFunction([&](BinaryFunction &BF) {
-      DataflowInfoManager DIM(BF, &RA, nullptr);
-      BranchLivenessInfo Info = BC.MIB->createBranchLivenessInfo(BF, DIM);
-      BF.fixBranches(&Info);
-    });
-  } else
-    forEachFunction([&](BinaryFunction &BF) { BF.fixBranches(nullptr); });
+  std::optional<RegAnalysis> RA;
+  if (ShouldRunRegisterAnalysis)
+    RA.emplace(BC, nullptr, nullptr);
+
+  for (auto &It : BC.getBinaryFunctions()) {
+    BinaryFunction &BF = It.second;
+    if (!BC.shouldEmit(BF) || !BF.isSimple())
+      continue;
+
+    DenseSet<const MCInst *> DeadFlagBranches =
+        RA ? computeDeadFlagBranches(BF, *RA) : DenseSet<const MCInst *>();
+    BF.fixBranches(&DeadFlagBranches);
+  }
   return Error::success();
 }
 
diff --git a/bolt/lib/Passes/BranchLivenessUtils.cpp b/bolt/lib/Passes/BranchLivenessUtils.cpp
new file mode 100644
index 0000000000000..9ab5de3570509
--- /dev/null
+++ b/bolt/lib/Passes/BranchLivenessUtils.cpp
@@ -0,0 +1,57 @@
+//===- bolt/Passes/BranchLivenessUtils.cpp -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Passes/BranchLivenessUtils.h"
+#include "bolt/Core/BinaryContext.h"
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/MCPlusBuilder.h"
+#include "bolt/Passes/DataflowInfoManager.h"
+#include "bolt/Passes/RegAnalysis.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/MC/MCRegister.h"
+
+namespace llvm {
+namespace bolt {
+
+bool hasShortRangeBranch(BinaryFunction &BF) {
+  BinaryContext &BC = BF.getBinaryContext();
+  if (!BC.isAArch64())
+    return false;
+
+  return llvm::any_of(BF, [&](BinaryBasicBlock &BB) {
+    return llvm::any_of(
+        BB, [&](MCInst &Inst) { return BC.MIB->isShortRangeBranch(Inst); });
+  });
+}
+
+DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
+                                                 RegAnalysis &RA) {
+  BinaryContext &BC = BF.getBinaryContext();
+  SmallVector<MCInst *, 4> Insts;
+  if (BC.isAArch64())
+    for (BinaryBasicBlock &BB : BF)
+      for (MCInst &Inst : BB)
+        if (BC.MIB->isShortRangeBranch(Inst))
+          Insts.push_back(&Inst);
+
+  DenseSet<const MCInst *> DeadFlagBranches;
+  if (Insts.empty())
+    return DeadFlagBranches;
+
+  DataflowInfoManager DIM(BF, &RA, nullptr);
+  LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+  const MCPhysReg FlagsReg = BC.MIB->getFlagsReg();
+  for (MCInst *Inst : Insts)
+    if (!LA.getLiveIn(*Inst).test(FlagsReg))
+      DeadFlagBranches.insert(Inst);
+  return DeadFlagBranches;
+}
+
+} // namespace bolt
+} // namespace llvm
diff --git a/bolt/lib/Passes/CMakeLists.txt b/bolt/lib/Passes/CMakeLists.txt
index ec012f05cc498..686dee6987a73 100644
--- a/bolt/lib/Passes/CMakeLists.txt
+++ b/bolt/lib/Passes/CMakeLists.txt
@@ -4,6 +4,7 @@ add_llvm_library(LLVMBOLTPasses
   AllocCombiner.cpp
   AsmDump.cpp
   BinaryPasses.cpp
+  BranchLivenessUtils.cpp
   CMOVConversion.cpp
   CacheMetrics.cpp
   DataflowAnalysis.cpp
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 0e8903697612d..497988728a527 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -11,10 +11,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/LongJmp.h"
-#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/ParallelUtilities.h"
-#include "bolt/Passes/DataflowInfoManager.h"
+#include "bolt/Passes/BranchLivenessUtils.h"
+#include "bolt/Passes/RegAnalysis.h"
 #include "bolt/Utils/CommandLineOpts.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/MathExtras.h"
 
 #define DEBUG_TYPE "longjmp"
@@ -664,8 +666,8 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
   return Error::success();
 }
 
-void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
-                                     const BranchLivenessInfo *BLI) {
+void LongJmpPass::relaxLocalBranches(
+    BinaryFunction &BF, const DenseSet<const MCInst *> *DeadFlagBranches) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -792,7 +794,6 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
 
       const MCSymbol *TargetSymbol = MIB->getTargetSymbol(*Inst);
       BB->eraseInstruction(BB->findInstruction(Inst));
-      BB->setOutputEndAddress(BB->getOutputEndAddress() - TrampolineSize);
 
       BinaryBasicBlock::BinaryBranchInfo BI;
       BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
@@ -842,7 +843,10 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, BLI);
+      bool MustPreserveFlags =
+          !DeadFlagBranches || !DeadFlagBranches->count(&Inst);
+      bool IsReversibleBranch =
+          MIB->isReversibleBranch(Inst, MustPreserveFlags);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -860,7 +864,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         {
           auto L = BC.scopeLock();
           MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(),
-                                      BC.Ctx.get(), BLI);
+                                      BC.Ctx.get(), MustPreserveFlags);
         }
         const uint64_t NewBBSize = BB->estimateSize();
 
@@ -956,22 +960,25 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
-  DenseMap<BinaryFunction *, BranchLivenessInfo> BranchLiveness;
-  if (opts::LivenessAnalysis) {
-    BinaryFunctionCallGraph CG = buildCallGraph(BC);
-    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
+  DenseMap<BinaryFunction *, DenseSet<const MCInst *>> DeadFlagBranches;
+  if (opts::FixBranchesWithLiveness) {
+    SmallVector<BinaryFunction *> Candidates;
     for (auto &It : BC.getBinaryFunctions()) {
       BinaryFunction &BF = It.second;
       if (!BC.shouldEmit(BF) || !BF.isSimple())
         continue;
-      DataflowInfoManager DIM(BF, &RA, nullptr);
-      BranchLiveness[&BF] = BC.MIB->createBranchLivenessInfo(BF, DIM);
+      if (hasShortRangeBranch(BF))
+        Candidates.push_back(&BF);
+    }
+    if (!Candidates.empty()) {
+      RegAnalysis RA(BC, nullptr, nullptr);
+      for (BinaryFunction *BF : Candidates)
+        DeadFlagBranches[BF] = computeDeadFlagBranches(*BF, RA);
     }
   }
-  auto getBranchLiveness =
-      [&](BinaryFunction &BF) -> const BranchLivenessInfo * {
-    auto It = BranchLiveness.find(&BF);
-    return It == BranchLiveness.end() ? nullptr : &It->second;
+  auto getDeadFlagBranches = [&](BinaryFunction &BF) {
+    auto It = DeadFlagBranches.find(&BF);
+    return It == DeadFlagBranches.end() ? nullptr : &It->second;
   };
 
   if (opts::CompactCodeModel) {
@@ -984,7 +991,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
         };
 
     ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF, getBranchLiveness(BF));
+      relaxLocalBranches(BF, getDeadFlagBranches(BF));
     };
 
     ParallelUtilities::runOnEachFunction(
@@ -1009,7 +1016,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches(getBranchLiveness(*Func));
+        Func->fixBranches(getDeadFlagBranches(*Func));
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 92b95b341af3f..f9bd174b38814 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -22,7 +22,6 @@
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/MCInstUtils.h"
 #include "bolt/Core/MCPlusBuilder.h"
-#include "bolt/Passes/DataflowInfoManager.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDisassembler/MCDisassembler.h"
@@ -2191,30 +2190,9 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  BranchLivenessInfo
-  createBranchLivenessInfo(BinaryFunction &BF,
-                           DataflowInfoManager &DIM) const override {
-    SmallVector<MCInst *> CompAndBranchInsts;
-    for (BinaryBasicBlock &BB : BF)
-      for (MCInst &Inst : BB)
-        if (isCompAndBranch(Inst))
-          CompAndBranchInsts.push_back(&Inst);
-    if (CompAndBranchInsts.empty())
-      return {};
-
-    BranchLivenessInfo Info;
-    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
-    for (MCInst *Inst : CompAndBranchInsts)
-      if (!LA.getLiveIn(*Inst).test(getFlagsReg()))
-        Info.FlagsDead.insert(Inst);
-    return Info;
-  }
-
-  bool
-  isReversibleBranch(const MCInst &Inst,
-                     const BranchLivenessInfo *BLI = nullptr) const override {
+  bool isReversibleBranch(const MCInst &Inst,
+                          bool MustPreserveFlags = true) const override {
     if (isCompAndBranch(Inst)) {
-      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
       if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0 &&
           MustPreserveFlags)
@@ -2226,10 +2204,11 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return MCPlusBuilder::isReversibleBranch(Inst);
   }
 
-  void reverseBranchCondition(
-      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
-      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
-    assert(isReversibleBranch(Inst, BLI) && "Irreversible branch");
+  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                              const MCSymbol *TBB, MCContext *Ctx,
+                              bool MustPreserveFlags = true) const override {
+    assert(isReversibleBranch(Inst, MustPreserveFlags) &&
+           "Irreversible branch");
 
     if (isTB(Inst) || isCB(Inst) || isCompAndBranch(Inst)) {
       bool ImmediateOutOfBounds = false;
diff --git a/bolt/lib/Target/AArch64/CMakeLists.txt b/bolt/lib/Target/AArch64/CMakeLists.txt
index e28ed0bd66ba1..1e171748aece6 100644
--- a/bolt/lib/Target/AArch64/CMakeLists.txt
+++ b/bolt/lib/Target/AArch64/CMakeLists.txt
@@ -29,11 +29,7 @@ add_llvm_library(LLVMBOLTTargetAArch64
   AArch64CommonTableGen
   )
 
-target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE
-  LLVMBOLTCore
-  LLVMBOLTPasses
-  LLVMBOLTUtils
-  )
+target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE LLVMBOLTCore LLVMBOLTUtils)
 
 include_directories(
   ${LLVM_MAIN_SRC_DIR}/lib/Target/AArch64
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index 34e8392f59743..082f9242adf22 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -162,9 +162,9 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  void reverseBranchCondition(
-      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
-      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
+  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                              const MCSymbol *TBB, MCContext *Ctx,
+                              bool MustPreserveFlags = true) const override {
     auto Opcode = getInvertedBranchOpcode(Inst.getOpcode());
     Inst.setOpcode(Opcode);
     replaceBranchTarget(Inst, TBB, Ctx);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 8dbe98e004dd2..ba84617dbd65f 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2811,9 +2811,9 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(CC));
   }
 
-  void reverseBranchCondition(
-      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
-      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
+  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                              const MCSymbol *TBB, MCContext *Ctx,
+                              bool MustPreserveFlags = true) const override {
     unsigned InvCC = getInvertedCondCode(getCondCode(Inst));
     assert(InvCC != X86::COND_INVALID && "invalid branch instruction");
     Inst.getOperand(Info->get(Inst.getOpcode()).NumOperands - 1).setImm(InvCC);
diff --git a/bolt/lib/Utils/CommandLineOpts.cpp b/bolt/lib/Utils/CommandLineOpts.cpp
index 82fde5788397a..4eaec32c7daf3 100644
--- a/bolt/lib/Utils/CommandLineOpts.cpp
+++ b/bolt/lib/Utils/CommandLineOpts.cpp
@@ -365,9 +365,9 @@ cl::opt<unsigned>
               cl::init(0), cl::ZeroOrMore, cl::cat(BoltCategory),
               cl::sub(cl::SubCommand::getAll()));
 
-cl::opt<bool> LivenessAnalysis(
-    "liveness-analysis",
-    cl::desc("use liveness analysis in FixupBranches and LongJmpPass"
+cl::opt<bool> FixBranchesWithLiveness(
+    "fix-branches-with-liveness",
+    cl::desc("use liveness analysis during branch fixup "
              "(needed for branch inversion on AArch64)"),
     cl::init(false), cl::cat(BoltCategory));
 
diff --git a/bolt/test/AArch64/compare-and-branch-inversion.S b/bolt/test/AArch64/compare-and-branch-inversion.S
index 0ea084e904101..871fce6c67f5f 100644
--- a/bolt/test/AArch64/compare-and-branch-inversion.S
+++ b/bolt/test/AArch64/compare-and-branch-inversion.S
@@ -15,7 +15,7 @@
 # RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=NO-LIVENESS
 #
 # RUN: llvm-bolt -v=1 %t -o %t.bolt --data %t.fdata --reorder-blocks=ext-tsp --compact-code-model \
-# RUN:   --liveness-analysis | FileCheck %s --check-prefix=BOLT-INFO-LIVENESS
+# RUN:   --fix-branches-with-liveness | FileCheck %s --check-prefix=BOLT-INFO-LIVENESS
 # RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=LIVENESS
 
 # COMMON: Disassembly of section .text:
diff --git a/bolt/unittests/Core/CMakeLists.txt b/bolt/unittests/Core/CMakeLists.txt
index b755afc8a66da..297dec7449202 100644
--- a/bolt/unittests/Core/CMakeLists.txt
+++ b/bolt/unittests/Core/CMakeLists.txt
@@ -24,7 +24,6 @@ target_link_libraries(CoreTests
   PRIVATE
   LLVMBOLTCore
   LLVMBOLTRewrite
-  LLVMBOLTPasses
   LLVMBOLTProfile
   LLVMBOLTUtils
   )
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 5fdaf47094770..1c206152c1fe8 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -18,9 +18,6 @@
 
 #include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Core/BinaryFunction.h"
-#include "bolt/Core/BinaryFunctionCallGraph.h"
-#include "bolt/Passes/BinaryPasses.h"
-#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
@@ -342,15 +339,9 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Underflows) {
                                   .addImm(0)
                                   .addExpr(MCSymbolRefExpr::create(
                                       TargetBB->getLabel(), *BC->Ctx.get())));
-  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
-  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
-  DataflowInfoManager DIM(*BF, &RA, nullptr);
-  BranchLivenessInfo BranchLiveness =
-      BC->MIB->createBranchLivenessInfo(*BF, DIM);
-
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/false));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &BranchLiveness);
+                                  BC->Ctx.get(), /*PreserveFlags=*/false);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSXri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::XZR);
@@ -383,15 +374,9 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
                                   .addImm(63)
                                   .addExpr(MCSymbolRefExpr::create(
                                       TargetBB->getLabel(), *BC->Ctx.get())));
-  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
-  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
-  DataflowInfoManager DIM(*BF, &RA, nullptr);
-  BranchLivenessInfo BranchLiveness =
-      BC->MIB->createBranchLivenessInfo(*BF, DIM);
-
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/false));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &BranchLiveness);
+                                  BC->Ctx.get(), /*PreserveFlags=*/false);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSWri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::WZR);
@@ -435,13 +420,7 @@ TEST_P(MCPlusBuilderTester, AArch64_IsReversibleBranch_LiveCondFlags) {
                                     .addReg(AArch64::X2)
                                     .addImm(13));
 
-  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
-  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
-  DataflowInfoManager DIM(*BF, &RA, nullptr);
-  BranchLivenessInfo BranchLiveness =
-      BC->MIB->createBranchLivenessInfo(*BF, DIM);
-
-  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
+  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/true));
 }
 
 TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {

>From 01278d952e31ecdad43a5c06a10205a30b654273 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Tue, 4 Aug 2026 13:07:18 +0100
Subject: [PATCH 5/6] wrap dense set in branch liveness info

---
 bolt/include/bolt/Core/BinaryFunction.h       |  4 ++-
 bolt/include/bolt/Core/BranchLiveness.h       | 30 +++++++++++++++++++
 .../include/bolt/Passes/BranchLivenessUtils.h | 13 ++++----
 bolt/include/bolt/Passes/LongJmp.h            | 10 +++----
 bolt/lib/Core/BinaryFunction.cpp              |  6 ++--
 bolt/lib/Passes/BinaryPasses.cpp              | 10 +++----
 bolt/lib/Passes/BranchLivenessUtils.cpp       | 15 +++++-----
 bolt/lib/Passes/LongJmp.cpp                   | 24 +++++++--------
 8 files changed, 68 insertions(+), 44 deletions(-)
 create mode 100644 bolt/include/bolt/Core/BranchLiveness.h

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 6128af7f5a941..4cd4c842a7153 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -67,6 +67,8 @@ class DWARFUnit;
 
 namespace bolt {
 
+struct BranchLivenessInfo;
+
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
 /// Types of macro-fusion alignment corrections.
@@ -2505,7 +2507,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches(const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
+  void fixBranches(const BranchLivenessInfo *BLI = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/BranchLiveness.h b/bolt/include/bolt/Core/BranchLiveness.h
new file mode 100644
index 0000000000000..80b9f3b272aac
--- /dev/null
+++ b/bolt/include/bolt/Core/BranchLiveness.h
@@ -0,0 +1,30 @@
+//===- bolt/Core/BranchLiveness.h -------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_CORE_BRANCHLIVENESS_H
+#define BOLT_CORE_BRANCHLIVENESS_H
+
+#include "llvm/ADT/DenseSet.h"
+
+namespace llvm {
+class MCInst;
+
+namespace bolt {
+
+struct BranchLivenessInfo {
+  DenseSet<const MCInst *> BranchesWithDeadFlags;
+
+  bool mustPreserveFlags(const MCInst &Inst) const {
+    return !BranchesWithDeadFlags.count(&Inst);
+  }
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif
diff --git a/bolt/include/bolt/Passes/BranchLivenessUtils.h b/bolt/include/bolt/Passes/BranchLivenessUtils.h
index aca9df6ef0b5c..e94cde4295afe 100644
--- a/bolt/include/bolt/Passes/BranchLivenessUtils.h
+++ b/bolt/include/bolt/Passes/BranchLivenessUtils.h
@@ -9,21 +9,18 @@
 #ifndef BOLT_PASSES_BRANCHLIVENESSUTILS_H
 #define BOLT_PASSES_BRANCHLIVENESSUTILS_H
 
-#include "llvm/ADT/DenseSet.h"
+#include "bolt/Core/BranchLiveness.h"
 
 namespace llvm {
-class MCInst;
-
 namespace bolt {
 class BinaryFunction;
 class RegAnalysis;
 
-/// Return true if \p BF has short-range branches.
-bool hasShortRangeBranch(BinaryFunction &BF);
+/// Return true if \p BF needs liveness info for branch transformations.
+bool needsBranchLiveness(BinaryFunction &BF);
 
-/// Return the branch instructions where the target flags register is dead.
-DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
-                                                 RegAnalysis &RA);
+/// Return liveness info required for branch transformations.
+BranchLivenessInfo computeBranchLiveness(BinaryFunction &BF, RegAnalysis &RA);
 
 } // namespace bolt
 } // namespace llvm
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index 9aedfcf44b731..e9f2b29ff9588 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -10,13 +10,12 @@
 #define BOLT_PASSES_LONGJMP_H
 
 #include "bolt/Passes/BinaryPasses.h"
-#include "llvm/ADT/DenseSet.h"
 
 namespace llvm {
-class MCInst;
-
 namespace bolt {
 
+struct BranchLivenessInfo;
+
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
 /// pull this pass inside BOLT because here we can do a better job at stub
@@ -77,9 +76,8 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(
-      BinaryFunction &BF,
-      const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
+  void relaxLocalBranches(BinaryFunction &BF,
+                          const BranchLivenessInfo *BLI = nullptr);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index be5f7f9d4cd0e..f9c48c56f9775 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -12,6 +12,7 @@
 
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/BinaryBasicBlock.h"
+#include "bolt/Core/BranchLiveness.h"
 #include "bolt/Core/DynoStats.h"
 #include "bolt/Core/HashUtilities.h"
 #include "bolt/Core/MCPlusBuilder.h"
@@ -3676,8 +3677,7 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches(
-    const DenseSet<const MCInst *> *DeadFlagBranches) {
+void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3737,7 +3737,7 @@ void BinaryFunction::fixBranches(
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
         bool MustPreserveFlags =
-            !DeadFlagBranches || !DeadFlagBranches->count(CondBranch);
+            BLI ? BLI->mustPreserveFlags(*CondBranch) : true;
         if (!MIB->isReversibleBranch(*CondBranch, MustPreserveFlags)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 2197c7342be6c..2637993ae4d4d 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -18,7 +18,6 @@
 #include "bolt/Passes/ReorderAlgorithm.h"
 #include "bolt/Passes/ReorderFunctions.h"
 #include "bolt/Utils/CommandLineOpts.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include <atomic>
@@ -562,7 +561,7 @@ Error FixupBranches::runOnFunctions(BinaryContext &BC) {
       opts::FixBranchesWithLiveness &&
       llvm::any_of(BC.getBinaryFunctions(), [&](auto &It) {
         BinaryFunction &BF = It.second;
-        return BC.shouldEmit(BF) && BF.isSimple() && hasShortRangeBranch(BF);
+        return BC.shouldEmit(BF) && BF.isSimple() && needsBranchLiveness(BF);
       });
 
   std::optional<RegAnalysis> RA;
@@ -574,9 +573,10 @@ Error FixupBranches::runOnFunctions(BinaryContext &BC) {
     if (!BC.shouldEmit(BF) || !BF.isSimple())
       continue;
 
-    DenseSet<const MCInst *> DeadFlagBranches =
-        RA ? computeDeadFlagBranches(BF, *RA) : DenseSet<const MCInst *>();
-    BF.fixBranches(&DeadFlagBranches);
+    BranchLivenessInfo BLI;
+    if (RA)
+      BLI = computeBranchLiveness(BF, *RA);
+    BF.fixBranches(&BLI);
   }
   return Error::success();
 }
diff --git a/bolt/lib/Passes/BranchLivenessUtils.cpp b/bolt/lib/Passes/BranchLivenessUtils.cpp
index 9ab5de3570509..6bc1c159377bc 100644
--- a/bolt/lib/Passes/BranchLivenessUtils.cpp
+++ b/bolt/lib/Passes/BranchLivenessUtils.cpp
@@ -19,7 +19,7 @@
 namespace llvm {
 namespace bolt {
 
-bool hasShortRangeBranch(BinaryFunction &BF) {
+bool needsBranchLiveness(BinaryFunction &BF) {
   BinaryContext &BC = BF.getBinaryContext();
   if (!BC.isAArch64())
     return false;
@@ -30,27 +30,26 @@ bool hasShortRangeBranch(BinaryFunction &BF) {
   });
 }
 
-DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
-                                                 RegAnalysis &RA) {
+BranchLivenessInfo computeBranchLiveness(BinaryFunction &BF, RegAnalysis &RA) {
   BinaryContext &BC = BF.getBinaryContext();
-  SmallVector<MCInst *, 4> Insts;
+  SmallVector<MCInst *> Insts;
   if (BC.isAArch64())
     for (BinaryBasicBlock &BB : BF)
       for (MCInst &Inst : BB)
         if (BC.MIB->isShortRangeBranch(Inst))
           Insts.push_back(&Inst);
 
-  DenseSet<const MCInst *> DeadFlagBranches;
+  BranchLivenessInfo BLI;
   if (Insts.empty())
-    return DeadFlagBranches;
+    return BLI;
 
   DataflowInfoManager DIM(BF, &RA, nullptr);
   LivenessAnalysis &LA = DIM.getLivenessAnalysis();
   const MCPhysReg FlagsReg = BC.MIB->getFlagsReg();
   for (MCInst *Inst : Insts)
     if (!LA.getLiveIn(*Inst).test(FlagsReg))
-      DeadFlagBranches.insert(Inst);
-  return DeadFlagBranches;
+      BLI.BranchesWithDeadFlags.insert(Inst);
+  return BLI;
 }
 
 } // namespace bolt
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 497988728a527..1b2cb912c574e 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -15,7 +15,6 @@
 #include "bolt/Passes/BranchLivenessUtils.h"
 #include "bolt/Passes/RegAnalysis.h"
 #include "bolt/Utils/CommandLineOpts.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/MathExtras.h"
 
@@ -666,8 +665,8 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
   return Error::success();
 }
 
-void LongJmpPass::relaxLocalBranches(
-    BinaryFunction &BF, const DenseSet<const MCInst *> *DeadFlagBranches) {
+void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
+                                     const BranchLivenessInfo *BLI) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -843,8 +842,7 @@ void LongJmpPass::relaxLocalBranches(
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool MustPreserveFlags =
-          !DeadFlagBranches || !DeadFlagBranches->count(&Inst);
+      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
       bool IsReversibleBranch =
           MIB->isReversibleBranch(Inst, MustPreserveFlags);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
@@ -960,25 +958,25 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
-  DenseMap<BinaryFunction *, DenseSet<const MCInst *>> DeadFlagBranches;
+  DenseMap<BinaryFunction *, BranchLivenessInfo> BranchLiveness;
   if (opts::FixBranchesWithLiveness) {
     SmallVector<BinaryFunction *> Candidates;
     for (auto &It : BC.getBinaryFunctions()) {
       BinaryFunction &BF = It.second;
       if (!BC.shouldEmit(BF) || !BF.isSimple())
         continue;
-      if (hasShortRangeBranch(BF))
+      if (needsBranchLiveness(BF))
         Candidates.push_back(&BF);
     }
     if (!Candidates.empty()) {
       RegAnalysis RA(BC, nullptr, nullptr);
       for (BinaryFunction *BF : Candidates)
-        DeadFlagBranches[BF] = computeDeadFlagBranches(*BF, RA);
+        BranchLiveness[BF] = computeBranchLiveness(*BF, RA);
     }
   }
-  auto getDeadFlagBranches = [&](BinaryFunction &BF) {
-    auto It = DeadFlagBranches.find(&BF);
-    return It == DeadFlagBranches.end() ? nullptr : &It->second;
+  auto getBranchLiveness = [&](BinaryFunction &BF) {
+    auto It = BranchLiveness.find(&BF);
+    return It == BranchLiveness.end() ? nullptr : &It->second;
   };
 
   if (opts::CompactCodeModel) {
@@ -991,7 +989,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
         };
 
     ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF, getDeadFlagBranches(BF));
+      relaxLocalBranches(BF, getBranchLiveness(BF));
     };
 
     ParallelUtilities::runOnEachFunction(
@@ -1016,7 +1014,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches(getDeadFlagBranches(*Func));
+        Func->fixBranches(getBranchLiveness(*Func));
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs

>From 42c20fdc5d034269b879df2df868461cf351960c Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Tue, 4 Aug 2026 13:41:06 +0100
Subject: [PATCH 6/6] nfc, stylistic changes

---
 bolt/lib/Core/BinaryFunction.cpp |  7 +++----
 bolt/lib/Passes/LongJmp.cpp      | 19 ++++++++-----------
 2 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index f9c48c56f9775..121a5e52c8c50 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3736,9 +3736,8 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        bool MustPreserveFlags =
-            BLI ? BLI->mustPreserveFlags(*CondBranch) : true;
-        if (!MIB->isReversibleBranch(*CondBranch, MustPreserveFlags)) {
+        bool PreserveFlags = BLI ? BLI->mustPreserveFlags(*CondBranch) : true;
+        if (!MIB->isReversibleBranch(*CondBranch, PreserveFlags)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3749,7 +3748,7 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
         MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
-                                    Ctx, MustPreserveFlags);
+                                    Ctx, PreserveFlags);
         return true;
       };
 
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 1b2cb912c574e..39743e2e17d90 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -842,9 +842,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
-      bool IsReversibleBranch =
-          MIB->isReversibleBranch(Inst, MustPreserveFlags);
+      bool PreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
+      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, PreserveFlags);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -862,7 +861,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         {
           auto L = BC.scopeLock();
           MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(),
-                                      BC.Ctx.get(), MustPreserveFlags);
+                                      BC.Ctx.get(), PreserveFlags);
         }
         const uint64_t NewBBSize = BB->estimateSize();
 
@@ -963,9 +962,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     SmallVector<BinaryFunction *> Candidates;
     for (auto &It : BC.getBinaryFunctions()) {
       BinaryFunction &BF = It.second;
-      if (!BC.shouldEmit(BF) || !BF.isSimple())
-        continue;
-      if (needsBranchLiveness(BF))
+      if (BC.shouldEmit(BF) && BF.isSimple() && needsBranchLiveness(BF))
         Candidates.push_back(&BF);
     }
     if (!Candidates.empty()) {
@@ -983,15 +980,15 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     BC.outs()
         << "BOLT-INFO: relaxing branches for compact code model (<128MB)\n";
 
+    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+      relaxLocalBranches(BF, getBranchLiveness(BF));
+    };
+
     ParallelUtilities::PredicateTy SkipPredicate =
         [&](const BinaryFunction &BF) {
           return !BC.shouldEmit(BF) || !BF.isSimple();
         };
 
-    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF, getBranchLiveness(BF));
-    };
-
     ParallelUtilities::runOnEachFunction(
         BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
         SkipPredicate, "RelaxLocalBranches");



More information about the llvm-commits mailing list