[CodeGen][MachineLoop] Fix getLoopID

Mirror the getLoopID implementation of LoopInfo in MachineLoopInfo.

getLoopID used findLoopControlBlock to detect the special case where there is a
single latch. However, findLoopControlBlock returns the exiting block if the
latch is not an exiting block. The middle end places the LoopID metadata on the
latch regardless of if it's an exiting block or not.

Author: kortbeek-snps

llvm/lib/CodeGen/MachineLoopInfo.cpp

@@ -181,48 +181,32 @@ MachineLoopInfo::findLoopPreheader(MachineLoop *L, bool SpeculativePreheader,
 
 MDNode *MachineLoop::getLoopID() const {
   MDNode *LoopID = nullptr;
-  if (const auto *MBB = findLoopControlBlock()) {
-    // If there is a single latch block, then the metadata
-    // node is attached to its terminating instruction.
+
+  // Go through the latch blocks and check the terminator for the metadata
+  SmallVector<MachineBasicBlock *, 4> LatchesBlocks;
+  getLoopLatches(LatchesBlocks);
+  for (const auto *MBB : LatchesBlocks) {
     const auto *BB = MBB->getBasicBlock();
     if (!BB)
       return nullptr;
-    if (const auto *TI = BB->getTerminator())
-      LoopID = TI->getMetadata(LLVMContext::MD_loop);
-  } else if (const auto *MBB = getHeader()) {
-    // There seem to be multiple latch blocks, so we have to
-    // visit all predecessors of the loop header and check
-    // their terminating instructions for the metadata.
-    if (const auto *Header = MBB->getBasicBlock()) {
-      // Walk over all blocks in the loop.
-      for (const auto *MBB : this->blocks()) {
-        const auto *BB = MBB->getBasicBlock();
-        if (!BB)
-          return nullptr;
-        const auto *TI = BB->getTerminator();
-        if (!TI)
-          return nullptr;
-        MDNode *MD = nullptr;
-        // Check if this terminating instruction jumps to the loop header.
-        for (const auto *Succ : successors(TI)) {
-          if (Succ == Header) {
-            // This is a jump to the header - gather the metadata from it.
-            MD = TI->getMetadata(LLVMContext::MD_loop);
-            break;
-          }
-        }
-        if (!MD)
-          continue;
-        if (!LoopID)
-          LoopID = MD;
-        else if (MD != LoopID)
-          return nullptr;
-      }
-    }
+    const auto *TI = BB->getTerminator();
+    if (!TI)
+      return nullptr;
+    MDNode *MD = TI->getMetadata(LLVMContext::MD_loop);
+
+    if (!MD)
+      return nullptr;
+
+    if (!LoopID)
+      LoopID = MD;
+    else if (MD != LoopID)
+      return nullptr;
   }
-  if (LoopID &&
-      (LoopID->getNumOperands() == 0 || LoopID->getOperand(0) != LoopID))
-    LoopID = nullptr;
+
+  if (!LoopID || LoopID->getNumOperands() == 0 ||
+      LoopID->getOperand(0) != LoopID)
+    return nullptr;
+
   return LoopID;
 }