[RISCV] Implement push-pop optimization

10 files changed, 318 insertions, 27 deletions

diff --git a/backend/include/rcc/CodeGen/MachineOptimizationPipeline.h b/backend/include/rcc/CodeGen/MachineOptimizationPipeline.h
index c436455..f60bbda 100644
--- a/backend/include/rcc/CodeGen/MachineOptimizationPipeline.h
+++ b/backend/include/rcc/CodeGen/MachineOptimizationPipeline.h
@@ -7,6 +7,8 @@
 #include "rcc/CodeGen/MachineScheduler.h"
 #include "rcc/CodeGen/TargetFrameLowering.h"
 
+#include <memory>
+
 namespace rcc {
 
 class TargetInstrInfo;
@@ -41,10 +43,10 @@ addSelectedSSAOptimizationPipeline(MachineFunctionPassManager &Pipeline,
 [[nodiscard]] Expected<void> addPreRegisterAllocationOptimizationPipeline(
     MachineFunctionPassManager &Pipeline, const MachineOptimizationPipelineTargetInfo &TargetInfo);
 
-[[nodiscard]] Expected<void>
-addPostRegisterAllocationPipeline(MachineFunctionPassManager &Pipeline,
-                                  const MachineOptimizationPipelineTargetInfo &TargetInfo,
-                                  MachineOptimizationPipelineOptions Options = {});
+[[nodiscard]] Expected<void> addPostRegisterAllocationPipeline(
+    MachineFunctionPassManager &Pipeline, const MachineOptimizationPipelineTargetInfo &TargetInfo,
+    MachineOptimizationPipelineOptions Options = {},
+    std::unique_ptr<MachineFunctionPass> PostRAPseudoExpansionTargetPass = nullptr);
 
 } // namespace rcc
 
diff --git a/backend/lib/CodeGen/MachineOptimizationPipeline.cpp b/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
index e61dfdc..cbc4b1a 100644
--- a/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
+++ b/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
@@ -77,10 +77,10 @@ Expected<void> addPreRegisterAllocationOptimizationPipeline(
                                        TargetInfo.SchedulerOptions);
 }
 
