[CodeGen] Track reserved register defs precisely

3 files changed, 33 insertions, 12 deletions

diff --git a/backend/lib/CodeGen/MachineLiveness.cpp b/backend/lib/CodeGen/MachineLiveness.cpp
index 528a675..7984e7d 100644
--- a/backend/lib/CodeGen/MachineLiveness.cpp
+++ b/backend/lib/CodeGen/MachineLiveness.cpp
@@ -846,11 +846,10 @@ struct LivenessRebuildPlan final {
   std::vector<OperandFlagPlan> Operands;
 };
 
-[[nodiscard]] Expected<bool> liveSetIntersects(const MachineOperand &Operand,
-                                               const RegisterLaneSet &Virtual,
-                                               const RegisterUnitLaneSet &Physical,
-                                               const MachineFunction &Function,
-                                               const TargetDescription &Target) {
+[[nodiscard]] Expected<bool>
+liveSetIntersects(const MachineOperand &Operand, const RegisterLaneSet &Virtual,
+                  const RegisterUnitLaneSet &Physical, const MachineFunction &Function,
+                  const TargetDescription &Target, bool TreatReservedAsLive) {
   const Register *Reg = Operand.getIf<Register>();
   if (Reg == nullptr || Reg->isNoRegister())
     return false;
@@ -860,7 +859,7 @@ struct LivenessRebuildPlan final {
       return std::unexpected(Lanes.error());
     return Virtual.intersects(*Reg, *Lanes);
   }
-  if (isReservedPhysicalRegister(*Reg, Function, Target))
+  if (TreatReservedAsLive && isReservedPhysicalRegister(*Reg, Function, Target))
     return true;
   auto Units = physicalOperandUnits(Operand, Target);
   if (!Units)
@@ -947,7 +946,7 @@ buildLivenessRebuildPlan(const MachineFunction &Function, const TargetDescriptio
         const MachineOperand *Operand = Instruction->operandAt(Index);
         if (Operand == nullptr || !Operand->isRegister() || !Operand->isDef())
           continue;
-        auto IsLive = liveSetIntersects(*Operand, Virtual, Physical, Function, Target);
+        auto IsLive = liveSetIntersects(*Operand, Virtual, Physical, Function, Target, false);
         if (!IsLive)
           return std::unexpected(IsLive.error());
         RegisterOperandFlags Flags = Operand->registerFlags();
@@ -1003,7 +1002,7 @@ buildLivenessRebuildPlan(const MachineFunction &Function, const TargetDescriptio
                                    !Operand->registerValue()->isNoRegister();
         bool IsKill = false;
         if (IsSemanticUse) {
-          auto IsLive = liveSetIntersects(*Operand, Virtual, Physical, Function, Target);
+          auto IsLive = liveSetIntersects(*Operand, Virtual, Physical, Function, Target, true);
           if (!IsLive)
             return std::unexpected(IsLive.error());
           IsKill = !*IsLive;
diff --git a/backend/lib/CodeGen/PostRAScheduler.cpp b/backend/lib/CodeGen/PostRAScheduler.cpp
index 204f05e..6a91239 100644
--- a/backend/lib/CodeGen/PostRAScheduler.cpp
+++ b/backend/lib/CodeGen/PostRAScheduler.cpp
@@ -182,7 +182,7 @@ buildBlockLivenessRepairPlan(MachineBasicBlock &Block, const TargetDescription &
       auto Units = physicalOperandUnits(*Operand, Target);
       if (!Units)
         return std::unexpected(Units.error());
-      if (!isReservedPhysicalRegister(*Reg, Block.function(), Target) && !intersects(Live, *Units))
+      if (!intersects(Live, *Units))
         Update.Flags.set(RegisterOperandFlag::Dead);
     }
 
diff --git a/backend/unittests/CodeGen/MachineLivenessTest.cpp b/backend/unittests/CodeGen/MachineLivenessTest.cpp
index 84a247f..cca7629 100644
--- a/backend/unittests/CodeGen/MachineLivenessTest.cpp
+++ b/backend/unittests/CodeGen/MachineLivenessTest.cpp
@@ -884,16 +884,38 @@ void testRISCVReservedRegisterLiveness(TestContext &Context) {
   auto Instruction = Immediate
                          ? Pending->build()
                          : rcc::Expected<rcc::MachineInstr *>{std::unexpected(Immediate.error())};
-  if (!Instruction) {
+  auto DeadDef =
+      Instruction ? Builder.buildInstr(rcc::Opcode::fromValue(
+                        static_cast<std::uint32_t>(rcc::riscv::InstructionOpcode::ADDI)))
+                  : rcc::Expected<rcc::MachineInstrBuilder>{std::unexpected(Instruction.error())};
+  auto DeadDefDestination = DeadDef ? DeadDef->addDef(rcc::Register::fromPhysical(*X0))
+                                    : rcc::Expected<void>{std::unexpected(DeadDef.error())};
+  auto DeadDefSource = DeadDefDestination
+                           ? DeadDef->addUse(rcc::Register::fromPhysical(*X0))
+                           : rcc::Expected<void>{std::unexpected(DeadDefDestination.error())};
+  auto DeadDefImmediate = DeadDefSource
+                              ? DeadDef->addImmediate(0)
+                              : rcc::Expected<void>{std::unexpected(DeadDefSource.error())};
+  auto DeadDefInstruction =
+      DeadDefImmediate
+          ? DeadDef->build()
+          : rcc::Expected<rcc::MachineInstr *>{std::unexpected(DeadDefImmediate.error())};
+  if (!Instruction || !DeadDefInstruction) {
     Context.expect(false, "RV64GC reserved-register liveness instruction must be built");
     return;
   }
 
   auto Rebuilt = rcc::rebuildMachineLiveness(*Fixture->Function, *Target);
   const rcc::MachineOperand *ZeroUse = (*Instruction)->operandAt(1);
+  const rcc::MachineOperand *ZeroDef = (*DeadDefInstruction)->operandAt(0);
+  const rcc::MachineOperand *DeadDefZeroUse = (*DeadDefInstruction)->operandAt(1);
   Context.expect(Rebuilt && Fixture->Blocks[0]->liveIns().empty() && ZeroUse != nullptr &&
-                     !ZeroUse->registerFlags().has(rcc::RegisterOperandFlag::Kill),
-                 "reserved registers must remain available without live-ins or kill flags");
+                     !ZeroUse->registerFlags().has(rcc::RegisterOperandFlag::Kill) &&
+                     ZeroDef != nullptr &&
+                     ZeroDef->registerFlags().has(rcc::RegisterOperandFlag::Dead) &&
+                     DeadDefZeroUse != nullptr &&
+                     !DeadDefZeroUse->registerFlags().has(rcc::RegisterOperandFlag::Kill),
+                 "reserved uses must remain available while dead defs retain precise liveness");
 }
 
 } // namespace