[llvm-branch-commits] [GlobalISel] Combiner: Observer-based DCE and retrying of combines (PR #102163)

[Amara Emerson via llvm-branch-commits]

================
@@ -45,61 +45,189 @@ cl::OptionCategory GICombinerOptionCategory(
 );
 } // end namespace llvm
 
-/// This class acts as the glue the joins the CombinerHelper to the overall
+/// This class acts as the glue that joins the CombinerHelper to the overall
 /// Combine algorithm. The CombinerHelper is intended to report the
 /// modifications it makes to the MIR to the GISelChangeObserver and the
-/// observer subclass will act on these events. In this case, instruction
-/// erasure will cancel any future visits to the erased instruction and
-/// instruction creation will schedule that instruction for a future visit.
-/// Other Combiner implementations may require more complex behaviour from
-/// their GISelChangeObserver subclass.
+/// observer subclass will act on these events.
 class Combiner::WorkListMaintainer : public GISelChangeObserver {
-  using WorkListTy = GISelWorkList<512>;
-  WorkListTy &WorkList;
+protected:
+#ifndef NDEBUG
   /// The instructions that have been created but we want to report once they
   /// have their operands. This is only maintained if debug output is requested.
-#ifndef NDEBUG
-  SetVector<const MachineInstr *> CreatedInstrs;
+  SmallSetVector<const MachineInstr *, 32> CreatedInstrs;
 #endif
+  using Level = CombinerInfo::ObserverLevel;
 
 public:
-  WorkListMaintainer(WorkListTy &WorkList) : WorkList(WorkList) {}
+  static std::unique_ptr<WorkListMaintainer>
+  create(Level Lvl, WorkListTy &WorkList, MachineRegisterInfo &MRI);
+
   virtual ~WorkListMaintainer() = default;
 
+  void reportFullyCreatedInstrs() {
+    LLVM_DEBUG({
+      for (auto *MI : CreatedInstrs) {
+        dbgs() << "Created: " << *MI;
+      }
+      CreatedInstrs.clear();
+    });
+  }
+
+  virtual void reset() = 0;
+  virtual void appliedCombine() = 0;
+};
+
+/// A configurable WorkListMaintainer implementation.
+/// The ObserverLevel determines how the WorkListMaintainer reacts to MIR
+/// changes.
+template <CombinerInfo::ObserverLevel Lvl>
+class Combiner::WorkListMaintainerImpl : public Combiner::WorkListMaintainer {
+  WorkListTy &WorkList;
+  MachineRegisterInfo &MRI;
+
+  // Defer handling these instructions until the combine finishes.
+  SmallSetVector<MachineInstr *, 32> DeferList;
+
+  // Track VRegs that (might) have lost a use.
+  SmallSetVector<Register, 32> LostUses;
+
+public:
+  WorkListMaintainerImpl(WorkListTy &WorkList, MachineRegisterInfo &MRI)
+      : WorkList(WorkList), MRI(MRI) {}
+
+  virtual ~WorkListMaintainerImpl() = default;
+
+  void reset() override {
+    DeferList.clear();
+    LostUses.clear();
+  }
+
   void erasingInstr(MachineInstr &MI) override {
-    LLVM_DEBUG(dbgs() << "Erasing: " << MI << "\n");
+    // MI will become dangling, remove it from all lists.
+    LLVM_DEBUG(dbgs() << "Erasing: " << MI; CreatedInstrs.remove(&MI));
     WorkList.remove(&MI);
+    if constexpr (Lvl != Level::Basic) {
+      DeferList.remove(&MI);
+      noteLostUses(MI);
+    }
   }
+
   void createdInstr(MachineInstr &MI) override {
-    LLVM_DEBUG(dbgs() << "Creating: " << MI << "\n");
-    WorkList.insert(&MI);
-    LLVM_DEBUG(CreatedInstrs.insert(&MI));
+    LLVM_DEBUG(dbgs() << "Creating: " << MI; CreatedInstrs.insert(&MI));
+    if constexpr (Lvl == Level::Basic)
+      WorkList.insert(&MI);
+    else
+      // Defer handling newly created instructions, because they don't have
+      // operands yet. We also insert them into the WorkList in reverse
+      // order so that they will be combined top down.
+      DeferList.insert(&MI);
   }
+
   void changingInstr(MachineInstr &MI) override {
-    LLVM_DEBUG(dbgs() << "Changing: " << MI << "\n");
-    WorkList.insert(&MI);
+    LLVM_DEBUG(dbgs() << "Changing: " << MI);
+    // Some uses might get dropped when MI is changed.
+    // For now, overapproximate by assuming all uses will be dropped.
+    // TODO: Is a more precise heuristic or manual tracking of use count
+    // decrements worth it?
+    if constexpr (Lvl != Level::Basic)
+      noteLostUses(MI);
   }
+
   void changedInstr(MachineInstr &MI) override {
-    LLVM_DEBUG(dbgs() << "Changed: " << MI << "\n");
-    WorkList.insert(&MI);
+    LLVM_DEBUG(dbgs() << "Changed: " << MI);
+    if constexpr (Lvl == Level::Basic)
+      WorkList.insert(&MI);
+    else
+      // Defer this for DCE
+      DeferList.insert(&MI);
   }
 
-  void reportFullyCreatedInstrs() {
-    LLVM_DEBUG(for (const auto *MI
-                    : CreatedInstrs) {
-      dbgs() << "Created: ";
-      MI->print(dbgs());
-    });
-    LLVM_DEBUG(CreatedInstrs.clear());
+  // Only track changes during the combine and then walk the def/use-chains once
+  // the combine is finished, because:
+  // - instructions might have multiple defs during the combine.
+  // - use counts aren't accurate during the combine.
+  void appliedCombine() override {
+    if constexpr (Lvl == Level::Basic)
+      return;
+
+    // DCE deferred instructions and add them to the WorkList bottom up.
+    while (!DeferList.empty()) {
+      MachineInstr &MI = *DeferList.pop_back_val();
+      if (tryDCE(MI, MRI))
+        continue;
+
+      if constexpr (Lvl >= Level::SinglePass)
+        addUsersToWorkList(MI);
+
+      WorkList.insert(&MI);
+    }
+
+    // Handle instructions that have lost a user.
+    while (!LostUses.empty()) {
+      Register Use = LostUses.pop_back_val();
+      MachineInstr *UseMI = MRI.getVRegDef(Use);
+      if (!UseMI)
+        continue;
+
+      // If DCE succeeds, UseMI's uses are added back to LostUses by
+      // erasingInstr.
+      if (tryDCE(*UseMI, MRI))
+        continue;
+
+      if constexpr (Lvl >= Level::SinglePass) {
+        // OneUse checks are relatively common, so we might be able to combine
+        // the single remaining user of this Reg.
+        if (MRI.hasOneNonDBGUser(Use))
+          WorkList.insert(&*MRI.use_instr_nodbg_begin(Use));
+
+        WorkList.insert(UseMI);
+      }
+    }
+  }
+
+  void noteLostUses(MachineInstr &MI) {
+    for (auto &Use : MI.explicit_uses()) {
+      if (!Use.isReg() || !Use.getReg().isVirtual())
+        continue;
+      LostUses.insert(Use.getReg());
+    }
+  }
+
+  void addUsersToWorkList(MachineInstr &MI) {
+    for (auto &Def : MI.defs()) {
+      Register DefReg = Def.getReg();
+      if (!DefReg.isVirtual())
+        continue;
+      for (auto &UseMI : MRI.use_nodbg_instructions(DefReg)) {
+        WorkList.insert(&UseMI);
+      }
+    }
   }
 };
 
+std::unique_ptr<Combiner::WorkListMaintainer>
+Combiner::WorkListMaintainer::create(Level Lvl, WorkListTy &WorkList,
+                                     MachineRegisterInfo &MRI) {
+  switch (Lvl) {
+  case Level::Basic:
+    return std::make_unique<WorkListMaintainerImpl<Level::Basic>>(WorkList,
+                                                                  MRI);
+  case Level::DCE:
+    return std::make_unique<WorkListMaintainerImpl<Level::DCE>>(WorkList, MRI);
+  case Level::SinglePass:
+    return std::make_unique<WorkListMaintainerImpl<Level::SinglePass>>(WorkList,
+                                                                       MRI);
+  }
+  llvm_unreachable("Illegal ObserverLevel");
----------------
aemerson wrote:

Put this into the default case.

https://github.com/llvm/llvm-project/pull/102163