-Expected<void>
-addPostRegisterAllocationPipeline(MachineFunctionPassManager &Pipeline,
-                                  const MachineOptimizationPipelineTargetInfo &TargetInfo,
-                                  MachineOptimizationPipelineOptions Options) {
+Expected<void> addPostRegisterAllocationPipeline(
+    MachineFunctionPassManager &Pipeline, const MachineOptimizationPipelineTargetInfo &TargetInfo,
+    MachineOptimizationPipelineOptions Options,
+    std::unique_ptr<MachineFunctionPass> PostRAPseudoExpansionTargetPass) {
   auto Added = addPass<MachineStackSlotColoringPass>(Pipeline);
   if (!Added)
     return std::unexpected(Added.error());
@@ -112,6 +112,11 @@ addPostRegisterAllocationPipeline(MachineFunctionPassManager &Pipeline,
   Added = addPass<ExpandPostRAPseudosPass>(Pipeline, TargetInfo.PseudoExpansion);
   if (!Added)
     return std::unexpected(Added.error());
+  if (PostRAPseudoExpansionTargetPass != nullptr) {
+    Added = Pipeline.addPass(std::move(PostRAPseudoExpansionTargetPass));
+    if (!Added)
+      return std::unexpected(Added.error());
+  }
   Added = addPass<PostRASchedulerPass>(Pipeline, TargetInfo.SchedulingModel,
                                        TargetInfo.SchedulerOptions);
   if (!Added)
diff --git a/backend/lib/CodeGen/MachinePseudoExpansion.cpp b/backend/lib/CodeGen/MachinePseudoExpansion.cpp
index 3d169cc..67a8ef8 100644
--- a/backend/lib/CodeGen/MachinePseudoExpansion.cpp
+++ b/backend/lib/CodeGen/MachinePseudoExpansion.cpp
@@ -328,7 +328,8 @@ expandImplementation(MachineFunction &Function, const TargetDescription &Target,
                             "pseudo expansion left an instruction owned by the completed phase");
     }
   }
-  if (Changed && Function.properties().has(MachineFunctionProperty::TracksLiveness)) {
+  if (Changed && Phase != MachinePseudoExpansionPhase::PostRA &&
+      Function.properties().has(MachineFunctionProperty::TracksLiveness)) {
     auto Rebuilt = rebuildMachineLiveness(Function, Target);
     if (!Rebuilt)
       return std::unexpected(Rebuilt.error());
diff --git a/backend/lib/Target/RISCV/CMakeLists.txt b/backend/lib/Target/RISCV/CMakeLists.txt
index 0e93ea9..9a3973d 100644
--- a/backend/lib/Target/RISCV/CMakeLists.txt
+++ b/backend/lib/Target/RISCV/CMakeLists.txt
@@ -131,6 +131,7 @@ add_library(
   RISCVOutliner.cpp
   RISCVPipeline.cpp
   RISCVPseudoExpansion.cpp
+  RISCVPushPopOptimizer.cpp
   RISCVRegisterAllocation.cpp
   RISCVScheduling.cpp
   RISCVSubtarget.cpp
@@ -218,6 +219,7 @@ rcc_backend_register_format_files(
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVOptimizationPipeline.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVPipeline.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVPseudoExpansion.cpp"
+  "${CMAKE_CURRENT_SOURCE_DIR}/RISCVPushPopOptimizer.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVRegisterAllocation.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVScheduling.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVSubtarget.cpp"
@@ -261,6 +263,7 @@ rcc_backend_register_tidy_files(
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVOptimizationPipeline.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVPipeline.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVPseudoExpansion.cpp"
+  "${CMAKE_CURRENT_SOURCE_DIR}/RISCVPushPopOptimizer.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVRegisterAllocation.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVScheduling.cpp"
   "${CMAKE_CURRENT_SOURCE_DIR}/RISCVSubtarget.cpp"
diff --git a/backend/lib/Target/RISCV/RISCV.h b/backend/lib/Target/RISCV/RISCV.h
index e7ab865..0b4139a 100644
--- a/backend/lib/Target/RISCV/RISCV.h
+++ b/backend/lib/Target/RISCV/RISCV.h
@@ -45,6 +45,8 @@ createMoveMergePass(const RISCVSubtarget &Subtarget);
 [[nodiscard]] std::unique_ptr<MachineFunctionPass>
 createMakeCompressiblePass(const RISCVSubtarget &Subtarget);
 [[nodiscard]] std::unique_ptr<MachineFunctionPass>
+createPushPopOptimizationPass(const RISCVSubtarget &Subtarget);
+[[nodiscard]] std::unique_ptr<MachineFunctionPass>
 createOptWInstrsPass(const RISCVSubtarget &Subtarget);
 [[nodiscard]] std::unique_ptr<MachineFunctionPass> createMergeBaseOffsetPass();
 [[nodiscard]] MachineOptimizationPipelineTargetInfo
diff --git a/backend/lib/Target/RISCV/RISCVFrameLowering.cpp b/backend/lib/Target/RISCV/RISCVFrameLowering.cpp
index 7ffaccb..8a40bc1 100644
--- a/backend/lib/Target/RISCV/RISCVFrameLowering.cpp
+++ b/backend/lib/Target/RISCV/RISCVFrameLowering.cpp
@@ -2280,6 +2280,12 @@ public:
         if (auto Result = Pending->addImplicitDef(Register::fromPhysical(Reg)); !Result)
           return std::unexpected(Result.error());
       }
+      for (const MachineOperand &Operand : Terminator.operands()) {
+        if (!Operand.isImplicit())
+          continue;
+        if (auto Result = Pending->addImplicitOperand(Operand.cloneDetached()); !Result)
+          return std::unexpected(Result.error());
+      }
       if (auto Result = setFrameFlag(*Pending, MachineInstrFlag::FrameDestroy); !Result)
         return std::unexpected(Result.error());
       auto Built = Pending->build();
diff --git a/backend/lib/Target/RISCV/RISCVInstrInfoPseudo.rccdesc b/backend/lib/Target/RISCV/RISCVInstrInfoPseudo.rccdesc
index a0767f9..19a0975 100644
--- a/backend/lib/Target/RISCV/RISCVInstrInfoPseudo.rccdesc
+++ b/backend/lib/Target/RISCV/RISCVInstrInfoPseudo.rccdesc
@@ -320,7 +320,7 @@ instructions:
     allow_extra_implicit_operands: true
     flags: [pseudo, produces_code, allowed_pre_legalizer, call, may_load, may_store, has_side_effects, not_duplicable, unmodeled_side_effects]
     size: 8
-    owner: post_ra_expansion
+    owner: target_late_expansion
     deadline: post_ra_finalized
   -
     id: 1223
@@ -333,7 +333,7 @@ instructions:
     allow_extra_implicit_operands: true
     flags: [pseudo, produces_code, allowed_pre_legalizer, terminator, branch, call, barrier, may_load, may_store, has_side_effects, not_duplicable, unmodeled_side_effects]
     size: 8
-    owner: post_ra_expansion
+    owner: target_late_expansion
     deadline: post_ra_finalized
   -
     id: 1269
@@ -347,7 +347,7 @@ instructions:
     allow_extra_implicit_operands: true
     flags: [pseudo, produces_code, allowed_pre_legalizer, call, may_load, may_store, has_side_effects, not_duplicable, unmodeled_side_effects]
     size: 8
-    owner: post_ra_expansion
+    owner: target_late_expansion
     deadline: post_ra_finalized
   -
     id: 1224
diff --git a/backend/lib/Target/RISCV/RISCVPipeline.cpp b/backend/lib/Target/RISCV/RISCVPipeline.cpp
index 331d81f..84fe012 100644
--- a/backend/lib/Target/RISCV/RISCVPipeline.cpp
+++ b/backend/lib/Target/RISCV/RISCVPipeline.cpp
@@ -121,7 +121,8 @@ template <typename PassT, typename... ArgumentT>
       optimizationPipelineInfo(Subtarget, Options.DebugMachineScheduler,
                                Options.DebugPrologueEpilogue),
       MachineOptimizationPipelineOptions{.UseTargetCopyInstructions =
-                                             Options.UseTargetCopyInstructions});
+                                             Options.UseTargetCopyInstructions},
+      createPushPopOptimizationPass(Subtarget));
   if (!Added)
     return std::unexpected(Added.error());
   Added = Pipeline.addPass(createMakeCompressiblePass(Subtarget));
@@ -194,7 +195,8 @@ Expected<void> addCanonicalPipeline(MachineFunctionPassManager &Pipeline,
         optimizationPipelineInfo(Subtarget, Options.DebugMachineScheduler,
                                  Options.DebugPrologueEpilogue),
         MachineOptimizationPipelineOptions{.UseTargetCopyInstructions =
-                                               Options.UseTargetCopyInstructions});
+                                               Options.UseTargetCopyInstructions},
+        createPushPopOptimizationPass(Subtarget));
     if (!Added)
       return std::unexpected(Added.error());
   }
