[RDA][ARM] collectKilledOperands across multiple blocks

[samparker]

Summary

Use MIOperand in collectLocalKilledOperands to make the search global, as we already have to search for global uses too. This allows us to delete more dead code when tail predicating.

[samparker]

Rebased and changed the way we collect the IT blocks.

[samparker]

• Fixed a bug in getUniqueReachingMIDef.
• During operand collection, visit all uses and not just killed ones.

[SjoerdMeijer]

Nice optimisation, looks very reasonable to me: LGTM

llvm/include/llvm/CodeGen/ReachingDefAnalysis.h

public:
   /// Return whether From can be moved forwards to just before To.
   bool isSafeToMoveForwards(MachineInstr *From, MachineInstr *To) const;
 
   /// Return whether From can be moved backwards to just after To.
   bool isSafeToMoveBackwards(MachineInstr *From, MachineInstr *To) const;
 
   /// Assuming MI is dead, recursively search the incoming operands which are
   /// killed by MI and collect those that would become dead.
-  void collectLocalKilledOperands(MachineInstr *MI, InstSet &Dead) const;
+  void collectKilledOperands(MachineInstr *MI, InstSet &Dead) const;
 
   /// Return whether removing this instruction will have no effect on the
   /// program, returning the redundant use-def chain.
   bool isSafeToRemove(MachineInstr *MI, InstSet &ToRemove) const;
 
   /// Return whether removing this instruction will have no effect on the
   /// program, ignoring the possible effects on some instructions, returning
   /// the redundant use-def chain.

llvm/lib/CodeGen/ReachingDefAnalysis.cpp

       if (!isSafeToRemove(I, Visited, ToRemove, Ignore))
         return false;
     }
   }
   ToRemove.insert(MI);
   return true;
 }
 
