[CodeGen] Complete register allocation infrastructure

1 files changed, 281 insertions, 0 deletions

diff --git a/backend/lib/CodeGen/RegAllocPBQP.cpp b/backend/lib/CodeGen/RegAllocPBQP.cpp
new file mode 100644
index 0000000..5e56728
--- /dev/null
+++ b/backend/lib/CodeGen/RegAllocPBQP.cpp
@@ -0,0 +1,281 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "rcc/CodeGen/RegAllocPBQP.h"
+
+#include "RegAllocBase.h"
+#include "rcc/CodeGen/GenericOpcodes.h"
+#include "rcc/CodeGen/LiveRegMatrix.h"
+#include "rcc/CodeGen/MachineBasicBlock.h"
+#include "rcc/CodeGen/MachineFunction.h"
+#include "rcc/CodeGen/MachineInstr.h"
+#include "rcc/CodeGen/MachineLiveness.h"
+#include "rcc/CodeGen/MachineOperand.h"
+#include "rcc/CodeGen/PBQP/Solver.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+#include <optional>
+#include <ranges>
+#include <set>
+#include <utility>
+#include <vector>
+
+namespace rcc {
+namespace {
+
+struct PBQPAllocationNode final {
+  Register VirtualReg;
+  pbqp::NodeID Node;
+  std::vector<std::optional<PhysicalRegister>> Options;
+};
+
+[[nodiscard]] std::set<std::pair<std::uint32_t, std::uint32_t>>
+copyRelatedRegisters(const MachineFunction &Function) {
+  std::set<std::pair<std::uint32_t, std::uint32_t>> Result;
+  for (const MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr;
+       Block = Block->next()) {
+    for (const MachineInstr *Instruction = Block->firstInstruction(); Instruction != nullptr;
+         Instruction = Instruction->next()) {
+      if (Instruction->opcode() != opcode(GenericOpcode::COPY) || Instruction->operandCount() < 2)
+        continue;
+      const Register *Destination = Instruction->operands()[0].getIf<Register>();
+      const Register *Source = Instruction->operands()[1].getIf<Register>();
+      if (Destination == nullptr || Source == nullptr || !Destination->isVirtual() ||
+          !Source->isVirtual() || *Destination == *Source)
+        continue;
+      std::uint32_t Left = Destination->opaqueValue();
+      std::uint32_t Right = Source->opaqueValue();
+      if (Right < Left)
+        std::swap(Left, Right);
+      Result.emplace(Left, Right);
+    }
+  }
+  return Result;
+}
+
+[[nodiscard]] pbqp::Cost boundedCost(std::uint64_t Value) noexcept {
+  const std::uint64_t Limit = static_cast<std::uint64_t>(pbqp::InfiniteCost - 1);
+  return static_cast<pbqp::Cost>(std::min(Value, Limit));
+}
+
+class PBQPAllocatorEngine final : public RegAllocBase {
+public:
+  PBQPAllocatorEngine(MachineFunction &Function, const TargetDescription &Target,
+                      const TargetRegisterAllocationInfo &TargetInfo,
+                      MachineAnalysisManager &Analyses, RegisterAllocationOptions Options)
+      : RegAllocBase(Function, Target, TargetInfo, Analyses, Options) {}
+
+private:
+  [[nodiscard]] Expected<AllocationRoundResult>
+  allocateRound(const LiveIntervalsAnalysis::Result &Intervals, VirtRegMap &Mapping) override {
+    LiveRegMatrix FixedMatrix(Function, Target, Intervals);
+    pbqp::Graph Problem;
+    std::vector<PBQPAllocationNode> Nodes;
+    for (const LiveInterval *Interval : orderedAllocationIntervals(Intervals, Function)) {
+      const Register VirtualReg = Interval->Reg;
+      const bool TerminatorDefined =
+          regalloc_detail::hasTerminatorDefinition(Function, Target, VirtualReg);
+      if (TerminatorDefined) {
+        auto Marked = Mapping.setSpillable(VirtualReg, false);
+        if (!Marked)
+          return std::unexpected(Marked.error());
+      }
+      auto ClassID = regalloc_detail::registerClassFor(Function, VirtualReg);
+      if (!ClassID)
+        return std::unexpected(ClassID.error());
+      auto Candidates = regalloc_detail::allocationOrder(Function, Target, VirtualReg);
+      if (!Candidates)
+        return std::unexpected(Candidates.error());
+      std::erase_if(*Candidates, [&](PhysicalRegister Candidate) {
+        return FixedMatrix.checkInterference(VirtualReg, Candidate).Kind !=
+               RegisterInterferenceKind::Free;
+      });
+
+      std::vector<std::optional<PhysicalRegister>> NodeOptions;
+      NodeOptions.reserve(Candidates->size() + 1U);
+      NodeOptions.push_back(std::nullopt);
+      for (const PhysicalRegister Candidate : *Candidates)
+        NodeOptions.push_back(Candidate);
+      pbqp::CostVector Costs(NodeOptions.size(), 0);
+      if (!Mapping.isSpillable(VirtualReg)) {
+        Costs[0] = pbqp::InfiniteCost;
+      } else {
+        const std::uint64_t TargetSpillCost = TargetInfo.spillCost(Function, *Interval, *ClassID);
+        const std::uint64_t Weight = std::max<std::uint64_t>(1, Interval->SpillWeight);
+        const std::uint64_t Scaled =
+            TargetSpillCost > std::numeric_limits<std::uint64_t>::max() / Weight
+                ? std::numeric_limits<std::uint64_t>::max()
+                : TargetSpillCost * Weight;
+        const std::uint64_t Adjusted =
+            Scaled == std::numeric_limits<std::uint64_t>::max() ? Scaled : Scaled + 1U;
+        Costs[0] = boundedCost(Adjusted);
+      }
+      auto Info = Function.registerInfo().virtualRegisterInfo(VirtualReg);
+      if (!Info)
+        return std::unexpected(Info.error());
+      for (std::size_t Index = 1; Index < NodeOptions.size(); ++Index) {
+        const auto CandidateOption = NodeOptions[Index];
+        if (!CandidateOption) {
+          return makeUnexpected(ErrorCode::InternalError,
+                                "PBQP allocation node has an unexpected spill option");
+        }
+        const PhysicalRegister Candidate = CandidateOption.value();
+        std::uint64_t Cost =
+            TargetInfo.physicalRegisterCost(Function, VirtualReg, Candidate, *ClassID, true);
+        const auto Hint = std::ranges::find(Info->get().AllocationHints, Candidate);
+        const std::uint64_t HintCost =
+            Hint == Info->get().AllocationHints.end()
+                ? 2U
+                : static_cast<std::uint64_t>(Hint - Info->get().AllocationHints.begin());
+        Cost = Cost > std::numeric_limits<std::uint64_t>::max() - HintCost
+                   ? std::numeric_limits<std::uint64_t>::max()
+                   : Cost + HintCost;
+        Costs[Index] = boundedCost(Cost);
+      }
+      auto NodeID = Problem.addNode(std::move(Costs), VirtualReg.opaqueValue());
+      if (!NodeID)
+        return std::unexpected(NodeID.error());
+      Nodes.push_back(PBQPAllocationNode{VirtualReg, *NodeID, std::move(NodeOptions)});
+    }
+
+    const auto CopyPairs = copyRelatedRegisters(Function);
+    for (std::size_t LeftIndex = 0; LeftIndex < Nodes.size(); ++LeftIndex) {
+      const PBQPAllocationNode &Left = Nodes[LeftIndex];
+      const LiveInterval *LeftInterval = Intervals.interval(Left.VirtualReg);
+      for (std::size_t RightIndex = LeftIndex + 1U; RightIndex < Nodes.size(); ++RightIndex) {
+        const PBQPAllocationNode &Right = Nodes[RightIndex];
+        const LiveInterval *RightInterval = Intervals.interval(Right.VirtualReg);
+        if (LeftInterval == nullptr || RightInterval == nullptr)
+          continue;
+        std::uint32_t LeftKey = Left.VirtualReg.opaqueValue();
+        std::uint32_t RightKey = Right.VirtualReg.opaqueValue();
+        if (RightKey < LeftKey)
+          std::swap(LeftKey, RightKey);
+        const bool CopyRelated =
+            Options.EnableCopyCoalescing && CopyPairs.contains({LeftKey, RightKey});
+        pbqp::CostMatrix Costs(Left.Options.size(), Right.Options.size(), 0);
+        bool HasCost = false;
+        for (std::size_t LeftOption = 1; LeftOption < Left.Options.size(); ++LeftOption) {
+          for (std::size_t RightOption = 1; RightOption < Right.Options.size(); ++RightOption) {
+            const auto LeftRegister = Left.Options[LeftOption];
+            const auto RightRegister = Right.Options[RightOption];
+            if (!LeftRegister || !RightRegister) {
+              return makeUnexpected(ErrorCode::InternalError,
+                                    "PBQP edge has an unexpected spill option");
+            }
+            const PhysicalRegister LeftReg = LeftRegister.value();
+            const PhysicalRegister RightReg = RightRegister.value();
+            if (regalloc_detail::intervalsOverlapOnPhysicalRegisters(
+                    *LeftInterval, LeftReg, *RightInterval, RightReg, Target.registers())) {
+              Costs.at(LeftOption, RightOption) = pbqp::InfiniteCost;
+              HasCost = true;
+            } else if (CopyRelated && LeftReg == RightReg) {
+              Costs.at(LeftOption, RightOption) = -1;
+              HasCost = true;
+            }
+          }
+        }
+        if (!HasCost)
+          continue;
+        const std::uint64_t StableKey =
+            (static_cast<std::uint64_t>(Left.VirtualReg.opaqueValue()) << 32U) |
+            Right.VirtualReg.opaqueValue();
+        auto Edge = Problem.addEdge(Left.Node, Right.Node, std::move(Costs), StableKey);
+        if (!Edge)
+          return std::unexpected(Edge.error());
+      }
+    }
+
+    pbqp::DeterministicSolver Solver;
+    auto Solution = Solver.solve(Problem);
+    if (!Solution)
+      return std::unexpected(Solution.error());
+
+    LiveRegMatrix AssignedMatrix(Function, Target, Intervals);
+    AllocationRoundResult Result;
+    for (const PBQPAllocationNode &Node : Nodes) {
+      const auto Selection = Solution->selection(Node.Node);
+      if (!Selection)
+        return makeUnexpected(ErrorCode::InternalError,
+                              "PBQP solver returned an invalid selection");
+      const std::size_t SelectedOption = Selection.value();
+      if (SelectedOption >= Node.Options.size()) {
+        return makeUnexpected(ErrorCode::InternalError,
+                              "PBQP solver selected an out-of-range option");
+      }
+      const auto AssignedOption = Node.Options[SelectedOption];
+      if (!AssignedOption) {
+        Result.Spilled.push_back(Node.VirtualReg);
+        continue;
+      }
+      const PhysicalRegister Assigned = AssignedOption.value();
+      auto MatrixAssignment = AssignedMatrix.assign(Node.VirtualReg, Assigned);
+      if (!MatrixAssignment)
+        return makeUnexpected(ErrorCode::InternalError,
+                              "PBQP solution violates its encoded interference constraints");
+      auto MapAssignment = Mapping.assign(Node.VirtualReg, Assigned);
+      if (!MapAssignment)
+        return std::unexpected(MapAssignment.error());
+    }
+    return Result;
+  }
+};
+
+[[nodiscard]] const MachinePassContract &pbqpContract() {
+  static const MachinePassContract Contract = [] {
+    MachinePassContract Value{
+        .Name = "PBQPAllocator",
+        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs,
+                     MachineFunctionProperty::TiedOpsRewritten,
+                     MachineFunctionProperty::TracksLiveness},
+        .Sets = {MachineFunctionProperty::NoVRegs},
+        .Clears = {},
+        .RequiredAnalyses = {MachineAnalysisManager::id<LiveIntervalsAnalysis>()},
+        .VerificationStage = PipelineStage::RegisterAllocated,
+    };
+    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
+    return Value;
+  }();
+  return Contract;
+}
+
+} // namespace
+
+const MachinePassContract &PBQPAllocatorPass::contract() const noexcept {
+  return pbqpContract();
+}
+
+Expected<MachinePassResult> PBQPAllocatorPass::run(MachineFunction &Function,
+                                                   const TargetDescription &Target,
+                                                   MachineAnalysisManager &Analyses) {
+  if (Info == nullptr)
+    return makeUnexpected(ErrorCode::InvalidArgument, "PBQPAllocator requires target behavior");
+  auto Required = regalloc_detail::requireAllocationProperties(Function, contract());
+  if (!Required)
+    return std::unexpected(Required.error());
+  auto Trial = regalloc_detail::cloneForAllocation(Function, Target);
+  if (!Trial)
+    return std::unexpected(Trial.error());
+  MachineAnalysisManager TrialAnalyses;
+  auto Registered = registerStage11MachineAnalyses(TrialAnalyses);
+  if (!Registered)
+    return std::unexpected(Registered.error());
+  PBQPAllocatorEngine TrialAllocator(*Trial->ClonedFunction, Target, *Info, TrialAnalyses, Options);
+  auto TrialRun = TrialAllocator.run();
+  if (!TrialRun)
+    return std::unexpected(TrialRun.error());
+  auto TrialVerification =
+      regalloc_detail::verifyAllocationTrial(*Trial->ClonedFunction, Target, contract());
+  if (!TrialVerification)
+    return std::unexpected(TrialVerification.error());
+  PBQPAllocatorEngine Allocator(Function, Target, *Info, Analyses, Options);
+  auto Allocated = Allocator.run();
+  if (!Allocated) {
+    return std::unexpected(regalloc_detail::withPrefix(
+        Allocated.error(), "PBQPAllocator diverged after a successful transaction preflight: "));
+  }
+  return regalloc_detail::passResult(*Allocated);
+}
+
+} // namespace rcc