diff --git a/backend/lib/Target/RISCV/RISCVPseudoExpansion.cpp b/backend/lib/Target/RISCV/RISCVPseudoExpansion.cpp
index fb09772..6945f68 100644
--- a/backend/lib/Target/RISCV/RISCVPseudoExpansion.cpp
+++ b/backend/lib/Target/RISCV/RISCVPseudoExpansion.cpp
@@ -585,19 +585,19 @@ public:
         return expandAddressPseudo(Instruction, Function, Target, true);
       return false;
     }
-    if (Phase != MachinePseudoExpansionPhase::PostRA)
-      return false;
-    if (Instruction.opcode() == opcode(GenericOpcode::COPY))
-      return expandCopy(Instruction, Function, Target);
-    auto Subregister = expandSubregisterCopy(Instruction, Function, Target);
-    if (!Subregister || *Subregister)
-      return Subregister;
-    if (Instruction.opcode() == instructionOpcode(InstructionOpcode::PseudoCALL))
-      return expandCall(Instruction, Function, Target, false);
-    if (Instruction.opcode() == instructionOpcode(InstructionOpcode::PseudoCALLReg))
-      return expandCall(Instruction, Function, Target, false);
-    if (Instruction.opcode() == instructionOpcode(InstructionOpcode::PseudoTAIL))
-      return expandCall(Instruction, Function, Target, true);
+    if (Phase == MachinePseudoExpansionPhase::PostRA) {
+      if (Instruction.opcode() == opcode(GenericOpcode::COPY))
+        return expandCopy(Instruction, Function, Target);
+      return expandSubregisterCopy(Instruction, Function, Target);
+    }
+    if (Phase == MachinePseudoExpansionPhase::TargetLate) {
+      if (Instruction.opcode() == instructionOpcode(InstructionOpcode::PseudoCALL))
+        return expandCall(Instruction, Function, Target, false);
+      if (Instruction.opcode() == instructionOpcode(InstructionOpcode::PseudoCALLReg))
+        return expandCall(Instruction, Function, Target, false);
+      if (Instruction.opcode() == instructionOpcode(InstructionOpcode::PseudoTAIL))
+        return expandCall(Instruction, Function, Target, true);
+    }
     return false;
   }
 };
