1 files changed, 248 insertions, 0 deletions
diff --git a/backend/lib/CodeGen/MachineMemoryFolding.cpp b/backend/lib/CodeGen/MachineMemoryFolding.cpp
new file mode 100644
index 0000000..0e15146
--- /dev/null
+++ b/backend/lib/CodeGen/MachineMemoryFolding.cpp
@@ -0,0 +1,248 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "rcc/CodeGen/MachineMemoryFolding.h"
+
+#include "rcc/CodeGen/InstructionDescriptor.h"
+#include "rcc/CodeGen/MachineBasicBlock.h"
+#include "rcc/CodeGen/MachineFunction.h"
+#include "rcc/CodeGen/MachineFunctionCloner.h"
+#include "rcc/CodeGen/MachineInstr.h"
+#include "rcc/CodeGen/MachineMemOperand.h"
+#include "rcc/CodeGen/MachineOperand.h"
+#include "rcc/CodeGen/MachineRegisterInfo.h"
+#include "rcc/CodeGen/TargetDescription.h"
+#include "rcc/CodeGen/TargetMemoryFolding.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <optional>
+#include <ranges>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace rcc {
+namespace {
+
+[[nodiscard]] const MachinePassContract &memoryFoldingContract() {
+  static const MachinePassContract Contract = [] {
+    MachinePassContract Value{
+        .Name = "machine-memory-folding",
+        .Requires = {MachineFunctionProperty::IsSSA, MachineFunctionProperty::Selected},
+        .Sets = {},
+        .Clears = {},
+        .RequiredAnalyses = {},
+        .VerificationStage = PipelineStage::Selected,
+    };
+    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
+    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands,
+                             MachinePassChange::Registers};
+    return Value;
+  }();
+  return Contract;
+}
+
+[[nodiscard]] Error prefixedError(const Error &Failure, std::string Prefix) {
+  Prefix.append(Failure.message());
+  return Error(Failure.code(), std::move(Prefix));
+}
+
+[[nodiscard]] bool hasProtectedAttachment(const MachineInstr &Instruction) noexcept {
+  return Instruction.debugInstrNumber() != 0 || Instruction.preInstrSymbol().has_value() ||
+         Instruction.postInstrSymbol().has_value() || Instruction.heapAllocMarker().has_value() ||
+         Instruction.pcSections().has_value() || Instruction.mmra().has_value() ||
+         Instruction.cfiType().has_value() || Instruction.deactivationSymbol().has_value() ||
+         Instruction.flags().has(MachineInstrFlag::BundledWithPred) ||
+         Instruction.flags().has(MachineInstrFlag::BundledWithSucc);
+}
+
+struct FoldCandidate final {
+  MachineInstr *MemoryInstruction = nullptr;
+  MachineInstr *Consumer = nullptr;
+  Register DefinedRegister;
+  std::size_t ConsumerOperand = 0;
+};
+
+[[nodiscard]] std::optional<FoldCandidate> findCandidate(MachineInstr &Instruction,
+                                                         MachineFunction &Function,
+                                                         const TargetDescription &Target) {
+  const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction.opcode());
+  if (Descriptor == nullptr || !Descriptor->Flags.has(InstructionFlag::MayLoad) ||
+      Descriptor->Flags.has(InstructionFlag::MayStore) ||
+      Descriptor->Flags.has(InstructionFlag::Call) ||
+      Descriptor->Flags.has(InstructionFlag::Barrier) ||
+      Descriptor->Flags.has(InstructionFlag::UnmodeledSideEffects) ||
+      hasProtectedAttachment(Instruction) || Instruction.memOperands().empty()) {
+    return std::nullopt;
+  }
+  for (const MachineMemOperand *Memory : Instruction.memOperands()) {
+    if (Memory == nullptr || Memory->isStore() || Memory->isAtomic() || Memory->isVolatile()) {
+      return std::nullopt;
+    }
+  }
+
+  MachineOperand *Definition = nullptr;
+  Register DefinedRegister;
+  for (std::size_t Index = 0; Index < Instruction.operandCount(); ++Index) {
+    MachineOperand *Operand = Instruction.operandAt(Index);
+    const Register *Reg = Operand == nullptr ? nullptr : Operand->getIf<Register>();
+    if (Reg == nullptr || !Operand->isDef())
+      continue;
+    if (Definition != nullptr || !Reg->isVirtual())
+      return std::nullopt;
+    Definition = Operand;
+    DefinedRegister = *Reg;
+  }
+  if (Definition == nullptr || !Function.registerInfo().hasOneUse(DefinedRegister, true)) {
+    return std::nullopt;
+  }
+  MachineOperand *Use = nullptr;
+  for (MachineOperand &Candidate : Function.registerInfo().useOperands(DefinedRegister, false)) {
+    Use = &Candidate;
+    break;
+  }
+  MachineInstr *Consumer = Use == nullptr ? nullptr : Use->parent();
+  if (Consumer == nullptr || Consumer->parent() != Instruction.parent() ||
+      Instruction.next() != Consumer || hasProtectedAttachment(*Consumer) || Use->isImplicit() ||
+      Use->tiedOperandIndex()) {
+    return std::nullopt;
+  }
+  for (std::size_t Index = 0; Index < Consumer->operandCount(); ++Index) {
+    if (Consumer->operandAt(Index) == Use)
+      return FoldCandidate{&Instruction, Consumer, DefinedRegister, Index};
+  }
+  return std::nullopt;
+}
+
+[[nodiscard]] bool containsMemoryDescription(const MachineInstr &Instruction,
+                                             const MachineMemOperandDesc &Description) noexcept {
+  return std::ranges::any_of(Instruction.memOperands(),
+                             [&Description](const MachineMemOperand *Memory) {
+                               return Memory != nullptr && Memory->description() == Description;
+                             });
+}
+
+[[nodiscard]] Expected<bool> applyFold(FoldCandidate Candidate, MachineFunction &Function,
+                                       const TargetDescription &Target,
+                                       const TargetMemoryFoldingInfo &Info) {
+  auto CanFold = Info.canFoldMemoryOperand(*Candidate.Consumer, Candidate.ConsumerOperand,
+                                           *Candidate.MemoryInstruction, Function, Target);
+  if (!CanFold)
+    return std::unexpected(CanFold.error());
+  if (!*CanFold)
+    return false;
+  std::vector<MachineMemOperandDesc> MemoryDescriptions;
+  for (const MachineMemOperand *Memory : Candidate.MemoryInstruction->memOperands()) {
+    if (Memory == nullptr) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "memory folding encountered a null memory operand");
+    }
+    MemoryDescriptions.push_back(Memory->description());
+  }
+  auto Folded = Info.foldMemoryOperand(*Candidate.Consumer, Candidate.ConsumerOperand,
+                                       *Candidate.MemoryInstruction, Function, Target);
+  if (!Folded)
+    return std::unexpected(Folded.error());
+  if (*Folded == nullptr || (*Folded)->parent() == nullptr ||
+      Candidate.MemoryInstruction->parent() != nullptr) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "target memory fold did not replace the consumer and erase the producer");
+  }
+  for (const MachineMemOperandDesc &Description : MemoryDescriptions) {
+    if (!containsMemoryDescription(**Folded, Description)) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "target memory fold did not migrate every memory operand");
+    }
+  }
+  auto FoldedSize = Info.getFoldedMemorySize(**Folded, Function, Target);
+  if (!FoldedSize)
+    return std::unexpected(FoldedSize.error());
+  if (!FoldedSize->has_value() || !FoldedSize->value().isKnown()) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "target folded instruction has no known folded memory size");
+  }
+  if (MemoryDescriptions.size() == 1 && MemoryDescriptions.front().Size.isKnown() &&
+      FoldedSize->value() != MemoryDescriptions.front().Size) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "target folded memory size disagrees with the migrated MMO");
+  }
+  if (!Function.registerInfo().useEmpty(Candidate.DefinedRegister, true) ||
+      Function.registerInfo().getUniqueVRegDef(Candidate.DefinedRegister) != nullptr) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "target memory fold retained the load result register");
+  }
+  auto Erased = Function.registerInfo().eraseVirtualRegister(Candidate.DefinedRegister);
+  if (!Erased)
+    return std::unexpected(Erased.error());
+  return true;
+}
+
+} // namespace
+
+const MachinePassContract &MachineMemoryFoldingPass::contract() const noexcept {
+  return memoryFoldingContract();
+}
+
+Expected<MachinePassResult> MachineMemoryFoldingPass::run(MachineFunction &Function,
+                                                          const TargetDescription &Target,
+                                                          MachineAnalysisManager &) {
+  if (!Function.properties().containsAll(contract().Requires)) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "memory folding requires selected SSA MachineIR");
+  }
+  auto Trial = cloneMachineFunction(
+      Function, Target,
+      MachineFunctionCloneOptions{.NewStableID = Function.stableID(), .NewName = std::nullopt});
+  if (!Trial) {
+    return std::unexpected(
+        prefixedError(Trial.error(), "memory folding 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(prefixedError(
+        Result.error(), "memory folding diverged after a successful transaction preflight: "));
+  }
+  return Result;
+}
+
+Expected<MachinePassResult>
+MachineMemoryFoldingPass::runImplementation(MachineFunction &Function,
+                                            const TargetDescription &Target) const {
+  MachinePassResult Result;
+  std::size_t FoldCount = 0;
+  bool LocalChange = true;
+  while (LocalChange && FoldCount < Options.MaximumFolds) {
+    LocalChange = false;
+    for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr && !LocalChange;
+         Block = Block->next()) {
+      for (MachineInstr *Instruction = Block->firstInstruction(); Instruction != nullptr;
+           Instruction = Instruction->next()) {
+        std::optional<FoldCandidate> Candidate = findCandidate(*Instruction, Function, Target);
+        if (!Candidate)
+          continue;
+        auto Folded = applyFold(*Candidate, Function, Target, *Info);
+        if (!Folded)
+          return std::unexpected(Folded.error());
+        if (!*Folded)
+          continue;
+        ++FoldCount;
+        Result.Changed = true;
+        LocalChange = true;
+        break;
+      }
+    }
+  }
+  if (Result.Changed) {
+    Result.Changes = {MachinePassChange::Instructions, MachinePassChange::Operands,
+                      MachinePassChange::Registers};
+    Result.Preserved = PreservedAnalyses::none();
+  } else {
+    Result.Preserved = PreservedAnalyses::all();
+  }
+  return Result;
+}
+
+} // namespace rcc