[CodeGen] Implement RISC-V peephole branch optimization

6 files changed, 359 insertions, 1 deletions

diff --git a/backend/include/rcc/CodeGen/MachinePeepholeOptimizer.h b/backend/include/rcc/CodeGen/MachinePeepholeOptimizer.h
new file mode 100644
index 0000000..9505625
--- /dev/null
+++ b/backend/include/rcc/CodeGen/MachinePeepholeOptimizer.h
@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#ifndef RCC_CODEGEN_MACHINEPEEPHOLEOPTIMIZER_H
+#define RCC_CODEGEN_MACHINEPEEPHOLEOPTIMIZER_H
+
+#include "rcc/CodeGen/MachinePassManager.h"
+
+namespace rcc {
+
+class TargetInstrInfo;
+
+class MachinePeepholeOptimizerPass final : public MachineFunctionPass {
+public:
+  explicit MachinePeepholeOptimizerPass(const TargetInstrInfo &Info) noexcept : Info(&Info) {}
+
+  [[nodiscard]] const MachinePassContract &contract() const noexcept override;
+  [[nodiscard]] Expected<MachinePassResult> run(MachineFunction &Function,
+                                                const TargetDescription &Target,
+                                                MachineFunctionAnalysisManager &Analyses) override;
+
+private:
+  [[nodiscard]] Expected<MachinePassResult>
+  runImplementation(MachineFunction &Function, const TargetDescription &Target) const;
+
+  const TargetInstrInfo *Info;
+};
+
+} // namespace rcc
+
+#endif // RCC_CODEGEN_MACHINEPEEPHOLEOPTIMIZER_H
diff --git a/backend/include/rcc/CodeGen/TargetInstrInfo.h b/backend/include/rcc/CodeGen/TargetInstrInfo.h
index 8457ede..aada102 100644
--- a/backend/include/rcc/CodeGen/TargetInstrInfo.h
+++ b/backend/include/rcc/CodeGen/TargetInstrInfo.h
@@ -73,6 +73,11 @@ public:
     return false;
   }
 