diff --git a/backend/lib/Target/RISCV/RISCVPushPopOptimizer.cpp b/backend/lib/Target/RISCV/RISCVPushPopOptimizer.cpp
new file mode 100644
index 0000000..16928fb
--- /dev/null
+++ b/backend/lib/Target/RISCV/RISCVPushPopOptimizer.cpp
@@ -0,0 +1,270 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "RISCV.h"
+#include "RISCVInstrInfo.h"
+#include "rcc/CodeGen/InstructionDescriptor.h"
+#include "rcc/CodeGen/LiveRegUnits.h"
+#include "rcc/CodeGen/MachineBasicBlock.h"
+#include "rcc/CodeGen/MachineFunction.h"
+#include "rcc/CodeGen/MachineIRBuilder.h"
+#include "rcc/CodeGen/MachineInstr.h"
+#include "rcc/CodeGen/MachineOperand.h"
+#include "rcc/CodeGen/RegisterMask.h"
+#include "rcc/CodeGen/TargetDescription.h"
+#include "rcc/Target/RISCV/RISCVRegisters.h"
+#include "rcc/Target/RISCV/RISCVTarget.h"
+
+#include <cstddef>
+#include <memory>
+#include <optional>
+#include <utility>
+
+namespace rcc::riscv {
+namespace {
+
+struct PushPopOpcodes final {
+  Opcode Pop;
+  Opcode PopRet;
+  Opcode PopRetZero;
+};
+
+[[nodiscard]] const MachinePassContract &pushPopContract() {
+  static const MachinePassContract Contract = [] {
+    MachinePassContract Value{
+        .Name = "riscv-push-pop-opt",
+        .Requires = {MachineFunctionProperty::NoPHIs, MachineFunctionProperty::NoVRegs,
+                     MachineFunctionProperty::FrameFinalized,
+                     MachineFunctionProperty::NoFrameIndices},
+        .Sets = {},
+        .Clears = {},
+        .RequiredAnalyses = {},
+        .VerificationStage = PipelineStage::PostRAFinalized,
+    };
+    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
+    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands};
+    return Value;
+  }();
+  return Contract;
+}
+
+[[nodiscard]] bool intersects(const LiveRegUnits &Left, const LiveRegUnits &Right) noexcept {
+  for (const auto &[Unit, Lanes] : Left.entries())
+    if (Right.intersects(Unit, Lanes))
+      return true;
+  return false;
+}
+
+[[nodiscard]] Expected<void> accumulateRegisterEffects(const MachineInstr &Instruction,
+                                                       const TargetDescription &Target,
+                                                       LiveRegUnits &Modified, LiveRegUnits &Used) {
+  for (const MachineOperand &Operand : Instruction.operands()) {
+    if (Operand.isDebugUse())
+      continue;
+    if (Operand.isRegister()) {
+      const std::optional<Register> Reg = Operand.registerValue();
+      if (!Reg || !Reg->isPhysical() || Reg->isNoRegister())
+        continue;
+      auto Units = physicalOperandUnits(Operand, Target);
+      if (!Units)
+        return std::unexpected(Units.error());
+      for (const PhysicalUnitAccess &Access : *Units) {
+        if (Operand.isDef())
+          Modified.add(Access.Unit, Access.Lanes);
+        if (Operand.isUse())
+          Used.add(Access.Unit, Access.Lanes);
+      }
+      continue;
+    }
+    if (const auto *Mask = Operand.getIf<RegisterMaskOperand>();
+        Mask != nullptr && Mask->Mask != nullptr) {
+      if (auto Added = addRegisterMaskDefs(Modified, *Mask->Mask, Target); !Added)
+        return std::unexpected(Added.error());
+      continue;
+    }
+    if (const auto *Mask = Operand.getIf<LiveOutRegisterMaskOperand>();
+        Mask != nullptr && Mask->Mask != nullptr) {
+      if (auto Added = addRegisterMaskUses(Used, *Mask->Mask, Target); !Added)
+        return std::unexpected(Added.error());
+    }
+  }
+  return {};
+}
+
+[[nodiscard]] MachineInstr *previousNonDebug(MachineInstr &Instruction) noexcept {
+  MachineInstr *Previous = Instruction.previous();
+  while (Previous != nullptr && Previous->isDebugInstr())
+    Previous = Previous->previous();
+  return Previous;
+}
+
+[[nodiscard]] MachineInstr *firstTerminator(MachineBasicBlock &Block,
+                                            const TargetDescription &Target) noexcept {
+  for (MachineInstr *Instruction = Block.firstInstruction(); Instruction != nullptr;
+       Instruction = Instruction->next()) {
+    const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction->opcode());
+    if (Descriptor != nullptr && Descriptor->Flags.has(InstructionFlag::Terminator))
+      return Instruction;
+  }
+  return nullptr;
+}
+
+[[nodiscard]] Expected<bool> removeZeroReturnAssignment(MachineInstr &Pop,
+                                                        const TargetDescription &Target,
+                                                        PhysicalRegister ReturnRegister,
+                                                        PhysicalRegister ZeroRegister) {
+  const PhysicalRegisterDesc *ReturnDescription =
+      Target.registers().physicalRegister(ReturnRegister);
+  if (ReturnDescription == nullptr)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V return register has no register-unit description");
+  const LiveRegUnits ReturnUnits = physicalRegisterCoverage(*ReturnDescription);
+  LiveRegUnits Modified;
+  LiveRegUnits Used;
+  for (MachineInstr *Instruction = previousNonDebug(Pop); Instruction != nullptr;
+       Instruction = previousNonDebug(*Instruction)) {
+    if (const std::optional<MachineCopyOperands> Copy = instrInfo().getCopyOperands(*Instruction)) {
+      const std::optional<Register> Destination = Copy->Destination->registerValue();
+      const std::optional<Register> Source = Copy->Source->registerValue();
+      if (Destination && Source && Destination->physicalRegister() == ReturnRegister &&
+          Source->physicalRegister() == ZeroRegister) {
+        MachineBasicBlock *Parent = Instruction->parent();
+        if (Parent == nullptr)
+          return makeUnexpected(ErrorCode::InternalError,
+                                "RISC-V return-value assignment became detached");
+        if (auto Removed = Parent->removeInstruction(*Instruction); !Removed)
+          return std::unexpected(Removed.error());
+        return true;
+      }
+    }
+    if (auto Accumulated = accumulateRegisterEffects(*Instruction, Target, Modified, Used);
+        !Accumulated) {
+      return std::unexpected(Accumulated.error());
+    }
+    if (intersects(Modified, ReturnUnits) || intersects(Used, ReturnUnits))
+      return false;
+  }
+  return false;
+}
+
+[[nodiscard]] Expected<MachineInstr *> buildPopReturn(MachineInstr &Pop, MachineInstr &Return,
+                                                      bool ReturnZero,
+                                                      const PushPopOpcodes &Opcodes,
+                                                      MachineFunction &Function,
+                                                      const TargetDescription &Target) {
+  const InstructionDescriptor *PopDescriptor = Target.instructionDescriptor(Pop.opcode());
+  if (PopDescriptor == nullptr || PopDescriptor->Operands.size() != 2U)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V CM_POP has an invalid instruction descriptor");
+  MachineIRBuilder Builder(Function, Target);
+  if (auto Point = Builder.setInsertPointBefore(Return); !Point)
+    return std::unexpected(Point.error());
+  auto Pending = Builder.buildInstr(ReturnZero ? Opcodes.PopRetZero : Opcodes.PopRet);
+  if (!Pending)
+    return std::unexpected(Pending.error());
+  for (std::size_t Index = 0; Index < PopDescriptor->Operands.size(); ++Index) {
+    const MachineOperand *Operand = Pop.operandAt(Index);
+    if (Operand == nullptr)
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "RISC-V CM_POP is missing an explicit operand");
+    if (auto Added = Pending->addOperand(Operand->cloneDetached()); !Added)
+      return std::unexpected(Added.error());
+  }
+  const std::size_t FirstDynamicImplicit = PopDescriptor->Operands.size() +
+                                           PopDescriptor->ImplicitDefs.size() +
+                                           PopDescriptor->ImplicitUses.size();
+  for (std::size_t Index = FirstDynamicImplicit; Index < Pop.operandCount(); ++Index) {
+    const MachineOperand *Operand = Pop.operandAt(Index);
+    if (Operand == nullptr || !Operand->isImplicit()) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "RISC-V CM_POP has a malformed dynamic implicit operand");
+    }
+    if (auto Added = Pending->addImplicitOperand(Operand->cloneDetached()); !Added)
+      return std::unexpected(Added.error());
+  }
+  if (auto Flags = Pending->setFlags({MachineInstrFlag::FrameDestroy}); !Flags)
+    return std::unexpected(Flags.error());
+  if (Return.debugLocation()) {
+    if (auto Location = Pending->setDebugLocation(*Return.debugLocation()); !Location)
+      return std::unexpected(Location.error());
+  }
+  return Pending->build();
+}
+
+class RISCVPushPopOptimizationPass final : public MachineFunctionPass {
+public:
+  explicit RISCVPushPopOptimizationPass(const RISCVSubtarget &Subtarget) noexcept
+      : Subtarget(&Subtarget) {}
+
+  [[nodiscard]] const MachinePassContract &contract() const noexcept override {
+    return pushPopContract();
+  }
+
+  [[nodiscard]] Expected<MachinePassResult> run(MachineFunction &Function,
+                                                const TargetDescription &Target,
+                                                MachineFunctionAnalysisManager &) override {
+    MachinePassResult Result;
+    if (!Subtarget->features().has(RISCVExtension::Zcmp)) {
+      Result.Preserved = PreservedAnalyses::all();
+      return Result;
+    }
+
+    const InstructionDescriptor *Pop = Target.instructions().descriptor("CM_POP");
+    const InstructionDescriptor *PopRet = Target.instructions().descriptor("CM_POPRET");
+    const InstructionDescriptor *PopRetZero = Target.instructions().descriptor("CM_POPRETZ");
+    if (Pop == nullptr || PopRet == nullptr || PopRetZero == nullptr) {
+      return makeUnexpected(ErrorCode::Unsupported,
+                            "RISC-V Zcmp push/pop instructions are not registered");
+    }
+    auto ReturnRegister = physicalRegister(PhysicalRegisterID::X10);
+    if (!ReturnRegister)
+      return std::unexpected(ReturnRegister.error());
+    auto ZeroRegister = physicalRegister(PhysicalRegisterID::X0);
+    if (!ZeroRegister)
+      return std::unexpected(ZeroRegister.error());
+    const PushPopOpcodes Opcodes{Pop->InstructionOpcode, PopRet->InstructionOpcode,
+                                 PopRetZero->InstructionOpcode};
+
+    for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr;
+         Block = Block->next()) {
+      MachineInstr *Return = firstTerminator(*Block, Target);
+      if (Return == nullptr ||
+          Return->opcode() != instructionOpcode(InstructionOpcode::PseudoRET)) {
+        continue;
+      }
+      MachineInstr *PopInstruction = previousNonDebug(*Return);
+      if (PopInstruction == nullptr || PopInstruction->opcode() != Opcodes.Pop ||
+          !PopInstruction->flags().has(MachineInstrFlag::FrameDestroy)) {
+        continue;
+      }
+      auto ReturnZero =
+          removeZeroReturnAssignment(*PopInstruction, Target, *ReturnRegister, *ZeroRegister);
+      if (!ReturnZero)
+        return std::unexpected(ReturnZero.error());
+      auto NewReturn =
+          buildPopReturn(*PopInstruction, *Return, *ReturnZero, Opcodes, Function, Target);
+      if (!NewReturn)
+        return std::unexpected(NewReturn.error());
+      if (auto Removed = Block->removeInstruction(*PopInstruction); !Removed)
+        return std::unexpected(Removed.error());
+      if (auto Removed = Block->removeInstruction(*Return); !Removed)
+        return std::unexpected(Removed.error());
+      Result.Changes.add(MachinePassChange::Instructions);
+      Result.Changes.add(MachinePassChange::Operands);
+    }
+    Result.Changed = !Result.Changes.empty();
+    Result.Preserved = Result.Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+    return Result;
+  }
+
+private:
+  const RISCVSubtarget *Subtarget;
+};
+
+} // namespace
+
+std::unique_ptr<MachineFunctionPass>
+createPushPopOptimizationPass(const RISCVSubtarget &Subtarget) {
+  return std::make_unique<RISCVPushPopOptimizationPass>(Subtarget);
+}
+
+} // namespace rcc::riscv