[CodeGen] Compile-time improvement for Machine Late Instructions Cleanup Pass.

[vpykhtin]

Summary

This is the follow-up for https://github.com/llvm/llvm-project/issues/61397 -
Machine Late Instructions Cleanup Pass takes long time to process huge testcase.

I've made an attempt to improve patch sent by Jonas Paulsson at the github issue, Jonas can you
please try this on your testcases to see if it has any compile-time benefit. Basically it joins
RegDefs and RegKills into single map to reduce lookup time.

I'm submitting original patch from Jonas and then mine so they can be compared for convenience.

[vpykhtin]

Mine version of the patch.

[vpykhtin]

Fixed debug build.

[jonpa]

Thanks for trying the idea of a single map!

On SystemZ, the two versions show practically identical runtimes, so there is no winner. Is there any difference at all on your machine?

I am thinking maybe it would be possible to lookup Reg only once by calling removeRedundantDef() with MI and also the KillMI. That way, clearKillsForDef() wouldn't have to duplicate the same lookup.

[vpykhtin]

On SystemZ, the two versions show practically identical runtimes, so there is no winner. Is there any difference at all on your machine?

Not really :) Something barely measurable like 2.13 and 2.09 seconds on the huge testcase. I'm curious which cases you use to compare to the current, non-map version?

I am thinking maybe it would be possible to lookup Reg only once by calling removeRedundantDef() with MI and also the KillMI. That way, clearKillsForDef() wouldn't have to duplicate the same lookup.

I need to dig more into what is going on there.

[jonpa]

In D146810#4221239, @vpykhtin wrote:

On SystemZ, the two versions show practically identical runtimes, so there is no winner. Is there any difference at all on your machine?

Not really :) Something barely measurable like 2.13 and 2.09 seconds on the huge testcase. I'm curious which cases you use to compare to the current, non-map version?

I am thinking maybe it would be possible to lookup Reg only once by calling removeRedundantDef() with MI and also the KillMI. That way, clearKillsForDef() wouldn't have to duplicate the same lookup.

I need to dig more into what is going on there.

If the map lookup time really matters, maybe it could be worth trying a vector of DefMI/KillMI pairs, so lookup time would be O(n) into the vector..? BTW, I really hope we don't have to use first/second accesses, something like RegInstrs[Reg].DefMI seems more readable...

If it's not possible to speed it up further, I think the simple version I tried is acceptable: it avoids the slowdowns in huge functions and is just barely slower on SPEC, which is then probably preferrable...

[vpykhtin]

If the map lookup time really matters, maybe it could be worth trying a vector of DefMI/KillMI pairs, so lookup time would be O(n) into the vector..?

I don't think it's that critical

BTW, I really hope we don't have to use first/second accesses, something like RegInstrs[Reg].DefMI seems more readable...

RegInstrs[Reg] is a lookup in the map, I don't know if compilers are good enough to optimize several identical map lookups.

If it's not possible to speed it up further, I think the simple version I tried is acceptable: it avoids the slowdowns in huge functions and is just barely slower on SPEC, which is then probably preferrable...

I agree, I'll abandon this review then. Maybe it worth to use SmallDenseMap in your version instead of std::map.

[jonpa]

In D146810#4236313, @vpykhtin wrote:

If the map lookup time really matters, maybe it could be worth trying a vector of DefMI/KillMI pairs, so lookup time would be O(n) into the vector..?

I don't think it's that critical

BTW, I really hope we don't have to use first/second accesses, something like RegInstrs[Reg].DefMI seems more readable...

RegInstrs[Reg] is a lookup in the map, I don't know if compilers are good enough to optimize several identical map lookups.

If it's not possible to speed it up further, I think the simple version I tried is acceptable: it avoids the slowdowns in huge functions and is just barely slower on SPEC, which is then probably preferrable...

I agree, I'll abandon this review then. Maybe it worth to use SmallDenseMap in your version instead of std::map.

I tried this patch again a final time, this time passing the RegInfo pointer to removeRedundantDef() and further into clearKillsForDef(). My idea here is that this would only mean a single lookup for both the def and the kill. This was however also not any faster at all.

I also tried having a vector of RegInfo:s for each MBB - instead of a map. The idea was to make the lookup time linear, but this was much slower in small functions (while perhaps good idea in huge ones).

I have tried different map types before, so I don't think there is anything to gain from that.

I will post my patch as a separate review...

[vpykhtin]

I tried this patch again a final time, this time passing the RegInfo pointer to removeRedundantDef() and further into clearKillsForDef(). My idea here is that this would only mean a single lookup for both the def and the kill. This was however also not any faster at all.

I also tried having a vector of RegInfo:s for each MBB - instead of a map. The idea was to make the lookup time linear, but this was much slower in small functions (while perhaps good idea in huge ones).

I have tried different map types before, so I don't think there is anything to gain from that.

