[CodeGen] Unify LLVM branch folding pipeline

1 files changed, 292 insertions, 4 deletions

diff --git a/backend/lib/CodeGen/MachineBranchOptimization.cpp b/backend/lib/CodeGen/MachineBranchOptimization.cpp
index d09de97..07a67ed 100644
--- a/backend/lib/CodeGen/MachineBranchOptimization.cpp
+++ b/backend/lib/CodeGen/MachineBranchOptimization.cpp
@@ -12,11 +12,15 @@
 #include "rcc/CodeGen/MachineFunctionRecords.h"
 #include "rcc/CodeGen/MachineInstr.h"
 #include "rcc/CodeGen/MachineJumpTableInfo.h"
+#include "rcc/CodeGen/MachineMemOperand.h"
+#include "rcc/CodeGen/MachineOperand.h"
+#include "rcc/CodeGen/RegisterMask.h"
 #include "rcc/CodeGen/TargetDescription.h"
 #include "rcc/CodeGen/TargetInstrInfo.h"
 
 #include <algorithm>
 #include <cstddef>
+#include <cstdint>
 #include <optional>
 #include <ranges>
 #include <string>
@@ -31,11 +35,11 @@ namespace {
   static const MachinePassContract Contract = [] {
     MachinePassContract Value{
         .Name = "machine-branch-optimization",
-        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs},
+        .Requires = {MachineFunctionProperty::NoPHIs},
         .Sets = {},
         .Clears = {},
         .RequiredAnalyses = {},
-        .VerificationStage = PipelineStage::Selected,
+        .VerificationStage = PipelineStage::FrameFinalized,
     };
     Value.ThreadSafety = PassThreadSafety::ThreadSafe;
     Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands,
@@ -80,6 +84,215 @@ namespace {
          Attributes.profileID().has_value() || isJumpTableTarget(Block, Function);
 }
 
+template <typename ValueT>
+[[nodiscard]] bool payloadEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
+  const ValueT *LeftValue = Left.getIf<ValueT>();
+  const ValueT *RightValue = Right.getIf<ValueT>();
+  return LeftValue != nullptr && RightValue != nullptr && *LeftValue == *RightValue;
+}
+
+template <typename MaskOperandT>
+[[nodiscard]] bool maskEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
+  const MaskOperandT *LeftValue = Left.getIf<MaskOperandT>();
+  const MaskOperandT *RightValue = Right.getIf<MaskOperandT>();
+  if (LeftValue == nullptr || RightValue == nullptr)
+    return false;
+  if (!LeftValue->Mask || !RightValue->Mask)
+    return LeftValue->Mask == RightValue->Mask;
+  return *LeftValue->Mask == *RightValue->Mask;
+}
+
+[[nodiscard]] bool operandEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
+  if (Left.kind() != Right.kind() || Left.targetFlags() != Right.targetFlags() ||
+      Left.tiedOperandIndex() != Right.tiedOperandIndex()) {
+    return false;
+  }
+  if (Left.kind() == MachineOperand::Kind::Register) {
+    const Register *LeftReg = Left.getIf<Register>();
+    const Register *RightReg = Right.getIf<Register>();
+    if (LeftReg == nullptr || RightReg == nullptr || *LeftReg != *RightReg ||
+        Left.subReg() != Right.subReg()) {
+      return false;
+    }
+    RegisterOperandFlags LeftFlags = Left.registerFlags();
+    RegisterOperandFlags RightFlags = Right.registerFlags();
+    for (const RegisterOperandFlag Flag :
+         {RegisterOperandFlag::Kill, RegisterOperandFlag::Dead, RegisterOperandFlag::Renamable}) {
+      LeftFlags.clear(Flag);
+      RightFlags.clear(Flag);
+    }
+    return LeftFlags == RightFlags;
+  }
+  switch (Left.kind()) {
+  case MachineOperand::Kind::Register:
+    return false;
+  case MachineOperand::Kind::Immediate:
+    return payloadEquals<std::int64_t>(Left, Right);
+  case MachineOperand::Kind::IntegerImmediate:
+    return payloadEquals<IntegerImmediateOperand>(Left, Right);
+  case MachineOperand::Kind::FloatingImmediate:
+    return payloadEquals<FloatingImmediateOperand>(Left, Right);
+  case MachineOperand::Kind::BasicBlock:
+    return payloadEquals<BasicBlockOperand>(Left, Right);
+  case MachineOperand::Kind::FrameIndex:
+    return payloadEquals<FrameIndexOperand>(Left, Right);
+  case MachineOperand::Kind::ConstantPoolIndex:
+    return payloadEquals<ConstantPoolOperand>(Left, Right);
+  case MachineOperand::Kind::JumpTableIndex:
+    return payloadEquals<JumpTableOperand>(Left, Right);
+  case MachineOperand::Kind::GlobalSymbol:
+    return payloadEquals<GlobalSymbolOperand>(Left, Right);
+  case MachineOperand::Kind::ExternalSymbol:
+    return payloadEquals<ExternalSymbolOperand>(Left, Right);
+  case MachineOperand::Kind::BlockAddress:
+    return payloadEquals<BlockAddressOperand>(Left, Right);
+  case MachineOperand::Kind::RegisterMask:
+    return maskEquals<RegisterMaskOperand>(Left, Right);
+  case MachineOperand::Kind::LiveOutRegisterMask:
+    return maskEquals<LiveOutRegisterMaskOperand>(Left, Right);
+  case MachineOperand::Kind::Metadata:
+    return payloadEquals<MetadataOperand>(Left, Right);
+  case MachineOperand::Kind::TargetIndex:
+    return payloadEquals<TargetIndexOperand>(Left, Right);
+  case MachineOperand::Kind::MCSymbol:
+    return payloadEquals<MCSymbolOperand>(Left, Right);
+  case MachineOperand::Kind::CFIIndex:
+    return payloadEquals<CFIIndexOperand>(Left, Right);
+  case MachineOperand::Kind::Intrinsic:
+    return payloadEquals<IntrinsicOperand>(Left, Right);
+  case MachineOperand::Kind::Predicate:
+    return payloadEquals<PredicateOperand>(Left, Right);
+  case MachineOperand::Kind::ShuffleMask:
+    return payloadEquals<ShuffleMaskOperand>(Left, Right);
+  case MachineOperand::Kind::DebugInstrRef:
+    return payloadEquals<DebugInstrRefOperand>(Left, Right);
+  case MachineOperand::Kind::LaneMask:
+    return payloadEquals<LaneMaskOperand>(Left, Right);
+  case MachineOperand::Kind::Count:
+    return false;
+  }
+  return false;
+}
+
+[[nodiscard]] Expected<bool> instructionEquals(const MachineInstr &Left, const MachineInstr &Right,
+                                               const TargetDescription &Target) {
+  const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Left.opcode());
+  if (Descriptor == nullptr || Target.instructionDescriptor(Right.opcode()) == nullptr) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "branch folding encountered an unknown instruction");
+  }
+  if (Left.opcode() != Right.opcode() || Left.flags() != Right.flags() ||
+      Left.operandCount() != Right.operandCount() ||
+      Left.memOperands().size() != Right.memOperands().size() ||
+      Left.sourceLocation() != Right.sourceLocation() ||
+      Left.debugLocation() != Right.debugLocation() || Left.origin() != Right.origin() ||
+      hasProtectedAttachment(Left) || hasProtectedAttachment(Right) ||
+      Descriptor->Flags.has(InstructionFlag::Call) ||
+      Descriptor->Flags.has(InstructionFlag::Convergent) ||
+      Descriptor->Flags.has(InstructionFlag::NotDuplicable) ||
+      Left.flags().has(MachineInstrFlag::NoMerge) ||
+      Left.flags().has(MachineInstrFlag::BundledWithPred) ||
+      Left.flags().has(MachineInstrFlag::BundledWithSucc)) {
+    return false;
+  }
+  for (std::size_t Index = 0; Index < Left.operandCount(); ++Index) {
+    const MachineOperand *LeftOperand = Left.operandAt(Index);
+    const MachineOperand *RightOperand = Right.operandAt(Index);
+    if (LeftOperand == nullptr || RightOperand == nullptr ||
+        !operandEquals(*LeftOperand, *RightOperand)) {
+      return false;
+    }
+  }
+  for (std::size_t Index = 0; Index < Left.memOperands().size(); ++Index) {
+    const MachineMemOperand *LeftMemory = Left.memOperands()[Index];
+    const MachineMemOperand *RightMemory = Right.memOperands()[Index];
+    if (LeftMemory == nullptr || RightMemory == nullptr ||
+        LeftMemory->description() != RightMemory->description()) {
+      return false;
+    }
+  }
+  return true;
+}
+
+[[nodiscard]] Expected<bool> blockEquals(const MachineBasicBlock &Left,
+                                         const MachineBasicBlock &Right,
+                                         const MachineFunction &Function,
+                                         const TargetDescription &Target) {
+  if (&Left == &Right || &Left == Function.firstBlock() || &Right == Function.firstBlock() ||
+      hasObservableBlockIdentity(Left, Function) || hasObservableBlockIdentity(Right, Function) ||
+      Left.attributes().alignment() != Right.attributes().alignment() ||
+      Left.attributes().sectionID() != Right.attributes().sectionID() ||
+      Left.attributes().callFrameSize() != Right.attributes().callFrameSize() ||
+      !std::ranges::equal(Left.liveIns(), Right.liveIns()) ||
+      !std::ranges::equal(Left.successors(), Right.successors()) ||
+      Left.instructionCount() != Right.instructionCount()) {
+    return false;
+  }
+  const MachineInstr *LeftInstruction = Left.firstInstruction();
+  const MachineInstr *RightInstruction = Right.firstInstruction();
+  while (LeftInstruction != nullptr && RightInstruction != nullptr) {
+    auto Equal = instructionEquals(*LeftInstruction, *RightInstruction, Target);
+    if (!Equal)
+      return std::unexpected(Equal.error());
+    if (!*Equal)
+      return false;
+    LeftInstruction = LeftInstruction->next();
+    RightInstruction = RightInstruction->next();
+  }
+  return LeftInstruction == nullptr && RightInstruction == nullptr;
+}
+
+[[nodiscard]] Expected<bool> canRedirectPredecessor(MachineBasicBlock &Predecessor,
+                                                    MachineBasicBlock &Duplicate,
+                                                    MachineFunction &Function,
+                                                    const TargetDescription &Target,
+                                                    const TargetInstrInfo &Info) {
+  auto Analysis = Info.analyzeBranch(Predecessor, Function, Target);
+  if (!Analysis)
+    return std::unexpected(Analysis.error());
+  if (Analysis->HasUnanalyzableTerminator)
+    return false;
+  if ((Analysis->Conditional != nullptr && Analysis->TrueTarget == &Duplicate &&
+       hasProtectedAttachment(*Analysis->Conditional)) ||
+      (Analysis->Unconditional != nullptr && Analysis->TrueTarget == &Duplicate &&
+       hasProtectedAttachment(*Analysis->Unconditional))) {
+    return false;
+  }
+  if (Analysis->TrueTarget == &Duplicate || Analysis->FalseTarget == &Duplicate)
+    return true;
+  return Analysis->Conditional == nullptr && Analysis->Unconditional == nullptr &&
+         Predecessor.successors().size() == 1 && Predecessor.next() == &Duplicate;
+}
+
+[[nodiscard]] Expected<bool>
+redirectPredecessor(MachineBasicBlock &Predecessor, MachineBasicBlock &Duplicate,
+                    MachineBasicBlock &Canonical, MachineFunction &Function,
+                    const TargetDescription &Target, const TargetInstrInfo &Info) {
+  auto Analysis = Info.analyzeBranch(Predecessor, Function, Target);
+  if (!Analysis)
+    return std::unexpected(Analysis.error());
+  auto Rewritten = Info.rewriteBranchTarget(Predecessor, Duplicate, Canonical, Function, Target);
+  if (!Rewritten)
+    return std::unexpected(Rewritten.error());
+  if (!*Rewritten) {
+    const bool IsImplicitEdge =
+        (Analysis->Conditional == nullptr && Analysis->Unconditional == nullptr &&
+         Predecessor.successors().size() == 1 && Predecessor.next() == &Duplicate) ||
+        (Analysis->Conditional != nullptr && Analysis->Unconditional == nullptr &&
+         Analysis->FalseTarget == &Duplicate);
+    if (!IsImplicitEdge)
+      return false;
+    auto Inserted = Info.insertBranch(Predecessor, Canonical, Function, Target);
+    if (!Inserted)
+      return std::unexpected(Inserted.error());
+  }
+  MachineCFGUpdater CFG = Function.cfgUpdater();
+  auto Redirected = CFG.replaceSuccessor(Predecessor, Duplicate, Canonical);
+  if (!Redirected)
+    return std::unexpected(Redirected.error());
+  return true;
+}
+
 [[nodiscard]] bool canDeleteBlock(const MachineBasicBlock &Block,
                                   const MachineFunction &Function) noexcept {
   if (hasObservableBlockIdentity(Block, Function))
@@ -288,6 +501,77 @@ namespace {
   return AnyChange;
 }
 
+[[nodiscard]] Expected<bool> mergeIdenticalTails(MachineFunction &Function,
+                                                 const TargetDescription &Target,
+                                                 const TargetInstrInfo &Info,
+                                                 std::size_t MaximumMergedBlocks) {
+  std::size_t MergedBlocks = 0;
+  bool Changed = false;
+  bool LocalChange = true;
+  while (LocalChange && MergedBlocks < MaximumMergedBlocks) {
+    LocalChange = false;
+    std::vector<MachineBasicBlock *> Blocks;
+    for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr; Block = Block->next())
+      Blocks.push_back(Block);
+    for (std::size_t DuplicateIndex = 1; DuplicateIndex < Blocks.size(); ++DuplicateIndex) {
+      MachineBasicBlock *Duplicate = Blocks[DuplicateIndex];
+      if (Duplicate == nullptr || !Duplicate->isLinked())
+        continue;
+      for (std::size_t CanonicalIndex = 1; CanonicalIndex < DuplicateIndex; ++CanonicalIndex) {
+        MachineBasicBlock *Canonical = Blocks[CanonicalIndex];
+        if (Canonical == nullptr || !Canonical->isLinked())
+          continue;
+        auto Equal = blockEquals(*Canonical, *Duplicate, Function, Target);
+        if (!Equal)
+          return std::unexpected(Equal.error());
+        if (!*Equal)
+          continue;
+
+        const std::vector<MachineBasicBlock *> Predecessors(Duplicate->predecessors().begin(),
+                                                            Duplicate->predecessors().end());
+        bool RedirectedAll = true;
+        for (MachineBasicBlock *Predecessor : Predecessors) {
+          if (Predecessor == nullptr || Predecessor == Canonical || Predecessor == Duplicate) {
+            RedirectedAll = false;
+            break;
+          }
+          auto CanRedirect =
+              canRedirectPredecessor(*Predecessor, *Duplicate, Function, Target, Info);
+          if (!CanRedirect)
+            return std::unexpected(CanRedirect.error());
+          if (!*CanRedirect) {
+            RedirectedAll = false;
+            break;
+          }
+        }
+        if (!RedirectedAll)
+          continue;
+        for (MachineBasicBlock *Predecessor : Predecessors) {
+          auto Redirected =
+              redirectPredecessor(*Predecessor, *Duplicate, *Canonical, Function, Target, Info);
+          if (!Redirected)
+            return std::unexpected(Redirected.error());
+          if (!*Redirected) {
+            return makeUnexpected(ErrorCode::CFGOptimizationFailed,
+                                  "tail merge predecessor changed after redirect preflight");
+          }
+        }
+        MachineCFGUpdater CFG = Function.cfgUpdater();
+        auto Deleted = CFG.deleteBlock(*Duplicate);
+        if (!Deleted)
+          return std::unexpected(Deleted.error());
+        ++MergedBlocks;
+        Changed = true;
+        LocalChange = true;
+        break;
+      }
+      if (LocalChange)
+        break;
+    }
+  }
+  return Changed;
+}
+
 } // namespace
 
 const MachinePassContract &MachineBranchOptimizationPass::contract() const noexcept {
@@ -299,7 +583,7 @@ Expected<MachinePassResult> MachineBranchOptimizationPass::run(MachineFunction &
                                                                MachineFunctionAnalysisManager &) {
   if (!Function.properties().containsAll(contract().Requires)) {
     return makeUnexpected(ErrorCode::VerificationError,
-                          "branch optimization requires selected PHI-free MachineIR");
+                          "branch folding requires PHI-free MachineIR");
   }
   auto Trial = cloneMachineFunction(
       Function, Target,
@@ -331,12 +615,16 @@ MachineBranchOptimizationPass::runImplementation(MachineFunction &Function,
   auto Threaded = threadForwardingBlocks(Function, Target, *Info, Options.MaximumThreadedEdges);
   if (!Threaded)
     return std::unexpected(Threaded.error());
+  auto Merged = mergeIdenticalTails(Function, Target, *Info, Options.MaximumMergedBlocks);
+  if (!Merged)
+    return std::unexpected(Merged.error());
   auto RemovedAfterThreading = removeUnreachableBlocks(Function);
   if (!RemovedAfterThreading)
     return std::unexpected(RemovedAfterThreading.error());
 
   MachinePassResult Result;
-  Result.Changed = *Folded || *RemovedBeforeThreading || *Threaded || *RemovedAfterThreading;
+  Result.Changed =
+      *Folded || *RemovedBeforeThreading || *Threaded || *Merged || *RemovedAfterThreading;
   if (!Result.Changed) {
     Result.Preserved = PreservedAnalyses::all();
     return Result;