[MachineLateInstrsCleanup] Improve compile time for huge functions.

[jonpa]

Summary

It was discovered that this pass can be slow on huge functions, meaning 20% compile time instead of the usual ~0.5% (with a test case spending ~19 mins on SystemZ just in the backend with llc).

The problem relates to the necessary clearing of the kill flags when a redundant instruction is removed. This was made by scanning backwards over the instructions until a use(/def) was found (I tried simply clearing all kill flags on the involved registers, but that did change a few cases of later optimizations, so it seems that this should not be simplified like that).

This patch replaces the scanning of instructions with a map from register to the last seen use (kill).

Compile time on the huge file has been remedied fully (-> 0.4%), but I see on SystemZ/SPEC a slight average compile time increase (0.46% -> 0.52%). This is small enough that it seems wise to go ahead and use the map to help the huge functions (or..?).

Some alternate experiments with this has already been made at https://reviews.llvm.org/D146810, but this turned out to be the simpler version with the same performance.

Fixes: https://github.com/llvm/llvm-project/issues/61397

[vpykhtin]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

46

If the number of map pairs is usually small I would use SmallDenseMap as it has pre-allocated storage, especially when you're placing it in a vector

53

Capitalize mi to match coding style

[craig.topper]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

47

this just std::map::lookup?

54

return mi && mi->isIdenticalTo(*MI)

[craig.topper]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

47

Nevermind. I guess lookup is unique to DenseMap.

[jonpa]

Updated per review -- see inline comment about the map type.

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

46

I tried replacing for SmallDenseMap here again, but no change in compile time at all that I could see.

[vpykhtin]

LGTM, thank you!

[jonpa]

In D147532#4250535, @vpykhtin wrote:

LGTM, thank you!

Thanks, Valery!

I'll wait a few days before committing in case there are any comments from anyone else on this...

[arsenm]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

46

I thought subclassing STL types was discouraged. I'd also default to using the LLVM map types if there's no observable difference anyway

[jonpa]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

46

My viewpoint was that if we use an LLVM map it looks like it's faster than STL, but if it isn't, we could use STL for now until we find a better way, and in the future it will be clear that this has not been improved upon yet.

I am fine either way, so I could change it if that's what the rest of you want...

[jonpa]

Use SmallDenseMap instead.

Iteration change necessary due to this (at "// Clear any entries in map that MI clobbers."), so you might as well take another look...

[jonpa]

Everybody happy with this?

[vpykhtin]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

240

Sorry, I didn't notice this change first.

You can still use erase by iterator (it's slightly faster than another lookup) because you're using make_early_inc_range.

[vpykhtin]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

240

bad advice - DefI is already incremented after Register Reg = DefI.first.

May be you can leave previous variant of iteration, you can increment the iterator like this:

MBBDefs.erase(DefI++) instead of DefI = MBBDefs.erase(DefI)

[jonpa]

Sorry for delay (vaccation). Patch updated -- see inline comments.

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

240

Yeah - I went for readability at first here as I did not see any change in compile time.

240

I did not find anything to convince me 100% that this use of erase with iterator is legal, so I will have to trust you on this ;-)

I think the readability is ok now as well with this.

[vpykhtin]

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

240

Now you're incrementing DefI after erase and in general this is incorrect for STL containers (though it's not the case here) because erase invalidates iterators to erased elements.

MBBDefs.erase(DefI++) works because arguments are fully evaluated before function call, so this is equivalent to:

auto DefIBefore = DefI;
++DefI; // Now DefI is valid after erase
MBBDefs.erase(DefIBefore);

[jonpa]

It seems that this would result in this loop:

for (auto DefI = MBBDefs.begin(); DefI != MBBDefs.end();) {
  Register Reg = DefI->first;
  if (MI.modifiesRegister(Reg, TRI)) {
    MBBDefs.erase(DefI++);
    MBBKills.erase(Reg);
    continue;
  } else if (MI.findRegisterUseOperandIdx(Reg, false /*isKill*/, TRI) != -1)
    // Keep track of the last use seen so far.
    MBBKills[Reg] = &MI;
  DefI++;
}

I find that acceptable, but the iteration logic is slightly complicated. That's why I chose to ignore the extra lookup when calling erase(Reg) in the first place:

for (auto DefI : llvm::make_early_inc_range(MBBDefs)) {
  Register Reg = DefI.first;
  if (MI.modifiesRegister(Reg, TRI)) {
    MBBDefs.erase(Reg);
    MBBKills.erase(Reg);
  } else if (MI.findRegisterUseOperandIdx(Reg, false /*isKill*/, TRI) != -1)
    // Keep track of the last use seen so far.
    MBBKills[Reg] = &MI;
}

I compared both these versions just now on a z15 with your test case, but found no difference:

      ---User Time---   --System Time--   --User+System--   ---Wall Time---  --- Name ---
itr:   4.8783 (  0.4%)   0.0011 (  0.0%)   4.8794 (  0.4%)   4.8833 (  0.4%)  Machine Late Instructions Cleanup Pass
reg:   4.8384 (  0.4%)   0.0009 (  0.0%)   4.8393 (  0.4%)   4.8439 (  0.4%)  Machine Late Instructions Cleanup Pass

I guess this could mean that this test case does a lot of reuse of the earlier Candidate, so it doesn't call erase() here that much, maybe.

On SPEC/SystemZ they both average 0.52% compile time across all files, so there is no difference there either.

