[CodeGen] Revisit instructions invalidated by legalization

2 files changed, 72 insertions, 12 deletions

diff --git a/backend/lib/CodeGen/MachinePipelinePasses.cpp b/backend/lib/CodeGen/MachinePipelinePasses.cpp
index 9c42bdb..9bde291 100644
--- a/backend/lib/CodeGen/MachinePipelinePasses.cpp
+++ b/backend/lib/CodeGen/MachinePipelinePasses.cpp
@@ -352,11 +352,13 @@ Expected<MachinePassResult> LegalizerPass::runImplementation(MachineFunction &Fu
   std::unordered_set<OpcodeState, OpcodeStateHash> SeenStates;
 
   std::vector<MachineInstr *> Worklist = layoutInstructions(Function);
-  std::unordered_set<MachineInstrStableID> DiscoveredInstructions;
+  std::unordered_set<MachineInstrStableID> PendingInstructions;
   for (const MachineInstr *Instruction : Worklist)
-    DiscoveredInstructions.insert(Instruction->stableID());
+    PendingInstructions.insert(Instruction->stableID());
   for (std::size_t Index = 0; Index < Worklist.size(); ++Index) {
     MachineInstr *Instruction = Worklist[Index];
+    if (Instruction != nullptr)
+      PendingInstructions.erase(Instruction->stableID());
     if (Instruction == nullptr || Instruction->parent() == nullptr)
       continue;
 
@@ -418,19 +420,16 @@ Expected<MachinePassResult> LegalizerPass::runImplementation(MachineFunction &Fu
     Result.Changed = true;
     Result.Changes.add(MachinePassChange::Instructions);
 
-    if (!InstructionWasErased) {
-      const InstructionDescriptor *LoweredDescriptor =
-          Target.instructionDescriptor(Instruction->opcode());
-      if (LoweredDescriptor == nullptr)
-        return std::unexpected(
-            passFailure("legalizer produced an unknown instruction descriptor"));
-      if (LoweredDescriptor->Flags.has(InstructionFlag::Generic))
-        Worklist.push_back(Instruction);
-    }
     const std::vector<MachineInstr *> Current = layoutInstructions(Function);
     for (MachineInstr *Candidate : Current) {
-      if (DiscoveredInstructions.insert(Candidate->stableID()).second)
+      const InstructionDescriptor *CandidateDescriptor =
+          Target.instructionDescriptor(Candidate->opcode());
+      if (CandidateDescriptor == nullptr)
+        return std::unexpected(passFailure("legalizer produced an unknown instruction descriptor"));
+      if (CandidateDescriptor->Flags.has(InstructionFlag::Generic) &&
+          PendingInstructions.insert(Candidate->stableID()).second) {
         Worklist.push_back(Candidate);
+      }
     }
   }
 
diff --git a/backend/unittests/CodeGen/MachinePipelinePassesTest.cpp b/backend/unittests/CodeGen/MachinePipelinePassesTest.cpp
index 0fe2e80..78c4cc7 100644
--- a/backend/unittests/CodeGen/MachinePipelinePassesTest.cpp
+++ b/backend/unittests/CodeGen/MachinePipelinePassesTest.cpp
@@ -64,6 +64,34 @@ public:
   }
 };
 