+  [[nodiscard]] virtual Expected<bool> optimizeCondBranch(MachineInstr &, MachineFunction &,
+                                                          const TargetDescription &) const {
+    return false;
+  }
+
   [[nodiscard]] virtual Expected<bool> tryConvertToThreeAddress(MachineInstr &, std::size_t,
                                                                 std::size_t, MachineFunction &,
                                                                 const TargetDescription &) const {
diff --git a/backend/lib/CodeGen/CMakeLists.txt b/backend/lib/CodeGen/CMakeLists.txt
index b95fe42..8524f1b 100644
--- a/backend/lib/CodeGen/CMakeLists.txt
+++ b/backend/lib/CodeGen/CMakeLists.txt
@@ -43,6 +43,7 @@ set(
   LivePhysRegs.cpp
   LiveRegUnits.cpp
   MachineMemoryFolding.cpp
+  MachinePeepholeOptimizer.cpp
   MachineMemOperand.cpp
   MachineMetadata.cpp
   MachineOperand.cpp
@@ -132,6 +133,7 @@ set(
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/LivePhysRegs.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/LiveRegUnits.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineMemoryFolding.h"
+  "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachinePeepholeOptimizer.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineMemOperand.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineMetadata.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineOperand.h"
diff --git a/backend/lib/CodeGen/MachineOptimizationPipeline.cpp b/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
index adec384..2badb78 100644
--- a/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
+++ b/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
@@ -14,6 +14,7 @@
 #include "rcc/CodeGen/MachineFinalization.h"
 #include "rcc/CodeGen/MachineLICM.h"
 #include "rcc/CodeGen/MachineMemoryFolding.h"
+#include "rcc/CodeGen/MachinePeepholeOptimizer.h"
 #include "rcc/CodeGen/MachinePseudoExpansion.h"
 #include "rcc/CodeGen/MachineScheduler.h"
 #include "rcc/CodeGen/MachineShrinkWrapping.h"
@@ -62,7 +63,13 @@ addSelectedSSAOptimizationPipeline(MachineFunctionPassManager &Pipeline,
   Added = addPass<MachineLICMPass>(Pipeline, MachineCodeMotionStage::SelectedSSA);
   if (!Added)
     return std::unexpected(Added.error());
-  return addPass<MachineSinkPass>(Pipeline, MachineCodeMotionStage::SelectedSSA);
+  Added = addPass<MachineSinkPass>(Pipeline, MachineCodeMotionStage::SelectedSSA);
+  if (!Added)
+    return std::unexpected(Added.error());
+  Added = addPass<MachinePeepholeOptimizerPass>(Pipeline, TargetInfo.InstrInfo);
+  if (!Added)
+    return std::unexpected(Added.error());
+  return addPass<MachineDCEPass>(Pipeline);
 }
 
 Expected<void> addPreRegisterAllocationOptimizationPipeline(
diff --git a/backend/lib/CodeGen/MachinePeepholeOptimizer.cpp b/backend/lib/CodeGen/MachinePeepholeOptimizer.cpp
new file mode 100644
index 0000000..f874699
--- /dev/null
+++ b/backend/lib/CodeGen/MachinePeepholeOptimizer.cpp
@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "rcc/CodeGen/MachinePeepholeOptimizer.h"
+
+#include "rcc/CodeGen/MachineBasicBlock.h"
+#include "rcc/CodeGen/MachineFunction.h"
+#include "rcc/CodeGen/MachineFunctionCloner.h"
+#include "rcc/CodeGen/MachineInstr.h"
+#include "rcc/CodeGen/TargetDescription.h"
+#include "rcc/CodeGen/TargetInstrInfo.h"
+
+#include <optional>
+
+namespace rcc {
+namespace {
+
+[[nodiscard]] const MachinePassContract &peepholeOptimizerContract() {
+  static const MachinePassContract Contract = [] {
+    MachinePassContract Value{
+        .Name = "peephole-opt",
+        .Requires = {MachineFunctionProperty::IsSSA},
+        .Sets = {},
+        .Clears = {},
+        .RequiredAnalyses = {},
+        .VerificationStage = PipelineStage::Selected,
+    };
+    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
+    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands,
+                             MachinePassChange::Registers};
+    return Value;
+  }();
+  return Contract;
+}
+
+} // namespace
+
+const MachinePassContract &MachinePeepholeOptimizerPass::contract() const noexcept {
+  return peepholeOptimizerContract();
+}
+
+Expected<MachinePassResult> MachinePeepholeOptimizerPass::run(MachineFunction &Function,
+                                                              const TargetDescription &Target,
+                                                              MachineFunctionAnalysisManager &) {
+  if (!Function.properties().containsAll(contract().Requires)) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "machine peephole optimizer requires SSA MachineIR");
+  }
+
+  auto Trial = cloneMachineFunction(
+      Function, Target,
+      MachineFunctionCloneOptions{.NewStableID = Function.stableID(), .NewName = std::nullopt});
+  if (!Trial) {
+    return std::unexpected(
+        withErrorContext(Trial.error(), "machine peephole transaction setup failed: "));
+  }
+  auto TrialResult = runImplementation(*Trial->ClonedFunction, Target);
+  if (!TrialResult)
+    return std::unexpected(TrialResult.error());
+
+  auto Result = runImplementation(Function, Target);
+  if (!Result) {
+    return std::unexpected(withErrorContext(
+        Result.error(), "machine peephole diverged after a successful transaction preflight: "));
+  }
+  return Result;
+}
+
+Expected<MachinePassResult>
+MachinePeepholeOptimizerPass::runImplementation(MachineFunction &Function,
+                                                const TargetDescription &Target) const {
+  MachinePassResult Result;
+  for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr; Block = Block->next()) {
+    for (MachineInstr *Instruction = Block->firstInstruction(); Instruction != nullptr;) {
+      MachineInstr *Next = Instruction->next();
+      if (!Instruction->isDebugInstr()) {
+        auto Changed = Info->optimizeCondBranch(*Instruction, Function, Target);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        Result.Changed |= *Changed;
+      }
+      Instruction = Next;
+    }
+  }
+  if (Result.Changed) {
+    Result.Changes = {MachinePassChange::Instructions, MachinePassChange::Operands,
+                      MachinePassChange::Registers};
+    Result.Preserved = PreservedAnalyses::none();
+  } else {
+    Result.Preserved = PreservedAnalyses::all();
+  }
+  return Result;
+}
+
+} // namespace rcc
diff --git a/backend/lib/Target/RISCV/RISCVInstrInfo.cpp b/backend/lib/Target/RISCV/RISCVInstrInfo.cpp
index c5cdd44..11387c1 100644
--- a/backend/lib/Target/RISCV/RISCVInstrInfo.cpp
+++ b/backend/lib/Target/RISCV/RISCVInstrInfo.cpp
@@ -595,6 +595,116 @@ selectLongBranchScratch(const MachineInstr &Instruction, const MachineBasicBlock
   return true;
 }
 
