[CodeGen] Implement machine optimization and scheduling

1 files changed, 514 insertions, 0 deletions

diff --git a/backend/lib/CodeGen/ScheduleDAG.cpp b/backend/lib/CodeGen/ScheduleDAG.cpp
new file mode 100644
index 0000000..8696f2e
--- /dev/null
+++ b/backend/lib/CodeGen/ScheduleDAG.cpp
@@ -0,0 +1,514 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "rcc/CodeGen/ScheduleDAG.h"
+
+#include "rcc/CodeGen/InstructionDescriptor.h"
+#include "rcc/CodeGen/MachineAnalyses.h"
+#include "rcc/CodeGen/MachineBasicBlock.h"
+#include "rcc/CodeGen/MachineFunction.h"
+#include "rcc/CodeGen/MachineInstr.h"
+#include "rcc/CodeGen/MachineMemOperand.h"
+#include "rcc/CodeGen/MachineOperand.h"
+#include "rcc/CodeGen/RegisterMask.h"
+#include "rcc/CodeGen/TargetDescription.h"
+#include "rcc/CodeGen/TargetRegisterInfo.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <optional>
+#include <ranges>
+#include <span>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace rcc {
+namespace {
+
+struct RegisterDependencyKey final {
+  bool IsUnit = false;
+  std::uint32_t Value = 0;
+
+  friend constexpr bool operator==(RegisterDependencyKey, RegisterDependencyKey) noexcept = default;
+};
+
+struct RegisterDependencyKeyHash final {
+  [[nodiscard]] std::size_t operator()(RegisterDependencyKey Key) const noexcept {
+    const std::uint64_t Value = (static_cast<std::uint64_t>(Key.IsUnit) << 32U) | Key.Value;
+    return std::hash<std::uint64_t>{}(Value);
+  }
+};
+
+struct RegisterDependencyState final {
+  struct Definition final {
+    std::size_t Unit = 0;
+    MachineInstr *Instruction = nullptr;
+    std::size_t OperandIndex = 0;
+  };
+
+  std::optional<Definition> LastDefinition;
+  std::vector<std::size_t> LastUses;
+};
+
+struct RegisterAccess final {
+  RegisterDependencyKey Key;
+  MachineInstr *Instruction = nullptr;
+  std::size_t OperandIndex = 0;
+};
+
+[[nodiscard]] bool keyLess(RegisterDependencyKey Left, RegisterDependencyKey Right) noexcept {
+  return Left.IsUnit == Right.IsUnit ? Left.Value < Right.Value : Left.IsUnit < Right.IsUnit;
+}
+
+[[nodiscard]] std::uint32_t saturatingAdd(std::uint32_t Left, std::uint32_t Right) noexcept {
+  return std::numeric_limits<std::uint32_t>::max() - Left < Right
+             ? std::numeric_limits<std::uint32_t>::max()
+             : Left + Right;
+}
+
+[[nodiscard]] std::uint64_t saturatingAdd(std::uint64_t Left, std::uint64_t Right) noexcept {
+  return std::numeric_limits<std::uint64_t>::max() - Left < Right
+             ? std::numeric_limits<std::uint64_t>::max()
+             : Left + Right;
+}
+
+[[nodiscard]] Expected<std::vector<RegisterDependencyKey>>
+operandKeys(const MachineOperand &Operand, const MachineFunction &Function,
+            const TargetDescription &Target, ScheduleDAGMode Mode) {
+  const Register *Reg = Operand.getIf<Register>();
+  if (Reg == nullptr || Reg->isNoRegister())
+    return std::vector<RegisterDependencyKey>{};
+  if (Reg->isVirtual()) {
+    if (Mode == ScheduleDAGMode::PostRegisterAllocation) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "post-RA scheduling region contains a virtual register");
+    }
+    if (!Function.registerInfo().isKnownVirtualRegister(*Reg)) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "scheduler encountered an unknown virtual register");
+    }
+    return std::vector<RegisterDependencyKey>{{false, Reg->opaqueValue()}};
+  }
+
+  std::optional<PhysicalRegister> Physical = Reg->physicalRegister();
+  if (!Physical)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "scheduler encountered an invalid physical register");
+  if (Operand.subReg().value() != 0) {
+    const std::optional<PhysicalRegister> SubReg =
+        Target.registers().subRegister(*Physical, Operand.subReg());
+    if (!SubReg) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "scheduler encountered an invalid physical subregister");
+    }
+    Physical = SubReg;
+  }
+  const PhysicalRegisterDesc *Description = Target.registers().physicalRegister(*Physical);
+  if (Description == nullptr)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "scheduler encountered an unknown physical register");
+  std::vector<RegisterDependencyKey> Result;
+  Result.reserve(Description->Units.size());
+  for (const RegisterUnitLane &Unit : Description->Units)
+    Result.push_back({true, Unit.Unit.value()});
+  std::ranges::sort(Result, {}, &RegisterDependencyKey::Value);
+  const auto [FirstDuplicate, End] = std::ranges::unique(Result);
+  Result.erase(FirstDuplicate, End);
+  return Result;
+}
+
+[[nodiscard]] Expected<std::vector<RegisterDependencyKey>>
+maskKeys(const RegisterMask &Mask, const TargetDescription &Target) {
+  std::vector<RegisterDependencyKey> Result;
+  for (const PhysicalRegisterDesc &Register : Target.registers().physicalRegisters()) {
+    if (!Mask.clobbers(Register.Reg))
+      continue;
+    for (const RegisterUnitLane &Unit : Register.Units)
+      Result.push_back({true, Unit.Unit.value()});
+  }
+  std::ranges::sort(Result, {}, &RegisterDependencyKey::Value);
+  const auto [FirstDuplicate, End] = std::ranges::unique(Result);
+  Result.erase(FirstDuplicate, End);
+  return Result;
+}
+
+[[nodiscard]] bool mayAlias(const MachineInstr &Earlier, const MachineInstr &Later,
+                            const MachineAliasAnalysis::Result &Aliases) noexcept {
+  if (Earlier.memOperands().empty() || Later.memOperands().empty())
+    return true;
+  for (const MachineMemOperand *Left : Earlier.memOperands()) {
+    if (Left == nullptr)
+      return true;
+    for (const MachineMemOperand *Right : Later.memOperands()) {
+      if (Right == nullptr || Aliases.alias(*Left, *Right) != MachineAliasResult::NoAlias)
+        return true;
+    }
+  }
+  return false;
+}
+
+[[nodiscard]] bool hasOrderedMemory(const MachineInstr &Instruction) noexcept {
+  return std::ranges::any_of(Instruction.memOperands(), [](const MachineMemOperand *Memory) {
+    return Memory == nullptr || Memory->isAtomic() || Memory->isVolatile();
+  });
+}
+
+[[nodiscard]] bool hasLoad(const MachineInstr &Instruction) noexcept {
+  return std::ranges::any_of(Instruction.memOperands(), [](const MachineMemOperand *Memory) {
+    return Memory == nullptr || Memory->isLoad();
+  });
+}
+
+[[nodiscard]] bool hasStore(const MachineInstr &Instruction) noexcept {
+  return std::ranges::any_of(Instruction.memOperands(), [](const MachineMemOperand *Memory) {
+    return Memory == nullptr || Memory->isStore();
+  });
+}
+
+} // namespace
+
+Expected<ScheduleDAG> ScheduleDAG::build(MachineBasicBlock &Block,
+                                         std::span<MachineInstr *const> Instructions,
+                                         const MachineFunction &Function,
+                                         const TargetDescription &Target,
+                                         const TargetSchedulingModel &Model, ScheduleDAGMode Mode) {
+  if (&Block.function() != &Function)
+    return makeUnexpected(ErrorCode::InvalidArgument,
+                          "schedule DAG block belongs to another function");
+  ScheduleDAG Result(Block);
+  auto Units = Result.buildUnits(Instructions, Target, Model);
+  if (!Units)
+    return std::unexpected(Units.error());
+  auto Registers = Result.buildRegisterDependencies(Function, Target, Model, Mode);
+  if (!Registers)
+    return std::unexpected(Registers.error());
+  auto Memory = Result.buildMemoryDependencies(Target);
+  if (!Memory)
+    return std::unexpected(Memory.error());
+  auto Order = Result.buildOrderDependencies(Target);
+  if (!Order)
+    return std::unexpected(Order.error());
+  auto Mutated = Model.hooks().mutateScheduleDAG(Result, Function, Target, Mode);
+  if (!Mutated)
+    return std::unexpected(Mutated.error());
+  Result.computeCriticalPaths();
+  return Result;
+}
+
+Expected<void> ScheduleDAG::buildUnits(std::span<MachineInstr *const> Instructions,
+                                       const TargetDescription &Target,
+                                       const TargetSchedulingModel &Model) {
+  MachineInstr *ExpectedInstruction = Instructions.empty() ? nullptr : Instructions.front();
+  for (std::size_t Index = 0; Index < Instructions.size();) {
+    MachineInstr *Instruction = Instructions[Index];
+    if (Instruction == nullptr || Instruction->parent() != Block)
+      return makeUnexpected(ErrorCode::InvalidArgument,
+                            "schedule region contains a detached or foreign instruction");
+    if (Instruction != ExpectedInstruction) {
+      return makeUnexpected(ErrorCode::InvalidArgument,
+                            "schedule region is not a contiguous instruction sequence");
+    }
+    if (Instruction->flags().has(MachineInstrFlag::BundledWithPred)) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "schedule region begins in the middle of an instruction bundle");
+    }
+
+    ScheduleUnit Unit;
+    Unit.ID = Units.size();
+    Unit.StableID = Instruction->stableID();
+    std::unordered_map<ProcessorResourceID, std::uint32_t> ResourceCycles;
+    for (;;) {
+      const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction->opcode());
+      if (Descriptor == nullptr) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "scheduler encountered an unknown instruction descriptor");
+      }
+      const SchedulingClassModelDesc *Class = Model.schedulingClass(Descriptor->SchedulingClass);
+      if (Class == nullptr) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "instruction references an unknown scheduling class");
+      }
+      Unit.MicroOps = saturatingAdd(Unit.MicroOps, Class->MicroOps);
+      Unit.Latency = std::max(Unit.Latency, Class->Latency);
+      for (const SchedulingResourceUse &Use : Class->Resources) {
+        ResourceCycles[Use.Resource] = saturatingAdd(ResourceCycles[Use.Resource], Use.Cycles);
+      }
+      Unit.Instructions.push_back(Instruction);
+      const bool HasSuccessor = Instruction->flags().has(MachineInstrFlag::BundledWithSucc);
+      ++Index;
+      if (!HasSuccessor)
+        break;
+      if (Index >= Instructions.size()) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "schedule region ends in the middle of an instruction bundle");
+      }
+      MachineInstr *Next = Instructions[Index];
+      if (Next == nullptr || Next->parent() != Block ||
+          !Next->flags().has(MachineInstrFlag::BundledWithPred) || Instruction->next() != Next) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "instruction bundle links are not reciprocal and contiguous");
+      }
+      Instruction = Next;
+    }
+    Unit.Resources.reserve(ResourceCycles.size());
+    for (const auto &[Resource, Cycles] : ResourceCycles)
+      Unit.Resources.push_back({Resource, Cycles});
+    std::ranges::sort(Unit.Resources, {}, &SchedulingResourceUse::Resource);
+    ExpectedInstruction = Unit.Instructions.back()->next();
+    Units.push_back(std::move(Unit));
+  }
+  return {};
+}
+
+Expected<void> ScheduleDAG::buildRegisterDependencies(const MachineFunction &Function,
+                                                      const TargetDescription &Target,
+                                                      const TargetSchedulingModel &Model,
+                                                      ScheduleDAGMode Mode) {
+  std::unordered_map<RegisterDependencyKey, RegisterDependencyState, RegisterDependencyKeyHash>
+      States;
+  for (ScheduleUnit &Unit : Units) {
+    std::vector<RegisterAccess> Uses;
+    std::vector<RegisterAccess> Definitions;
+    for (MachineInstr *Instruction : Unit.Instructions) {
+      for (std::size_t OperandIndex = 0; OperandIndex < Instruction->operandCount();
+           ++OperandIndex) {
+        const MachineOperand *Operand = Instruction->operandAt(OperandIndex);
+        if (Operand == nullptr)
+          continue;
+        if (const auto *Mask = Operand->getIf<RegisterMaskOperand>()) {
+          if (!Mask->Mask)
+            return makeUnexpected(ErrorCode::VerificationError,
+                                  "scheduler encountered a null register mask");
+          auto Keys = maskKeys(*Mask->Mask, Target);
+          if (!Keys)
+            return std::unexpected(Keys.error());
+          for (const RegisterDependencyKey Key : *Keys)
+            Definitions.push_back({Key, Instruction, OperandIndex});
+          continue;
+        }
+        if (const auto *Mask = Operand->getIf<LiveOutRegisterMaskOperand>()) {
+          if (!Mask->Mask)
+            return makeUnexpected(ErrorCode::VerificationError,
+                                  "scheduler encountered a null live-out register mask");
+          for (const PhysicalRegisterDesc &Register : Target.registers().physicalRegisters()) {
+            if (!Mask->Mask->preserves(Register.Reg))
+              continue;
+            for (const RegisterUnitLane &RegisterUnit : Register.Units)
+              Uses.push_back({{true, RegisterUnit.Unit.value()}, Instruction, OperandIndex});
+          }
+          continue;
+        }
+        if (!Operand->isRegister() || Operand->isDebugUse() ||
+            Operand->registerFlags().has(RegisterOperandFlag::Undef)) {
+          continue;
+        }
+        auto Keys = operandKeys(*Operand, Function, Target, Mode);
+        if (!Keys)
+          return std::unexpected(Keys.error());
+        if (Operand->isUse()) {
+          for (const RegisterDependencyKey Key : *Keys)
+            Uses.push_back({Key, Instruction, OperandIndex});
+        }
+        if (Operand->isDef()) {
+          for (const RegisterDependencyKey Key : *Keys)
+            Definitions.push_back({Key, Instruction, OperandIndex});
+        }
+      }
+    }
+    std::ranges::stable_sort(Uses, [](const RegisterAccess &Left, const RegisterAccess &Right) {
+      return keyLess(Left.Key, Right.Key);
+    });
+    std::ranges::stable_sort(Definitions,
+                             [](const RegisterAccess &Left, const RegisterAccess &Right) {
+                               return keyLess(Left.Key, Right.Key);
+                             });
+
+    for (std::size_t Begin = 0; Begin < Uses.size();) {
+      std::size_t End = Begin + 1;
+      while (End < Uses.size() && Uses[End].Key == Uses[Begin].Key)
+        ++End;
+      RegisterDependencyState &State = States[Uses[Begin].Key];
+      if (State.LastDefinition) {
+        std::uint32_t Latency = 0;
+        for (std::size_t Index = Begin; Index < End; ++Index) {
+          const InstructionDescriptor *Producer =
+              Target.instructionDescriptor(State.LastDefinition->Instruction->opcode());
+          const InstructionDescriptor *Consumer =
+              Target.instructionDescriptor(Uses[Index].Instruction->opcode());
+          if (Producer == nullptr || Consumer == nullptr) {
+            return makeUnexpected(ErrorCode::VerificationError,
+                                  "scheduler latency lookup encountered an unknown instruction");
+          }
+          auto OperandLatency =
+              Model.operandLatency(Producer->SchedulingClass, State.LastDefinition->OperandIndex,
+                                   Consumer->SchedulingClass, Uses[Index].OperandIndex);
+          if (!OperandLatency)
+            return std::unexpected(OperandLatency.error());
+          Latency = std::max(Latency, *OperandLatency);
+        }
+        addDependency(State.LastDefinition->Unit, Unit.ID,
+                      Uses[Begin].Key.IsUnit ? ScheduleDependencyKind::RegisterUnit
+                                             : ScheduleDependencyKind::Data,
+                      Latency);
+      }
+      if (State.LastUses.empty() || State.LastUses.back() != Unit.ID)
+        State.LastUses.push_back(Unit.ID);
+      Begin = End;
+    }
+
+    for (std::size_t Begin = 0; Begin < Definitions.size();) {
+      std::size_t End = Begin + 1;
+      while (End < Definitions.size() && Definitions[End].Key == Definitions[Begin].Key) {
+        ++End;
+      }
+      RegisterDependencyState &State = States[Definitions[Begin].Key];
+      if (State.LastDefinition) {
+        addDependency(State.LastDefinition->Unit, Unit.ID, ScheduleDependencyKind::Output, 0);
+      }
+      for (const std::size_t Use : State.LastUses)
+        addDependency(Use, Unit.ID, ScheduleDependencyKind::Anti, 0);
+      State.LastUses.clear();
+      const RegisterAccess &Last = Definitions[End - 1];
+      State.LastDefinition =
+          RegisterDependencyState::Definition{Unit.ID, Last.Instruction, Last.OperandIndex};
+      Begin = End;
+    }
+  }
+  return {};
+}
+
+Expected<void> ScheduleDAG::addArtificialDependency(std::size_t Source, std::size_t Destination,
+                                                    std::uint32_t Latency) {
+  if (Source >= Units.size() || Destination >= Units.size() || Source >= Destination) {
+    return makeUnexpected(
+        ErrorCode::InvalidArgument,
+        "target schedule mutation requires a forward dependency between known units");
+  }
+  addDependency(Source, Destination, ScheduleDependencyKind::Artificial, Latency);
+  return {};
+}
+
+Expected<void> ScheduleDAG::buildMemoryDependencies(const TargetDescription &Target) {
+  MachineAliasAnalysis::Result Aliases;
+  for (std::size_t LaterIndex = 0; LaterIndex < Units.size(); ++LaterIndex) {
+    for (std::size_t EarlierIndex = 0; EarlierIndex < LaterIndex; ++EarlierIndex) {
+      bool NeedsEdge = false;
+      for (MachineInstr *Earlier : Units[EarlierIndex].Instructions) {
+        const InstructionDescriptor *EarlierDescriptor =
+            Target.instructionDescriptor(Earlier->opcode());
+        if (EarlierDescriptor == nullptr)
+          return makeUnexpected(ErrorCode::VerificationError,
+                                "scheduler encountered an unknown instruction descriptor");
+        const bool EarlierMemory = EarlierDescriptor->Flags.has(InstructionFlag::MayLoad) ||
+                                   EarlierDescriptor->Flags.has(InstructionFlag::MayStore) ||
+                                   !Earlier->memOperands().empty();
+        if (!EarlierMemory)
+          continue;
+        for (MachineInstr *Later : Units[LaterIndex].Instructions) {
+          const InstructionDescriptor *LaterDescriptor =
+              Target.instructionDescriptor(Later->opcode());
+          if (LaterDescriptor == nullptr)
+            return makeUnexpected(ErrorCode::VerificationError,
+                                  "scheduler encountered an unknown instruction descriptor");
+          const bool LaterMemory = LaterDescriptor->Flags.has(InstructionFlag::MayLoad) ||
+                                   LaterDescriptor->Flags.has(InstructionFlag::MayStore) ||
+                                   !Later->memOperands().empty();
+          if (!LaterMemory)
+            continue;
+          if (hasOrderedMemory(*Earlier) || hasOrderedMemory(*Later) ||
+              ((hasStore(*Earlier) || hasStore(*Later)) && mayAlias(*Earlier, *Later, Aliases)) ||
+              (Earlier->memOperands().empty() || Later->memOperands().empty()) ||
+              (!hasLoad(*Earlier) && !hasStore(*Earlier)) ||
+              (!hasLoad(*Later) && !hasStore(*Later))) {
+            NeedsEdge = true;
+            break;
+          }
+        }
+        if (NeedsEdge)
+          break;
+      }
+      if (NeedsEdge)
+        addDependency(EarlierIndex, LaterIndex, ScheduleDependencyKind::Memory, 0);
+    }
+  }
+  return {};
+}
+
+Expected<void> ScheduleDAG::buildOrderDependencies(const TargetDescription &Target) {
+  for (std::size_t Barrier = 0; Barrier < Units.size(); ++Barrier) {
+    bool IsBarrier = false;
+    for (const MachineInstr *Instruction : Units[Barrier].Instructions) {
+      const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction->opcode());
+      if (Descriptor == nullptr) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "scheduler encountered an unknown instruction descriptor");
+      }
+      const InstructionFlags Flags = Descriptor->Flags;
+      IsBarrier |= Flags.has(InstructionFlag::Terminator) || Flags.has(InstructionFlag::Branch) ||
+                   Flags.has(InstructionFlag::Call) || Flags.has(InstructionFlag::Return) ||
+                   Flags.has(InstructionFlag::Barrier) ||
+                   Flags.has(InstructionFlag::HasSideEffects) ||
+                   Flags.has(InstructionFlag::UnmodeledSideEffects) ||
+                   Flags.has(InstructionFlag::Convergent) || Flags.has(InstructionFlag::Phi) ||
+                   Flags.has(InstructionFlag::Meta);
+    }
+    if (!IsBarrier)
+      continue;
+    for (std::size_t Earlier = 0; Earlier < Barrier; ++Earlier)
+      addDependency(Earlier, Barrier, ScheduleDependencyKind::Barrier, 0);
+    for (std::size_t Later = Barrier + 1; Later < Units.size(); ++Later)
+      addDependency(Barrier, Later, ScheduleDependencyKind::Barrier, 0);
+  }
+  return {};
+}
+
+void ScheduleDAG::addDependency(std::size_t Source, std::size_t Destination,
+                                ScheduleDependencyKind Kind, std::uint32_t Latency) {
+  if (Source >= Units.size() || Destination >= Units.size() || Source == Destination)
+    return;
+  const auto Existing = std::ranges::find_if(
+      Units[Source].Successors, [Destination, Kind](const ScheduleDependency &Dependency) {
+        return Dependency.Destination == Destination && Dependency.Kind == Kind;
+      });
+  if (Existing != Units[Source].Successors.end()) {
+    Existing->Latency = std::max(Existing->Latency, Latency);
+    const auto Predecessor = std::ranges::find_if(
+        Units[Destination].Predecessors, [Source, Kind](const ScheduleDependency &Dependency) {
+          return Dependency.Source == Source && Dependency.Kind == Kind;
+        });
+    if (Predecessor != Units[Destination].Predecessors.end())
+      Predecessor->Latency = Existing->Latency;
+    return;
+  }
+  const ScheduleDependency Dependency{Source, Destination, Kind, Latency};
+  Units[Source].Successors.push_back(Dependency);
+  Units[Destination].Predecessors.push_back(Dependency);
+}
+
+void ScheduleDAG::computeCriticalPaths() {
+  for (ScheduleUnit &Unit : Units) {
+    Unit.Depth = 0;
+    for (const ScheduleDependency &Dependency : Unit.Predecessors) {
+      Unit.Depth =
+          std::max(Unit.Depth, saturatingAdd(Units[Dependency.Source].Depth,
+                                             static_cast<std::uint64_t>(Dependency.Latency)));
+    }
+  }
+  for (auto Position = Units.rbegin(); Position != Units.rend(); ++Position) {
+    Position->Height = Position->Latency;
+    for (const ScheduleDependency &Dependency : Position->Successors) {
+      Position->Height =
+          std::max(Position->Height, saturatingAdd(Units[Dependency.Destination].Height,
+                                                   static_cast<std::uint64_t>(Dependency.Latency)));
+    }
+  }
+  for (ScheduleUnit &Unit : Units) {
+    std::ranges::sort(Unit.Predecessors, {}, &ScheduleDependency::Source);
+    std::ranges::sort(Unit.Successors, {}, &ScheduleDependency::Destination);
+  }
+}
+
+} // namespace rcc