The situation may be different on other systems, I suppose. I was using GCC version used to build clang: 12.2.1. Do you see any compile time difference at all on your platform?

I would personally vote for the readability (ignore the erase(Reg) potential overhead), but I will let you the reviewer(s) decide :)

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

240

Sorry,I thought you were implying that increment after erase() was fine with this SmallDenseMap...

[vpykhtin]

Ok, since this isn't critical I don't have strong preference to erase by the iterator, I'm fine with shorter and nicer version with make_early_inc_range :) Thank you :)

[jonpa]

In D147532#4314784, @vpykhtin wrote:

Ok, since this isn't critical I don't have strong preference to erase by the iterator, I'm fine with shorter and nicer version with make_early_inc_range :) Thank you :)

Thanks for review!

I pushed it, but ran into a buildbot failure, which I fixed with 10f0158.

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

 STATISTIC(NumRemoved, "Number of redundant instructions removed.");
 
 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;
+  // Data structures to map regs to their definitions and kills per MBB.
+  struct Reg2MIMap : public std::map<Register, MachineInstr *> {
+    MachineInstr *get(Register Reg) {
+      auto I = find(Reg);
+      return I != end() ? I->second : nullptr;
+    }
+
+    bool hasIdentical(Register Reg, MachineInstr *ArgMI) {
+      MachineInstr *MI = get(Reg);
+      return MI && MI->isIdenticalTo(*ArgMI);
+    }
+  };
+
+  std::vector<Reg2MIMap> RegDefs;
+  std::vector<Reg2MIMap> RegKills;
 
   // Walk through the instructions in MBB and remove any redundant
   // instructions.
   bool processBlock(MachineBasicBlock *MBB);
 
+  void removeRedundantDef(MachineInstr *MI);
+  void clearKillsForDef(Register Reg, MachineBasicBlock *MBB,
+                        MachineBasicBlock::iterator I,
+                        BitVector &VisitedPreds);
+
 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());
+  RegKills.clear();
+  RegKills.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,
-                             MachineBasicBlock::iterator I,
-                             BitVector &VisitedPreds,
-                             const TargetRegisterInfo *TRI) {
+void MachineLateInstrsCleanup::
+clearKillsForDef(Register Reg, MachineBasicBlock *MBB,
+                 BitVector &VisitedPreds) {
   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)
+
+  // Kill flag in MBB
+  if (MachineInstr *KillMI = RegKills[MBB->getNumber()].get(Reg)) {
+    KillMI->clearRegisterKills(Reg, TRI);
+
+  // Def in MBB (missing kill flag)
+  if (MachineInstr *DefMI = RegDefs[MBB->getNumber()].get(Reg))
+    if (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);
+      clearKillsForDef(Reg, Pred, Pred->end(), VisitedPreds);
 }
 
-static void removeRedundantDef(MachineInstr *MI,
-                               const TargetRegisterInfo *TRI) {
+void MachineLateInstrsCleanup::removeRedundantDef(MachineInstr *MI) {
   Register Reg = MI->getOperand(0).getReg();
   BitVector VisitedPreds(MI->getMF()->getNumBlockIDs());
-  clearKillsForDef(Reg, MI->getParent(), MI->getIterator(), VisitedPreds, TRI);
+  clearKillsForDef(Reg, MI->getParent(), MI->getIterator(), VisitedPreds);
   MI->eraseFromParent();
   ++NumRemoved;
 }
 
 // Return true if MI is a potential candidate for reuse/removal and if so
 // also the register it defines in DefedReg.  A candidate is a simple
 // instruction that does not touch memory, has only one register definition
 // and the only reg it may use is FrameReg. Typically this is an immediate
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()];
+  Reg2MIMap &MBBDefs = RegDefs[MBB->getNumber()];
+  Reg2MIMap &MBBKills = RegKills[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);
+                return RegDefs[Pred->getNumber()].hasIdentical(Reg, DefMI);
               })) {
         MBBDefs[Reg] = DefMI;
         LLVM_DEBUG(dbgs() << "Reusable instruction from pred(s): in "
                           << printMBBReference(*MBB) << ":  " << *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();
+      MBBKills.clear();
       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)) {
-        LLVM_DEBUG(dbgs() << "Removing redundant instruction in "
-                          << printMBBReference(*MBB) << ":  " << MI;);
-        removeRedundantDef(&MI, TRI);
-        Changed = true;
-        continue;
-      }
-    }
+    if (IsCandidate && MBBDefs.hasIdentical(DefedReg, &MI)) {
+      removeRedundantDef(&MI);
 
     // 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))
+      if (MI.modifiesRegister(Reg, TRI)) {
         DefI = MBBDefs.erase(DefI);
-      else
-        ++DefI;
+        MBBKills.erase(Reg);
+        continue;
+      }
+
+      // Keep track of the last use seen so far.
+      if (MI.findRegisterUseOperandIdx(Reg, false /*isKill*/, TRI) != -1)
+        MBBKills[Reg] = &MI;
+
     }
 
     // Record this MI for potential later reuse.
     if (IsCandidate) {
       LLVM_DEBUG(dbgs() << "Found interesting instruction in "
                         << printMBBReference(*MBB) << ":  " << MI;);
       MBBDefs[DefedReg] = &MI;
+      assert(!MBBKills.count(DefedReg) && "Should already have been removed.");
     }
   }
 
   return Changed;
 }