I will post my patch as a separate review...

Ok, thank you, I'll close this one.

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

 
 namespace {
 
 class MachineLateInstrsCleanup : public MachineFunctionPass {
   const TargetRegisterInfo *TRI;
   const TargetInstrInfo *TII;
 
   // Data structures to map regs to their definitions per MBB.
-  using Reg2DefMap = std::map<Register, MachineInstr*>;
-  std::vector<Reg2DefMap> RegDefs;
+  struct RegInfo {
+    MachineInstr *DefMI = nullptr;
+    MachineInstr *KillMI = nullptr;
+  };
+  using Reg2MIMap = SmallDenseMap<Register, RegInfo>;
+  std::vector<Reg2MIMap> BBRegInfo;
 
   // Walk through the instructions in MBB and remove any redundant
   // instructions.
   bool processBlock(MachineBasicBlock *MBB);
 
+  void removeRedundantDef(MachineInstr *MI, const TargetRegisterInfo *TRI);
+  void clearKillsForDef(Register Reg, MachineBasicBlock *MBB,
+                        MachineBasicBlock::iterator I, BitVector &VisitedPreds,
+                        const TargetRegisterInfo *TRI);
+
 public:
   static char ID; // Pass identification, replacement for typeid
 
   MachineLateInstrsCleanup() : MachineFunctionPass(ID) {
     initializeMachineLateInstrsCleanupPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {

 
 bool MachineLateInstrsCleanup::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   TRI = MF.getSubtarget().getRegisterInfo();
   TII = MF.getSubtarget().getInstrInfo();
 
-  RegDefs.clear();
-  RegDefs.resize(MF.getNumBlockIDs());
+  for (auto &RI : BBRegInfo)
+    RI.clear();
+  BBRegInfo.resize(MF.getNumBlockIDs());
 
   // Visit all MBBs in an order that maximises the reuse from predecessors.
   bool Changed = false;
   ReversePostOrderTraversal<MachineFunction *> RPOT(&MF);
   for (MachineBasicBlock *MBB : RPOT)
     Changed |= processBlock(MBB);
 
   return Changed;
 }
 
 // Clear any previous kill flag on Reg found before I in MBB. Walk backwards
 // in MBB and if needed continue in predecessors until a use/def of Reg is
 // encountered. This seems to be faster in practice than tracking kill flags
 // in a map.
-static void clearKillsForDef(Register Reg, MachineBasicBlock *MBB,
+void MachineLateInstrsCleanup::clearKillsForDef(Register Reg,
+                                                MachineBasicBlock *MBB,
                                                 MachineBasicBlock::iterator I,
                                                 BitVector &VisitedPreds,
                                                 const TargetRegisterInfo *TRI) {
   VisitedPreds.set(MBB->getNumber());
-  while (I != MBB->begin()) {
-    --I;
-    bool Found = false;
-    for (auto &MO : I->operands())
-      if (MO.isReg() && TRI->regsOverlap(MO.getReg(), Reg)) {
-        if (MO.isDef())
-          return;
-        if (MO.readsReg()) {
-          MO.setIsKill(false);
-          Found = true; // Keep going for an implicit kill of the super-reg.
-        }
-      }
-    if (Found)
+
+  auto &Map = BBRegInfo[MBB->getNumber()];
+  auto RI = Map.find(Reg);
+  if (RI != Map.end()) {
+    if (RI->second.KillMI) {
+      // Kill flag in MBB
+      RI->second.KillMI->clearRegisterKills(Reg, TRI);
+    // Def in MBB (missing kill flag)
+    if (RI->second.DefMI->getParent() == MBB)
       return;
   }
 
   // If an earlier def is not in MBB, continue in predecessors.
   if (!MBB->isLiveIn(Reg))
     MBB->addLiveIn(Reg);
   assert(!MBB->pred_empty() && "Predecessor def not found!");
   for (MachineBasicBlock *Pred : MBB->predecessors())
     if (!VisitedPreds.test(Pred->getNumber()))
       clearKillsForDef(Reg, Pred, Pred->end(), VisitedPreds, TRI);
 }
 
-static void removeRedundantDef(MachineInstr *MI,
-                               const TargetRegisterInfo *TRI) {
+void MachineLateInstrsCleanup::removeRedundantDef(
+    MachineInstr *MI, const TargetRegisterInfo *TRI) {
   Register Reg = MI->getOperand(0).getReg();
   BitVector VisitedPreds(MI->getMF()->getNumBlockIDs());
   clearKillsForDef(Reg, MI->getParent(), MI->getIterator(), VisitedPreds, TRI);
   MI->eraseFromParent();
   ++NumRemoved;
 }
 
 // Return true if MI is a potential candidate for reuse/removal and if so
if (MO.isReg()) {
                  MO.isGlobal() || MO.isSymbol()))
       return false;
   }
   return DefedReg.isValid();
 }
 
 bool MachineLateInstrsCleanup::processBlock(MachineBasicBlock *MBB) {
   bool Changed = false;
-  Reg2DefMap &MBBDefs = RegDefs[MBB->getNumber()];
+  auto &MBBRegInfo = BBRegInfo[MBB->getNumber()];
 
   // Find reusable definitions in the predecessor(s).
   if (!MBB->pred_empty() && !MBB->isEHPad()) {
     MachineBasicBlock *FirstPred = *MBB->pred_begin();
-    for (auto [Reg, DefMI] : RegDefs[FirstPred->getNumber()])
-      if (llvm::all_of(
-              drop_begin(MBB->predecessors()),
-              [&, &Reg = Reg, &DefMI = DefMI](const MachineBasicBlock *Pred) {
-                auto PredDefI = RegDefs[Pred->getNumber()].find(Reg);
-                return PredDefI != RegDefs[Pred->getNumber()].end() &&
-                       DefMI->isIdenticalTo(*PredDefI->second);
+    for (auto &[Reg, RI] : BBRegInfo[FirstPred->getNumber()]) {
+      if (!RI.DefMI)
+        continue;
+      if (llvm::all_of(drop_begin(MBB->predecessors()),
+                       [&, &Reg = Reg,
+                        &DefMI = RI.DefMI](const MachineBasicBlock *Pred) {
+                         auto &PredMap = BBRegInfo[Pred->getNumber()];
+                         auto PredDefI = PredMap.find(Reg);
+                         return PredDefI != PredMap.end() &&
+                                PredDefI->second.DefMI &&
+                                DefMI->isIdenticalTo(*PredDefI->second.DefMI);
                        })) {
-        MBBDefs[Reg] = DefMI;
+        MBBRegInfo[Reg].DefMI = RI.DefMI;
         LLVM_DEBUG(dbgs() << "Reusable instruction from pred(s): in "
-                          << printMBBReference(*MBB) << ":  " << *DefMI;);
-      }
+                          << printMBBReference(*MBB) << ":  " << *RI.DefMI;);
+      }
   }
 
   // Process MBB.
   MachineFunction *MF = MBB->getParent();
   const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
   Register FrameReg = TRI->getFrameRegister(*MF);
   for (MachineInstr &MI : llvm::make_early_inc_range(*MBB)) {
     // If FrameReg is modified, no previous load-address instructions (using
     // it) are valid.
     if (MI.modifiesRegister(FrameReg, TRI)) {
-      MBBDefs.clear();
+      for (auto I = MBBRegInfo.begin(); I != MBBRegInfo.end();) {
+        auto ThisI = I++;
+        if (!ThisI->second.KillMI)
+          MBBRegInfo.erase(ThisI);
+        else
+          ThisI->second.DefMI = nullptr;
+      }
       continue;
     }
 
     Register DefedReg;
     bool IsCandidate = isCandidate(&MI, DefedReg, FrameReg);
 
     // Check for an earlier identical and reusable instruction.
     if (IsCandidate) {
-      auto DefI = MBBDefs.find(DefedReg);
-      if (DefI != MBBDefs.end() && MI.isIdenticalTo(*DefI->second)) {
+      auto I = MBBRegInfo.find(DefedReg);
+      if (I != MBBRegInfo.end() && I->second.DefMI &&
+          MI.isIdenticalTo(*I->second.DefMI)) {
         LLVM_DEBUG(dbgs() << "Removing redundant instruction in "
                           << printMBBReference(*MBB) << ":  " << MI;);
         removeRedundantDef(&MI, TRI);
         Changed = true;
         continue;
       }
     }
 
     // Clear any entries in map that MI clobbers.
-    for (auto DefI = MBBDefs.begin(); DefI != MBBDefs.end();) {
-      Register Reg = DefI->first;
-      if (MI.modifiesRegister(Reg, TRI))
-        DefI = MBBDefs.erase(DefI);
-      else
-        ++DefI;
+    for (auto I = MBBRegInfo.begin(); I != MBBRegInfo.end();) {
+      auto ThisI = I++;
+      if (!ThisI->second.DefMI)
+        continue;
+      Register Reg = ThisI->first;
+      if (MI.modifiesRegister(Reg, TRI)) {
+        MBBRegInfo.erase(ThisI);
+        continue;
+      }
+      if (MI.findRegisterUseOperandIdx(Reg, false /*isKill*/, TRI) != -1)
+        ThisI->second.KillMI = &MI;
     }
 
     // Record this MI for potential later reuse.
     if (IsCandidate) {
       LLVM_DEBUG(dbgs() << "Found interesting instruction in "
                         << printMBBReference(*MBB) << ":  " << MI;);
-      MBBDefs[DefedReg] = &MI;
+      MBBRegInfo[DefedReg].DefMI = &MI;
     }
   }
 
   return Changed;
 }