[ARM] Allow findLoopPreheader to return headers with multiple loop successors

[dmgreen]

Summary

The findLoopPreheader function will currently not find a preheader if it branches to multiple different loop headers. This patch adds an option to relax that, allowing ARMLowOverheadLoops to process more loops successfully. This helps with WhileLoopStart setup instructions that can branch/fallthrough to the low overhead loop and to branch to a separate loop from the same preheader (but I don't believe it is possible for both loops to be low overhead loops).

[samtebbs]

llvm/include/llvm/CodeGen/MachineLoopInfo.h

114

[samtebbs]

We're go for committing after fixing the little nitpick in the comment :)

[dmgreen]

Thanks

llvm/include/llvm/CodeGen/MachineLoopInfo.h

public:
 
   LoopInfoBase<MachineBasicBlock, MachineLoop>& getBase() { return LI; }
 
   /// Find the block that either is the loop preheader, or could
   /// speculatively be used as the preheader. This is e.g. useful to place
   /// loop setup code. Code that cannot be speculated should not be placed
   /// here. SpeculativePreheader is controlling whether it also tries to
   /// find the speculative preheader if the regular preheader is not present.
-  MachineBasicBlock *findLoopPreheader(MachineLoop *L,
-                                       bool SpeculativePreheader = false) const;
+  /// With FindMultiLoopPreheader = false, nullptr will be returned if the found
+  /// preheader the preheader of multiple loops.
+  MachineBasicBlock *
+  findLoopPreheader(MachineLoop *L, bool SpeculativePreheader = false,
+                    bool FindMultiLoopPreheader = false) const;
 
   /// The iterator interface to the top-level loops in the current function.
   using iterator = LoopInfoBase<MachineBasicBlock, MachineLoop>::iterator;
   inline iterator begin() const { return LI.begin(); }
   inline iterator end() const { return LI.end(); }
   bool empty() const { return LI.empty(); }
 
   /// Return the innermost loop that BB lives in. If a basic block is in no loop

llvm/lib/CodeGen/MachineLoopInfo.cpp

DebugLoc MachineLoop::getStartLoc() const {
   if (MachineBasicBlock *HeadMBB = getHeader())
     if (const BasicBlock *HeadBB = HeadMBB->getBasicBlock())
       return HeadBB->getTerminator()->getDebugLoc();
 
   return DebugLoc();
 }
 
 MachineBasicBlock *
-MachineLoopInfo::findLoopPreheader(MachineLoop *L,
-                                   bool SpeculativePreheader) const {
+MachineLoopInfo::findLoopPreheader(MachineLoop *L, bool SpeculativePreheader,
+                                   bool FindMultiLoopPreheader) const {
   if (MachineBasicBlock *PB = L->getLoopPreheader())
     return PB;
 
   if (!SpeculativePreheader)
     return nullptr;
 
   MachineBasicBlock *HB = L->getHeader(), *LB = L->getLoopLatch();
   if (HB->pred_size() != 2 || HB->hasAddressTaken())
     return nullptr;
   // Find the predecessor of the header that is not the latch block.
   MachineBasicBlock *Preheader = nullptr;
   for (MachineBasicBlock *P : HB->predecessors()) {
     if (P == LB)
       continue;
     // Sanity.
     if (Preheader)
       return nullptr;
     Preheader = P;
   }
 
   // Check if the preheader candidate is a successor of any other loop
   // headers. We want to avoid having two loop setups in the same block.
+  if (!FindMultiLoopPreheader) {
     for (MachineBasicBlock *S : Preheader->successors()) {
       if (S == HB)
         continue;
       MachineLoop *T = getLoopFor(S);
       if (T && T->getHeader() == S)
         return nullptr;
     }
+  }
   return Preheader;
 }
 
 bool MachineLoop::isLoopInvariant(MachineInstr &I) const {
   MachineFunction *MF = I.getParent()->getParent();
   MachineRegisterInfo *MRI = &MF->getRegInfo();
   const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
 

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

void ProcessLoop() {
         [this, &GetPredecessor] (MachineBasicBlock *MBB) -> void {
         Order.push_back(MBB);
         if (MBB->pred_size() == 1)
           GetPredecessor(*MBB->pred_begin());
       };
 
       if (auto *Preheader = ML.getLoopPreheader())
         GetPredecessor(Preheader);
-      else if (auto *Preheader = MLI.findLoopPreheader(&ML, true))
+      else if (auto *Preheader = MLI.findLoopPreheader(&ML, true, true))
         GetPredecessor(Preheader);
     }
   };
 
   struct PredicatedMI {
     MachineInstr *MI = nullptr;
     SetVector<MachineInstr*> Predicates;
 
struct LowOverheadLoop {
     LowOverheadLoop(MachineLoop &ML, MachineLoopInfo &MLI,
                     ReachingDefAnalysis &RDA, const TargetRegisterInfo &TRI,
                     const ARMBaseInstrInfo &TII)
         : ML(ML), MLI(MLI), RDA(RDA), TRI(TRI), TII(TII),
           TPNumElements(MachineOperand::CreateImm(0)) {
       MF = ML.getHeader()->getParent();
       if (auto *MBB = ML.getLoopPreheader())
         Preheader = MBB;
-      else if (auto *MBB = MLI.findLoopPreheader(&ML, true))
+      else if (auto *MBB = MLI.findLoopPreheader(&ML, true, true))
         Preheader = MBB;
       VPTState::reset();
     }
 
     // If this is an MVE instruction, check that we know how to use tail
     // predication with it. Record VPT blocks and return whether the
     // instruction is valid for tail predication.
     bool ValidateMVEInst(MachineInstr *MI);

 bool ARMLowOverheadLoops::ProcessLoop(MachineLoop *ML) {
 
   bool Changed = false;
 
   // Process inner loops first.
   for (auto I = ML->begin(), E = ML->end(); I != E; ++I)
     Changed |= ProcessLoop(*I);
 
-  LLVM_DEBUG(dbgs() << "ARM Loops: Processing loop containing:\n";
-             if (auto *Preheader = ML->getLoopPreheader())
-               dbgs() << " - " << Preheader->getName() << "\n";
-             else if (auto *Preheader = MLI->findLoopPreheader(ML))
-               dbgs() << " - " << Preheader->getName() << "\n";
-             else if (auto *Preheader = MLI->findLoopPreheader(ML, true))
-               dbgs() << " - " << Preheader->getName() << "\n";
-             for (auto *MBB : ML->getBlocks())
-               dbgs() << " - " << MBB->getName() << "\n";
-            );
+  LLVM_DEBUG({
+    dbgs() << "ARM Loops: Processing loop containing:\n";
+      dbgs() << " - Preheader: " << printMBBReference(*Preheader) << "\n";
+    else if (auto *Preheader = MLI->findLoopPreheader(ML, true, true))
+      dbgs() << " - Preheader: " << printMBBReference(*Preheader) << "\n";
+      dbgs() << " - Block: " << printMBBReference(*MBB) << "\n";
+  });
 
   // Search the given block for a loop start instruction. If one isn't found,
   // and there's only one predecessor block, search that one too.
   std::function<MachineInstr*(MachineBasicBlock*)> SearchForStart =
     [&SearchForStart](MachineBasicBlock *MBB) -> MachineInstr* {
     for (auto &MI : *MBB) {
       if (isLoopStart(MI))
         return &MI;

llvm/test/CodeGen/Thumb2/mve-memtp-loop.ll

 ; CHECK-LABEL: test11:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    .save {r4, lr}
 ; CHECK-NEXT:    push {r4, lr}
 ; CHECK-NEXT:    cmp.w r2, #-1
 ; CHECK-NEXT:    it gt
 ; CHECK-NEXT:    popgt {r4, pc}
 ; CHECK-NEXT:  .LBB10_1: @ %prehead
-; CHECK-NEXT:    add.w r3, r2, #15
 ; CHECK-NEXT:    mov r12, r1
 ; CHECK-NEXT:    mov r4, r0
-; CHECK-NEXT:    lsr.w lr, r3, #4
 ; CHECK-NEXT:    mov r3, r2
-; CHECK-NEXT:    subs.w lr, lr, #0
-; CHECK-NEXT:    beq .LBB10_3
-; CHECK-NEXT:    b .LBB10_2
+; CHECK-NEXT:    wlstp.8 lr, r3, .LBB10_3
 ; CHECK-NEXT:  .LBB10_2: @ =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    vctp.8 r3
-; CHECK-NEXT:    subs r3, #16
-; CHECK-NEXT:    vpstt
-; CHECK-NEXT:    vldrbt.u8 q0, [r12], #16
-; CHECK-NEXT:    vstrbt.8 q0, [r4], #16
-; CHECK-NEXT:    subs.w lr, lr, #1
-; CHECK-NEXT:    bne .LBB10_2
-; CHECK-NEXT:    b .LBB10_3
+; CHECK-NEXT:    vldrb.u8 q0, [r12], #16
+; CHECK-NEXT:    vstrb.8 q0, [r4], #16
+; CHECK-NEXT:    letp lr, .LBB10_2
 ; CHECK-NEXT:  .LBB10_3: @ %for.body
 ; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    ldrb r3, [r0], #1
 ; CHECK-NEXT:    subs r2, #2
 ; CHECK-NEXT:    strb r3, [r1], #1
 ; CHECK-NEXT:    bne .LBB10_3
 ; CHECK-NEXT:  @ %bb.4: @ %for.cond.cleanup
 ; CHECK-NEXT:    pop {r4, pc}