+class BackwardInvalidatingLegalizer final : public rcc::TargetLegalizerInfo {
+public:
+  [[nodiscard]] rcc::Expected<rcc::LegalizationAction>
+  classify(const rcc::MachineInstr &Instruction, const rcc::MachineFunction &,
+           const rcc::TargetDescription &) const override {
+    return Instruction.opcode() == rcc::opcode(rcc::GenericOpcode::G_SUB)
+               ? rcc::LegalizationAction::Custom
+               : rcc::LegalizationAction::Legal;
+  }
+
+  [[nodiscard]] rcc::Expected<void> lower(rcc::MachineInstr &Instruction,
+                                          rcc::MachineFunction &Function,
+                                          const rcc::TargetDescription &,
+                                          rcc::LegalizationAction) const override {
+    for (rcc::MachineInstr *Candidate : Function.allInstructions()) {
+      if (Candidate != &Instruction &&
+          Candidate->opcode() == rcc::opcode(rcc::GenericOpcode::G_ADD)) {
+        auto Invalidated = Candidate->setOpcode(rcc::opcode(rcc::GenericOpcode::G_SUB));
+        if (!Invalidated)
+          return std::unexpected(Invalidated.error());
+        break;
+      }
+    }
+    return Instruction.setOpcode(
+        rcc::Opcode::fromValue(static_cast<std::uint32_t>(rcc::riscv::InstructionOpcode::ADD)));
+  }
+};
+
 class MutatingFailingSelector final : public rcc::TargetInstructionSelector {
 public:
   [[nodiscard]] rcc::Expected<void> canSelect(const rcc::MachineInstr &,
@@ -408,6 +436,38 @@ void testLegalizerCycleDiagnostic(TestContext &Context) {
                  "a failed legalization cycle must not publish Legalized");
 }
 
+void testLegalizerRevisitsInvalidatedInstructions(TestContext &Context) {
+  auto Target = rcc::riscv::createTargetDescription();
+  auto Fixture = createFixture(Context, "pipeline_legalizer_revisit", 7);
+  if (!Target || !Fixture.Function || Fixture.Entry == nullptr)
+    return;
+  const rcc::Register Left = createRegister(Context, *Fixture.Function, 64);
+  const rcc::Register Right = createRegister(Context, *Fixture.Function, 64);
+  const rcc::Register FirstResult = createRegister(Context, *Fixture.Function, 64);
+  const rcc::Register SecondResult = createRegister(Context, *Fixture.Function, 64);
+  rcc::MachineIRBuilder Builder(*Fixture.Function, *Target);
+  if (Left.isNoRegister() || Right.isNoRegister() || FirstResult.isNoRegister() ||
+      SecondResult.isNoRegister() || !Builder.setInsertPoint(*Fixture.Entry) ||
+      !buildAdd(Context, Builder, FirstResult, Left, Right) ||
+      !buildAdd(Context, Builder, SecondResult, Left, Right))
+    return;
+  rcc::MachineInstr *First = Fixture.Entry->firstInstruction();
+  rcc::MachineInstr *Second = First == nullptr ? nullptr : First->next();
+  if (Second == nullptr || !Second->setOpcode(rcc::opcode(rcc::GenericOpcode::G_SUB)))
+    return;
+
+  Fixture.Function->properties() = {rcc::MachineFunctionProperty::IsSSA,
+                                    rcc::MachineFunctionProperty::TracksLiveness};
+  BackwardInvalidatingLegalizer Legalizer;
+  rcc::LegalizerPass Pass(Legalizer);
+  rcc::MachineAnalysisManager Analyses;
+  auto Run = Pass.run(*Fixture.Function, *Target, Analyses);
+  const rcc::Opcode LegalOpcode =
+      rcc::Opcode::fromValue(static_cast<std::uint32_t>(rcc::riscv::InstructionOpcode::ADD));
+  Context.expect(Run && First->opcode() == LegalOpcode && Second->opcode() == LegalOpcode,
+                 "legalizer must revisit an earlier instruction invalidated by a later lowering");
+}
+
 void testInstructionSelectionTransaction(TestContext &Context) {
   auto Target = rcc::riscv::createTargetDescription();
   auto Fixture = createFixture(Context, "pipeline_selection_transaction", 8);
@@ -621,6 +681,7 @@ int main() {
   testBankConflict(Context);
   testBankClassOnlyChange(Context);
   testLegalizerCycleDiagnostic(Context);
+  testLegalizerRevisitsInvalidatedInstructions(Context);
   testInstructionSelectionTransaction(Context);
   testInstructionSelectionOrder(Context);
   testInstructionSelectionWorklist(Context);