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

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).

Differential Revision: https://reviews.llvm.org/D102747

include/llvm/CodeGen/MachineLoopInfo.h

@@ -110,8 +110,11 @@ public:
   /// 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 is 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;

lib/CodeGen/MachineLoopInfo.cpp

@@ -115,8 +115,8 @@ DebugLoc MachineLoop::getStartLoc() const {
 }
 
 MachineBasicBlock *
-MachineLoopInfo::findLoopPreheader(MachineLoop *L,
-                                   bool SpeculativePreheader) const {
+MachineLoopInfo::findLoopPreheader(MachineLoop *L, bool SpeculativePreheader,
+                                   bool FindMultiLoopPreheader) const {
   if (MachineBasicBlock *PB = L->getLoopPreheader())
     return PB;
 
@@ -139,12 +139,14 @@ MachineLoopInfo::findLoopPreheader(MachineLoop *L,
 
   // 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.
-  for (MachineBasicBlock *S : Preheader->successors()) {
-    if (S == HB)
-      continue;
-    MachineLoop *T = getLoopFor(S);
-    if (T && T->getHeader() == S)
-      return nullptr;
+  if (!FindMultiLoopPreheader) {
+    for (MachineBasicBlock *S : Preheader->successors()) {
+      if (S == HB)
+        continue;
+      MachineLoop *T = getLoopFor(S);
+      if (T && T->getHeader() == S)
+        return nullptr;
+    }
   }
   return Preheader;
 }

lib/Target/ARM/ARMLowOverheadLoops.cpp

@@ -158,7 +158,7 @@ namespace {
 
       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);
     }
   };
@@ -386,7 +386,7 @@ namespace {
       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();
     }
@@ -1197,16 +1197,15 @@ bool ARMLowOverheadLoops::ProcessLoop(MachineLoop *ML) {
   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";
+    if (auto *Preheader = ML->getLoopPreheader())
+      dbgs() << " - Preheader: " << printMBBReference(*Preheader) << "\n";
+    else if (auto *Preheader = MLI->findLoopPreheader(ML, true, true))
+      dbgs() << " - Preheader: " << printMBBReference(*Preheader) << "\n";
+    for (auto *MBB : ML->getBlocks())
+      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.

test/CodeGen/Thumb2/mve-memtp-branch.ll

@@ -117,10 +117,7 @@ define i32 @a(i8 zeroext %b, [3 x i8]* nocapture readonly %c, [3 x i32]* nocaptu
 ; CHECK-NEXT:    add.w r5, r12, r3
 ; CHECK-NEXT:    rsb.w r3, r3, #108
 ; CHECK-NEXT:    add.w r4, r5, #19
-; CHECK-NEXT:    add.w r5, r3, #15
-; CHECK-NEXT:    lsrs r5, r5, #4
-; CHECK-NEXT:    subs.w lr, r5, #0
-; CHECK-NEXT:    beq .LBB0_13
+; CHECK-NEXT:    wlstp.8 lr, r3, .LBB0_13
 ; CHECK-NEXT:    b .LBB0_23
 ; CHECK-NEXT:  .LBB0_13: @ %for.cond
 ; CHECK-NEXT:    @ =>This Loop Header: Depth=1
@@ -190,12 +187,8 @@ define i32 @a(i8 zeroext %b, [3 x i8]* nocapture readonly %c, [3 x i32]* nocaptu
 ; CHECK-NEXT:    b .LBB0_12
 ; CHECK-NEXT:  .LBB0_23: @ Parent Loop BB0_13 Depth=1
 ; CHECK-NEXT:    @ => This Inner Loop Header: Depth=2
-; CHECK-NEXT:    vctp.8 r3
-; CHECK-NEXT:    subs r3, #16
-; CHECK-NEXT:    vpst
-; CHECK-NEXT:    vstrbt.8 q3, [r4], #16
-; CHECK-NEXT:    subs.w lr, lr, #1
-; CHECK-NEXT:    bne .LBB0_23
+; CHECK-NEXT:    vstrb.8 q3, [r4], #16
+; CHECK-NEXT:    letp lr, .LBB0_23
 ; CHECK-NEXT:    b .LBB0_13
 entry:
   %cmp = icmp ugt i8 %b, 1

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

@@ -233,23 +233,14 @@ define void @test11(i8* nocapture %x, i8* nocapture %y, i32 %n) {
 ; 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