-void ReachingDefAnalysis::collectLocalKilledOperands(MachineInstr *MI,
-                                                     InstSet &Dead) const {
+void ReachingDefAnalysis::collectKilledOperands(MachineInstr *MI,
   Dead.insert(MI);
   auto IsDead = [this, &Dead](MachineInstr *Def, int PhysReg) {
     unsigned LiveDefs = 0;
     for (auto &MO : Def->operands()) {
       if (!isValidRegDef(MO))
         continue;
       if (!MO.isDead())
         ++LiveDefs;
     }
 
     if (LiveDefs > 1)
       return false;
 
     SmallPtrSet<MachineInstr*, 4> Uses;
     getGlobalUses(Def, PhysReg, Uses);
     for (auto *Use : Uses)
       if (!Dead.count(Use))
         return false;
     return true;
   };
 
   for (auto &MO : MI->operands()) {
     if (!isKilledRegUse(MO))
       continue;
-    if (MachineInstr *Def = getReachingLocalMIDef(MI, MO.getReg()))
+    if (MachineInstr *Def = getMIOperand(MI, MO))
       if (IsDead(Def, MO.getReg()))
-        collectLocalKilledOperands(Def, Dead);
+        collectKilledOperands(Def, Dead);
   }
 }
 
 bool ReachingDefAnalysis::isSafeToDefRegAt(MachineInstr *MI,
                                            int PhysReg) const {
   SmallPtrSet<MachineInstr*, 1> Ignore;
   return isSafeToDefRegAt(MI, PhysReg, Ignore);
 }

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

 void ARMLowOverheadLoops::IterationCountDCE(LowOverheadLoop &LoLoop) {
   if (!LoLoop.IsTailPredicationLegal())
     return;
 
   MachineInstr *Def = RDA->getMIOperand(LoLoop.Start, 0);
   if (!Def)
     return;
 
-  // Collect IT blocks.
+  // Collect the dead code and the MBBs in which they reside.
+  SmallPtrSet<MachineInstr*, 4> Killed;
+  RDA->collectKilledOperands(Def, Killed);
+  SmallPtrSet<MachineBasicBlock*, 2> BasicBlocks;
+  for (auto *MI : Killed)
+    BasicBlocks.insert(MI->getParent());
+
+  // Collect IT blocks in all affected basic blocks.
   std::map<MachineInstr *, SmallPtrSet<MachineInstr *, 2>> ITBlocks;
-  std::map<MachineInstr *, MachineInstr *> Predicates;
-  MachineInstr *IT = nullptr;
-  for (auto &MI : *Def->getParent()) {
-    if (MI.getOpcode() == ARM::t2IT)
-      IT = &MI;
-    else if (TII->getPredicate(MI) != ARMCC::AL) {
-      ITBlocks[IT].insert(&MI);
-      Predicates[&MI] = IT;
+  for (auto *MBB : BasicBlocks) {
+    for (auto &MI : *MBB) {
+      if (MI.getOpcode() != ARM::t2IT)
+        continue;
+      RDA->getReachingLocalUses(&MI, ARM::ITSTATE, ITBlocks[&MI]);
     }
   }
 
   // If we're removing all of the instructions within an IT block, then
   // also remove the IT instruction.
   SmallPtrSet<MachineInstr*, 2> ModifiedITs;
-  SmallPtrSet<MachineInstr*, 2> DeadITs;
-  SmallPtrSet<MachineInstr*, 4> Killed;
-  RDA->collectLocalKilledOperands(Def, Killed);
   for (auto *MI : Killed) {
-    if (!Predicates.count(MI))
-      continue;
-
-    MachineInstr *IT = Predicates[MI];
-    auto &CurrentBlock = ITBlocks[IT];
-    CurrentBlock.erase(MI);
-    ModifiedITs.insert(IT);
-    if (CurrentBlock.empty()) {
-      DeadITs.insert(IT);
-      ModifiedITs.erase(IT);
+    if (MachineOperand *MO = MI->findRegisterUseOperand(ARM::ITSTATE)) {
+      MachineInstr *IT = RDA->getMIOperand(MI, *MO);
+      auto &CurrentBlock = ITBlocks[IT];
+      CurrentBlock.erase(MI);
+      if (CurrentBlock.empty())
+        ModifiedITs.erase(IT);
+      else
+        ModifiedITs.insert(IT);
     }
   }
 
   // Delete the killed instructions only if we don't have any IT blocks that
   // need to be modified because we need to fixup the mask.
   // TODO: Handle cases where IT blocks are modified.
   if (ModifiedITs.empty()) {
     LLVM_DEBUG(dbgs() << "ARM Loops: Will remove iteration count:\n";

llvm/test/CodeGen/Thumb2/LowOverheadLoops/it-block-mov.mir

body:             |
   ; CHECK-LABEL: name: cond_trip_count
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
   ; CHECK:   liveins: $lr, $r0, $r1, $r2, $r4
   ; CHECK:   frame-setup tPUSH 14 /* CC::al */, $noreg, killed $r4, killed $lr, implicit-def $sp, implicit $sp
   ; CHECK:   frame-setup CFI_INSTRUCTION def_cfa_offset 8
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $lr, -4
   ; CHECK:   frame-setup CFI_INSTRUCTION offset $r4, -8
-  ; CHECK:   tCMPi8 renamable $r1, 4, 14 /* CC::al */, $noreg, implicit-def $cpsr
-  ; CHECK:   renamable $r3 = t2MOVi 4, 14 /* CC::al */, $noreg, $noreg
-  ; CHECK:   t2IT 11, 8, implicit-def $itstate
-  ; CHECK:   dead $r3 = tMOVr renamable $r1, 11 /* CC::lt */, killed $cpsr, implicit killed renamable $r3, implicit killed $itstate
   ; CHECK:   tCMPi8 renamable $r1, 2, 14 /* CC::al */, $noreg, implicit-def $cpsr
   ; CHECK:   renamable $r12 = t2MOVi 4, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   tBcc %bb.2, 2 /* CC::hs */, killed $cpsr
   ; CHECK: bb.1:
   ; CHECK:   liveins: $r2
   ; CHECK:   renamable $s0 = VLDRS %const.0, 0, 14 /* CC::al */, $noreg
   ; CHECK:   VSTRS killed renamable $s0, killed renamable $r2, 0, 14 /* CC::al */, $noreg
   ; CHECK:   tPOP_RET 14 /* CC::al */, $noreg, def $r4, def $pc