[CodeGen] Verify local physical live-ins

2 files changed, 103 insertions, 3 deletions

diff --git a/backend/lib/CodeGen/MachineLiveness.cpp b/backend/lib/CodeGen/MachineLiveness.cpp
index 7984e7d..9882b11 100644
--- a/backend/lib/CodeGen/MachineLiveness.cpp
+++ b/backend/lib/CodeGen/MachineLiveness.cpp
@@ -900,6 +900,66 @@ liveSetIntersects(const MachineOperand &Operand, const RegisterLaneSet &Virtual,
   return addOperandUse(Virtual, Physical, Operand, Function, Target);
 }
 
+[[nodiscard]] Expected<RegisterUnitLaneSet>
+requiredLocalPhysicalLiveIns(const MachineBasicBlock &Block, const MachineFunction &Function,
+                             const TargetDescription &Target) {
+  RegisterLaneSet Virtual;
+  RegisterUnitLaneSet Physical;
+  for (const MachineInstr *Instruction = Block.lastInstruction(); Instruction != nullptr;
+       Instruction = Instruction->previous()) {
+    const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction->opcode());
+    if (Descriptor == nullptr) {
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "liveness verification encountered an unknown instruction descriptor");
+    }
+    const bool IsPhi = Descriptor->Flags.has(InstructionFlag::Phi);
+    for (const MachineOperand &Operand : Instruction->operands()) {
+      const auto *MaskOperand = Operand.getIf<LiveOutRegisterMaskOperand>();
+      if (MaskOperand == nullptr)
+        continue;
+      if (!MaskOperand->Mask) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "liveness verification encountered a null live-out register mask");
+      }
+      auto Added = addRegisterMaskUses(Physical, *MaskOperand->Mask, Target);
+      if (!Added)
+        return std::unexpected(Added.error());
+    }
+
+    for (const MachineOperand &Operand : Instruction->operands()) {
+      if (Operand.isRegister() && Operand.isDef()) {
+        auto Removed = removeOperandFromLiveSet(Operand, Virtual, Physical, Function, Target);
+        if (!Removed)
+          return std::unexpected(Removed.error());
+        continue;
+      }
+      const auto *MaskOperand = Operand.getIf<RegisterMaskOperand>();
+      if (MaskOperand == nullptr)
+        continue;
+      if (!MaskOperand->Mask) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "liveness verification encountered a null register mask");
+      }
+      RegisterUnitLaneSet Clobbers;
+      auto Added = addRegisterMaskDefs(Clobbers, *MaskOperand->Mask, Target);
+      if (!Added)
+        return std::unexpected(Added.error());
+      Physical.subtract(Clobbers);
+    }
+
+    if (IsPhi)
+      continue;
+    for (const MachineOperand &Operand : Instruction->operands()) {
+      if (!Operand.isRegister() || !Operand.isUse())
+        continue;
+      auto Added = addOperandToLiveSet(Operand, Virtual, Physical, Function, Target);
+      if (!Added)
+        return std::unexpected(Added.error());
+    }
+  }
+  return Physical;
+}
+
 [[nodiscard]] Expected<LivenessRebuildPlan>
 buildLivenessRebuildPlan(const MachineFunction &Function, const TargetDescription &Target) {
   auto Variables = computeLiveVariables(Function, Target);
@@ -1792,14 +1852,24 @@ Expected<void> verifyTrackedMachineLiveness(const MachineFunction &Function,
   auto Plan = buildLivenessRebuildPlan(Function, Target);
   if (!Plan)
     return std::unexpected(Plan.error());
-  for (const BlockLivenessPlan &BlockPlan : Plan->Blocks) {
-    auto Actual = liveInUnits(BlockPlan.OriginalLiveIns, Target);
+  for (const MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr;
+       Block = Block->next()) {
+    auto Actual = liveInUnits(Block->liveIns(), Target);
     if (!Actual)
       return std::unexpected(Actual.error());
-    auto Expected = liveInUnits(BlockPlan.RebuiltLiveIns, Target);
+    auto Expected = requiredLocalPhysicalLiveIns(*Block, Function, Target);
     if (!Expected)
       return std::unexpected(Expected.error());
     if (!isSubset(*Expected, *Actual)) {
+      const auto Missing = std::ranges::find_if(Expected->entries(), [&Actual](const auto &Entry) {
+        return !Actual->contains(Entry.first, Entry.second);
+      });
+      if (Missing != Expected->entries().end()) {
+        return makeUnexpected(ErrorCode::VerificationError,
+                              "TracksLiveness block " + std::to_string(Block->stableID().value()) +
+                                  " is missing required physical register unit " +
+                                  std::to_string(Missing->first.value()));
+      }
       return makeUnexpected(ErrorCode::VerificationError,
                             "TracksLiveness block is missing required physical live-ins");
     }
diff --git a/backend/unittests/CodeGen/MachineLivenessTest.cpp b/backend/unittests/CodeGen/MachineLivenessTest.cpp
index cca7629..50e6ffc 100644
--- a/backend/unittests/CodeGen/MachineLivenessTest.cpp
+++ b/backend/unittests/CodeGen/MachineLivenessTest.cpp
@@ -685,6 +685,35 @@ void testLivenessRebuildAndVerifier(TestContext &Context, const rcc::TargetDescr
                  "a second liveness rebuild must repair stale derived state");
 }
 
+void testTrackedLiveThroughPhysicalRegister(TestContext &Context,
+                                            const rcc::TargetDescription &Target) {
+  auto Fixture = createFunction("tracked-live-through", 11, 2);
+  if (!Fixture) {
+    Context.expect(false, "tracked live-through fixture creation must succeed");
+    return;
+  }
+  auto &Function = *Fixture->Function;
+  auto &Predecessor = *Fixture->Blocks[0];
+  auto &Successor = *Fixture->Blocks[1];
+  const rcc::Register G0 = rcc::Register::fromPhysical(physical(rcc::test::TestRegister::G0));
+  auto Edge =
+      Function.cfgUpdater().addSuccessor(Predecessor, Successor, rcc::BranchProbability::one());
+  auto Use = buildUse(Function, Target, Successor, G0);
+  auto LiveIn = Successor.addLiveIn(physical(rcc::test::TestRegister::G0));
+  if (!Edge || !Use || !LiveIn) {
+    Context.expect(false, "tracked live-through MachineIR must be constructed");
+    return;
+  }
+
+  Function.properties().set(rcc::MachineFunctionProperty::TracksLiveness);
+  Context.expect(rcc::verifyTrackedMachineLiveness(Function, Target).has_value() &&
+                     Predecessor.liveIns().empty(),
+                 "a predecessor must not declare a physical register used only by its successor");
+  Successor.clearLiveIns();
+  Context.expect(!rcc::verifyTrackedMachineLiveness(Function, Target),
+                 "a block must declare a physical register that it reads before defining");
+}
+
 void testFailedRebuildIsAtomic(TestContext &Context, const rcc::TargetDescription &Target) {
   auto Fixture = createFunction("failed-rebuild", 9, 2);
   if (!Fixture) {
@@ -933,6 +962,7 @@ int main() {
   testLaneSubrangesAndEarlyClobber(Context, *Target);
   testPhysicalUnitsAndRegMasks(Context, *Target);
   testLivenessRebuildAndVerifier(Context, *Target);
+  testTrackedLiveThroughPhysicalRegister(Context, *Target);
   testFailedRebuildIsAtomic(Context, *Target);
   testAnalysisRegistrationAndInvalidation(Context, *Target);
   testRISCVLiveness(Context);