+[[nodiscard]] bool loadImmediate(const MachineInstr *Instruction, std::int64_t &Immediate) {
+  if (Instruction == nullptr ||
+      Instruction->opcode() != instructionOpcode(InstructionOpcode::ADDI) ||
+      !isPhysicalRegister(Instruction->operandAt(1), PhysicalRegisterID::X0)) {
+    return false;
+  }
+  const auto *Value = Instruction->operandAt(2) == nullptr
+                          ? nullptr
+                          : Instruction->operandAt(2)->getIf<std::int64_t>();
+  if (Value == nullptr)
+    return false;
+  Immediate = *Value;
+  return true;
+}
+
+[[nodiscard]] bool loadImmediateSource(const MachineRegisterInfo &Registers,
+                                       const MachineOperand &Operand, std::int64_t &Immediate) {
+  const Register *Reg = Operand.getIf<Register>();
+  if (Reg == nullptr)
+    return false;
+  if (const std::optional<PhysicalRegister> Physical = Reg->physicalRegister()) {
+    if (Physical->number() != static_cast<std::uint32_t>(PhysicalRegisterID::X0))
+      return false;
+    Immediate = 0;
+    return true;
+  }
+  return Reg->isVirtual() && loadImmediate(Registers.getVRegDef(*Reg), Immediate);
+}
+
+[[nodiscard]] std::optional<BranchConditionCode> branchCondition(Opcode Value) noexcept {
+  if (Value == instructionOpcode(InstructionOpcode::BEQ))
+    return BranchConditionCode::Equal;
+  if (Value == instructionOpcode(InstructionOpcode::BNE))
+    return BranchConditionCode::NotEqual;
+  if (Value == instructionOpcode(InstructionOpcode::BLT))
+    return BranchConditionCode::SignedLessThan;
+  if (Value == instructionOpcode(InstructionOpcode::BGE))
+    return BranchConditionCode::SignedGreaterEqual;
+  if (Value == instructionOpcode(InstructionOpcode::BLTU))
+    return BranchConditionCode::UnsignedLessThan;
+  if (Value == instructionOpcode(InstructionOpcode::BGEU))
+    return BranchConditionCode::UnsignedGreaterEqual;
+  return std::nullopt;
+}
+
+[[nodiscard]] bool evaluateBranch(BranchConditionCode Condition, std::int64_t Left,
+                                  std::int64_t Right) noexcept {
+  switch (Condition) {
+  case BranchConditionCode::Equal:
+    return Left == Right;
+  case BranchConditionCode::NotEqual:
+    return Left != Right;
+  case BranchConditionCode::SignedLessThan:
+    return Left < Right;
+  case BranchConditionCode::SignedGreaterEqual:
+    return Left >= Right;
+  case BranchConditionCode::UnsignedLessThan:
+    return static_cast<std::uint64_t>(Left) < static_cast<std::uint64_t>(Right);
+  case BranchConditionCode::UnsignedGreaterEqual:
+    return static_cast<std::uint64_t>(Left) >= static_cast<std::uint64_t>(Right);
+  }
+  return false;
+}
+
+[[nodiscard]] Expected<void> setRegisterUse(MachineOperand &Operand, Register Reg) {
+  auto Changed = Operand.setRegister(Reg);
+  if (!Changed)
+    return std::unexpected(Changed.error());
+  Changed = Operand.setRegisterFlags({RegisterOperandFlag::Use});
+  if (!Changed)
+    return std::unexpected(Changed.error());
+  Changed = Operand.setSubReg({});
+  if (!Changed)
+    return std::unexpected(Changed.error());
+  Operand.setTargetFlags(0);
+  return {};
+}
+
+[[nodiscard]] Expected<void> copyRegisterUse(MachineOperand &Destination,
+                                             const MachineOperand &Source) {
+  const Register *Reg = Source.getIf<Register>();
+  if (Reg == nullptr || !Source.isUse())
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V branch optimization requires register uses");
+  auto Changed = Destination.setRegister(*Reg);
+  if (!Changed)
+    return std::unexpected(Changed.error());
+  Changed = Destination.setRegisterFlags(Source.registerFlags());
+  if (!Changed)
+    return std::unexpected(Changed.error());
+  Changed = Destination.setSubReg(Source.subReg());
+  if (!Changed)
+    return std::unexpected(Changed.error());
+  Destination.setTargetFlags(Source.targetFlags());
+  return {};
+}
+
+[[nodiscard]] Expected<void> clearKillFlags(MachineRegisterInfo &Registers, Register Reg) {
+  for (MachineOperand &Operand : Registers.useOperands(Reg)) {
+    RegisterOperandFlags Flags = Operand.registerFlags();
+    if (!Flags.has(RegisterOperandFlag::Kill))
+      continue;
+    Flags.clear(RegisterOperandFlag::Kill);
+    auto Changed = Operand.setRegisterFlags(Flags);
+    if (!Changed)
+      return std::unexpected(Changed.error());
+  }
+  return {};
+}
+
 class RISCVInstrInfo final : public TargetInstrInfo {
 public:
   [[nodiscard]] Expected<MachineBranchAnalysis>
@@ -1085,6 +1195,116 @@ public:
     return false;
   }
 
