From e5fdceb9b869f010a2a75638fd38aaf1867f3008 Mon Sep 17 00:00:00 2001 From: Richard-Qin-X Date: Sun, 9 Aug 2026 04:40:17 +0800 Subject: [CodeGen] Complete register allocation infrastructure --- backend/lib/CodeGen/RegAllocPBQP.cpp | 281 +++++++++++++++++++++++++++++++++++ 1 file changed, 281 insertions(+) create mode 100644 backend/lib/CodeGen/RegAllocPBQP.cpp (limited to 'backend/lib/CodeGen/RegAllocPBQP.cpp') 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 +#include +#include +#include +#include +#include +#include +#include + +namespace rcc { +namespace { + +struct PBQPAllocationNode final { + Register VirtualReg; + pbqp::NodeID Node; + std::vector> Options; +}; + +[[nodiscard]] std::set> +copyRelatedRegisters(const MachineFunction &Function) { + std::set> 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(); + const Register *Source = Instruction->operands()[1].getIf(); + 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(pbqp::InfiniteCost - 1); + return static_cast(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 + allocateRound(const LiveIntervalsAnalysis::Result &Intervals, VirtRegMap &Mapping) override { + LiveRegMatrix FixedMatrix(Function, Target, Intervals); + pbqp::Graph Problem; + std::vector 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> 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(1, Interval->SpillWeight); + const std::uint64_t Scaled = + TargetSpillCost > std::numeric_limits::max() / Weight + ? std::numeric_limits::max() + : TargetSpillCost * Weight; + const std::uint64_t Adjusted = + Scaled == std::numeric_limits::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(Hint - Info->get().AllocationHints.begin()); + Cost = Cost > std::numeric_limits::max() - HintCost + ? std::numeric_limits::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(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()}, + .VerificationStage = PipelineStage::RegisterAllocated, + }; + Value.ThreadSafety = PassThreadSafety::ThreadSafe; + return Value; + }(); + return Contract; +} + +} // namespace + +const MachinePassContract &PBQPAllocatorPass::contract() const noexcept { + return pbqpContract(); +} + +Expected 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 -- cgit v1.2.3