[CodeGen][BranchFolding] Add options to control common hoisting and block reordering (#205704)

The BranchFolder pass performs several independent optimizations: tail
 merging, common-code hoisting, and branch optimization (which includes
 basic-block reordering). Currently only tail merging is individually
controllable -- it is disabled for targets that require a structured
CFG,
because tail merging can make the CFG irreducible. Common-code hoisting
and
 basic-block reordering always run whenever the pass executes.

 Those two sub-phases do not change CFG edges (so they cannot affect
 reducibility); they only change block layout. A target whose register
 allocation is sensitive to the final block layout may need the rest of
 branch folding while suppressing reordering and/or hoisting. This is a
 layout / register-allocation concern, distinct from the structured-CFG
 (reducibility) concern that gates tail merging.

 This change makes common-code hoisting and basic-block reordering
 individually controllable:

  - BranchFolder gains an EnableBasicBlockReordering flag (with a
setBasicBlockReordering() setter) that gates the two block-reordering
    transforms in OptimizeBlock.
- BranchFolderLegacy gains constructor parameters for both sub-phases,
and a
createBranchFolder(EnableCommonHoist, EnableBasicBlockReordering)
factory
is exposed so a target can request a safe subset from its pass pipeline.
  - Two hidden command-line flags, -branch-folder-hoist-common-code and
-branch-folder-reorder-blocks, override the configured values at the
start
    of OptimizeFunction (mirroring -enable-tail-merge).

Defaults enable both sub-phases, so existing behavior is unchanged for
all
 in-tree targets.

 New MIR tests exercise each flag via -run-pass=branch-folder.

 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

llvm/include/llvm/CodeGen/Passes.h

@@ -279,6 +279,13 @@ LLVM_ABI extern char &PostRASchedulerID;
 /// branches.
 LLVM_ABI extern char &BranchFolderPassID;
 
+/// createBranchFolder - Create the BranchFolder pass, optionally disabling the
+/// common-code hoisting and/or basic-block reordering sub-phases. Default
+/// enables both (full BranchFolding behavior).
+LLVM_ABI FunctionPass *
+createBranchFolder(bool EnableCommonHoist = true,
+                   bool EnableBasicBlockReordering = true);
+
 /// BranchRelaxation - This pass replaces branches that need to jump further
 /// than is supported by a branch instruction.
 LLVM_ABI extern char &BranchRelaxationPassID;

llvm/lib/CodeGen/BranchFolding.cpp

@@ -78,6 +78,20 @@ static cl::opt<cl::boolOrDefault>
     FlagEnableTailMerge("enable-tail-merge",
                         cl::init(cl::boolOrDefault::BOU_UNSET), cl::Hidden);
 
+// Override the common-code hoisting sub-phase of BranchFolding. Unset by
+// default, in which case the value configured by the caller is used.
+static cl::opt<cl::boolOrDefault> FlagEnableHoistCommonCode(
+    "branch-folder-hoist-common-code", cl::init(cl::boolOrDefault::BOU_UNSET),
+    cl::Hidden,
+    cl::desc("Override common-code hoisting in the BranchFolding pass"));
+
+// Override the basic-block reordering sub-phase of BranchFolding. Unset by
+// default, in which case the value configured by the caller is used.
+static cl::opt<cl::boolOrDefault> FlagEnableBlockReordering(
+    "branch-folder-reorder-blocks", cl::init(cl::boolOrDefault::BOU_UNSET),
+    cl::Hidden,
+    cl::desc("Override basic-block reordering in the BranchFolding pass"));
+
 // Throttle for huge numbers of predecessors (compile speed problems)
 static cl::opt<unsigned>
 TailMergeThreshold("tail-merge-threshold",
@@ -94,10 +108,16 @@ namespace {
 
   /// BranchFolderPass - Wrap branch folder in a machine function pass.
 class BranchFolderLegacy : public MachineFunctionPass {
+  bool EnableCommonHoist;
+  bool EnableBasicBlockReordering;
+
 public:
   static char ID;
 
-  explicit BranchFolderLegacy() : MachineFunctionPass(ID) {}
+  explicit BranchFolderLegacy(bool EnableCommonHoist = true,
+                              bool EnableBasicBlockReordering = true)
+      : MachineFunctionPass(ID), EnableCommonHoist(EnableCommonHoist),
+        EnableBasicBlockReordering(EnableBasicBlockReordering) {}
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
@@ -141,6 +161,7 @@ PreservedAnalyses BranchFolderPass::run(MachineFunction &MF,
   MBFIWrapper MBBFreqInfo(MBFI);
   BranchFolder Folder(EnableTailMerge, /*CommonHoist=*/true, MBBFreqInfo, MBPI,
                       PSI);
+  Folder.setBasicBlockReordering(true);
   if (Folder.OptimizeFunction(MF, MF.getSubtarget().getInstrInfo(),
                               MF.getSubtarget().getRegisterInfo()))
     return getMachineFunctionPassPreservedAnalyses();
@@ -160,9 +181,10 @@ bool BranchFolderLegacy::runOnMachineFunction(MachineFunction &MF) {
   MBFIWrapper MBBFreqInfo(
       getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI());
   BranchFolder Folder(
-      EnableTailMerge, /*CommonHoist=*/true, MBBFreqInfo,
+      EnableTailMerge, EnableCommonHoist, MBBFreqInfo,
       getAnalysis<MachineBranchProbabilityInfoWrapperPass>().getMBPI(),
       &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI());
+  Folder.setBasicBlockReordering(EnableBasicBlockReordering);
   return Folder.OptimizeFunction(MF, MF.getSubtarget().getInstrInfo(),
                                  MF.getSubtarget().getRegisterInfo());
 }
@@ -171,8 +193,9 @@ BranchFolder::BranchFolder(bool DefaultEnableTailMerge, bool CommonHoist,
                            MBFIWrapper &FreqInfo,
                            const MachineBranchProbabilityInfo &ProbInfo,
                            ProfileSummaryInfo *PSI, unsigned MinTailLength)
-    : EnableHoistCommonCode(CommonHoist), MinCommonTailLength(MinTailLength),
-      MBBFreqInfo(FreqInfo), MBPI(ProbInfo), PSI(PSI) {
+    : EnableHoistCommonCode(CommonHoist), EnableBasicBlockReordering(true),
+      MinCommonTailLength(MinTailLength), MBBFreqInfo(FreqInfo), MBPI(ProbInfo),
+      PSI(PSI) {
   switch (FlagEnableTailMerge) {
   case cl::boolOrDefault::BOU_UNSET:
     EnableTailMerge = DefaultEnableTailMerge;
@@ -235,6 +258,16 @@ bool BranchFolder::OptimizeFunction(MachineFunction &MF,
   if (!UpdateLiveIns)
     MRI.invalidateLiveness();
 
+  // Command-line flags take final precedence over the caller-configured values,
+  // letting individual BranchFolding sub-phases be toggled (for tests and for
+  // targets that only want a safe subset of the optimization).
+  if (FlagEnableHoistCommonCode != cl::boolOrDefault::BOU_UNSET)
+    EnableHoistCommonCode =
+        FlagEnableHoistCommonCode == cl::boolOrDefault::BOU_TRUE;
+  if (FlagEnableBlockReordering != cl::boolOrDefault::BOU_UNSET)
+    EnableBasicBlockReordering =
+        FlagEnableBlockReordering == cl::boolOrDefault::BOU_TRUE;
+
   bool MadeChange = false;
 
   // Recalculate EH scope membership.
@@ -1526,8 +1559,8 @@ ReoptimizeBlock:
     // We consider it more likely that execution will stay in the function (e.g.
     // due to loops) than it is to exit it.  This asserts in loops etc, moving
     // the assert condition out of the loop body.
-    if (MBB->succ_empty() && !PriorCond.empty() && !PriorFBB &&
-        MachineFunction::iterator(PriorTBB) == FallThrough &&
+    if (EnableBasicBlockReordering && MBB->succ_empty() && !PriorCond.empty() &&
+        !PriorFBB && MachineFunction::iterator(PriorTBB) == FallThrough &&
         !MBB->canFallThrough()) {
       bool DoTransform = true;
 
@@ -1723,7 +1756,7 @@ ReoptimizeBlock:
   // If the prior block doesn't fall through into this block, and if this
   // block doesn't fall through into some other block, see if we can find a
   // place to move this block where a fall-through will happen.
-  if (!PrevBB.canFallThrough()) {
+  if (EnableBasicBlockReordering && !PrevBB.canFallThrough()) {
     // Now we know that there was no fall-through into this block, check to
     // see if it has a fall-through into its successor.
     bool CurFallsThru = MBB->canFallThrough();
@@ -2182,3 +2215,8 @@ bool BranchFolder::HoistCommonCodeInSuccs(MachineBasicBlock *MBB) {
   ++NumHoist;
   return true;
 }
+
+FunctionPass *llvm::createBranchFolder(bool EnableCommonHoist,
+                                       bool EnableBasicBlockReordering) {
+  return new BranchFolderLegacy(EnableCommonHoist, EnableBasicBlockReordering);
+}

llvm/lib/CodeGen/BranchFolding.h

@@ -46,6 +46,13 @@ class TargetRegisterInfo;
                           MachineLoopInfo *mli = nullptr,
                           bool AfterPlacement = false);
 
+    /// Enable or disable the basic-block reordering sub-phase of branch
+    /// optimization. Enabled by default; targets sensitive to block layout
+    /// (e.g. those with structured-CFG register allocation) can disable it.
+    void setBasicBlockReordering(bool Enable) {
+      EnableBasicBlockReordering = Enable;
+    }
+
   private:
     class MergePotentialsElt {
       unsigned Hash;
@@ -119,6 +126,7 @@ class TargetRegisterInfo;
     bool AfterBlockPlacement = false;
     bool EnableTailMerge = false;
     bool EnableHoistCommonCode = false;
+    bool EnableBasicBlockReordering = false;
     bool UpdateLiveIns = false;
     unsigned MinCommonTailLength;
     const TargetInstrInfo *TII = nullptr;

llvm/test/CodeGen/AArch64/branch-folder-hoist-disable.mir

@@ -0,0 +1,46 @@
+# RUN: llc -o - %s -mtriple=aarch64 -run-pass branch-folder -verify-machineinstrs | FileCheck %s --check-prefix=HOIST
+# RUN: llc -o - %s -mtriple=aarch64 -run-pass branch-folder -branch-folder-hoist-common-code=false -verify-machineinstrs | FileCheck %s --check-prefix=NOHOIST
+
+# Check that -branch-folder-hoist-common-code=false disables the common-code
+# hoisting sub-phase of BranchFolding. With hoisting enabled (default) the
+# common "ADDXri $x0, 1, 0" is hoisted from both successors into bb.0; with
+# hoisting disabled it must remain duplicated in bb.1 and bb.2.
+
+name: func
+tracksRegLiveness: true
+body: |
+  bb.0:
+    ; HOIST-LABEL: name: func
+    ; HOIST-LABEL: bb.0:
+    ; HOIST: $x0 = ADDXri $x0, 1, 0
+    ; HOIST: CBZX $x1, %bb.2
+    ;
+    ; NOHOIST-LABEL: name: func
+    ; NOHOIST-LABEL: bb.0:
+    ; NOHOIST-NOT: $x0 = ADDXri $x0, 1, 0
+    ; NOHOIST: CBZX $x1, %bb.2
+    liveins: $x0, $x1
+    CBZX $x1, %bb.2
+
+  bb.1:
+    ; HOIST-LABEL: bb.1:
+    ; HOIST-NOT: $x0 = ADDXri $x0, 1, 0
+    ;
+    ; NOHOIST-LABEL: bb.1:
+    ; NOHOIST: $x0 = ADDXri $x0, 1, 0
+    liveins: $x0
+    $x0 = ADDXri $x0, 1, 0
+    $x0 = ADDXri $x0, 2, 0
+    RET_ReallyLR implicit $x0
+
+  bb.2:
+    ; HOIST-LABEL: bb.2:
+    ; HOIST-NOT: $x0 = ADDXri $x0, 1, 0
+    ;
+    ; NOHOIST-LABEL: bb.2:
+    ; NOHOIST: $x0 = ADDXri $x0, 1, 0
+    liveins: $x0
+    $x0 = ADDXri $x0, 1, 0
+    $x0 = ADDXri $x0, 3, 0
+    RET_ReallyLR implicit $x0
+...

llvm/test/CodeGen/AArch64/branch-folder-reorder-disable.mir

@@ -0,0 +1,55 @@
+# RUN: llc -o - %s -mtriple=aarch64 -run-pass branch-folder -verify-machineinstrs | FileCheck %s --check-prefix=REORDER
+# RUN: llc -o - %s -mtriple=aarch64 -run-pass branch-folder -branch-folder-reorder-blocks=false -verify-machineinstrs | FileCheck %s --check-prefix=NOREORDER
+
+# Check that -branch-folder-reorder-blocks=false disables the basic-block
+# reordering sub-phase of branch optimization. bb.1 is a non-fall-through return
+# block sitting between the conditional branch in bb.0 and its taken target bb.2.
+#
+# With reordering enabled (default) BranchFolding reverses the bb.0 condition
+# (CBZX -> CBNZX) so it can fall through, and relocates the blocks. With
+# reordering disabled the original layout (CBZX; bb.1 is the return block) is
+# preserved.
+
+# REORDER-LABEL: name: func
+# REORDER: bb.0:
+# REORDER: CBNZX $x1, %bb.2
+# REORDER: bb.1:
+# REORDER: $x0 = ADDXri $x0, 1, 0
+# REORDER: $x0 = ADDXri $x0, 2, 0
+# REORDER: bb.2:
+# REORDER-NOT: ADDXri
+# REORDER: RET_ReallyLR implicit $x0
+
+# NOREORDER-LABEL: name: func
+# NOREORDER: bb.0:
+# NOREORDER: CBZX $x1, %bb.2
+# NOREORDER: bb.1:
+# NOREORDER: RET_ReallyLR implicit $x0
+# NOREORDER: bb.2:
+# NOREORDER: $x0 = ADDXri $x0, 1, 0
+# NOREORDER: $x0 = ADDXri $x0, 2, 0
+# NOREORDER: RET_ReallyLR implicit $x0
+
+name: func
+tracksRegLiveness: true
+body: |
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: $x0, $x1
+    CBZX $x1, %bb.2
+
+  bb.1:
+    liveins: $x0
+    RET_ReallyLR implicit $x0
+
+  bb.2:
+    successors: %bb.3
+    liveins: $x0
+    $x0 = ADDXri $x0, 1, 0
+    B %bb.3
+
+  bb.3:
+    liveins: $x0
+    $x0 = ADDXri $x0, 2, 0
+    RET_ReallyLR implicit $x0
+...