+  [[nodiscard]] Expected<bool> optimizeCondBranch(MachineInstr &Instruction,
+                                                  MachineFunction &Function,
+                                                  const TargetDescription &) const override {
+    if (&Instruction.function() != &Function)
+      return makeUnexpected(ErrorCode::InvalidArgument,
+                            "RISC-V branch optimization references another function");
+    const auto Condition = branchCondition(Instruction.opcode());
+    if (!Condition)
+      return false;
+    MachineOperand *Left = Instruction.operandAt(0);
+    MachineOperand *Right = Instruction.operandAt(1);
+    if (Left == nullptr || Right == nullptr || Instruction.operandAt(2) == nullptr ||
+        Instruction.operandAt(2)->getIf<BasicBlockOperand>() == nullptr) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "RISC-V conditional branch has malformed operands");
+    }
+
+    MachineRegisterInfo &Registers = Function.registerInfo();
+    std::int64_t LeftConstant = 0;
+    std::int64_t RightConstant = 0;
+    if (loadImmediateSource(Registers, *Left, LeftConstant) &&
+        loadImmediateSource(Registers, *Right, RightConstant)) {
+      const InstructionOpcode NewOpcode = evaluateBranch(*Condition, LeftConstant, RightConstant)
+                                              ? InstructionOpcode::BEQ
+                                              : InstructionOpcode::BNE;
+      auto Zero = physicalRegister(PhysicalRegisterID::X0);
+      if (!Zero)
+        return std::unexpected(Zero.error());
+      auto Changed = Instruction.setOpcode(instructionOpcode(NewOpcode));
+      if (!Changed)
+        return std::unexpected(Changed.error());
+      Changed = setRegisterUse(*Left, Register::fromPhysical(*Zero));
+      if (!Changed)
+        return std::unexpected(Changed.error());
+      Changed = setRegisterUse(*Right, Register::fromPhysical(*Zero));
+      if (!Changed)
+        return std::unexpected(Changed.error());
+      return true;
+    }
+
+    const bool IsEquality =
+        *Condition == BranchConditionCode::Equal || *Condition == BranchConditionCode::NotEqual;
+    if (IsEquality)
+      return false;
+    const bool IsSigned = *Condition == BranchConditionCode::SignedLessThan ||
+                          *Condition == BranchConditionCode::SignedGreaterEqual;
+    const auto SearchConstant = [&Instruction](std::int64_t Value) -> Register {
+      for (MachineInstr *Candidate = Instruction.previous(); Candidate != nullptr;
+           Candidate = Candidate->previous()) {
+        std::int64_t CandidateValue = 0;
+        if (!loadImmediate(Candidate, CandidateValue) || CandidateValue != Value)
+          continue;
+        const MachineOperand *Definition = Candidate->operandAt(0);
+        const Register *Reg = Definition == nullptr ? nullptr : Definition->getIf<Register>();
+        if (Reg != nullptr && Reg->isVirtual())
+          return *Reg;
+      }
+      return {};
+    };
+    const std::optional<Opcode> NewOpcode = reversedOpcode(Instruction.opcode());
+    if (!NewOpcode)
+      return false;
+
+    const Register *LeftReg = Left->getIf<Register>();
+    if (loadImmediateSource(Registers, *Left, LeftConstant) && LeftConstant != 0 &&
+        LeftReg != nullptr && LeftReg->isVirtual() && Registers.hasOneUse(*LeftReg, false) &&
+        (IsSigned || LeftConstant != -1)) {
+      const Register Replacement = SearchConstant(LeftConstant + 1);
+      if (!Replacement.isNoRegister()) {
+        MachineOperand SavedRight = Right->cloneDetached();
+        auto Changed = Instruction.setOpcode(*NewOpcode);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        Changed = copyRegisterUse(*Left, SavedRight);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        Changed = setRegisterUse(*Right, Replacement);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        Changed = clearKillFlags(Registers, Replacement);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        return true;
+      }
+    }
+
+    const Register *RightReg = Right->getIf<Register>();
+    if (loadImmediateSource(Registers, *Right, RightConstant) && RightConstant != 0 &&
+        RightReg != nullptr && RightReg->isVirtual() && Registers.hasOneUse(*RightReg, false)) {
+      const Register Replacement = SearchConstant(RightConstant - 1);
+      if (!Replacement.isNoRegister()) {
+        MachineOperand SavedLeft = Left->cloneDetached();
+        auto Changed = Instruction.setOpcode(*NewOpcode);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        Changed = setRegisterUse(*Left, Replacement);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        Changed = copyRegisterUse(*Right, SavedLeft);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        Changed = clearKillFlags(Registers, Replacement);
+        if (!Changed)
+          return std::unexpected(Changed.error());
+        return true;
+      }
+    }
+    return false;
+  }
+
   [[nodiscard]] Expected<std::uint32_t>
   instructionSize(const MachineInstr &Instruction, const MachineFunction &,
                   const TargetDescription &Target) const override {