Preserve loop metadata when splitting exit blocks

LoopUtils.cpp contains a utility that splits an loop exit block, so that the new block contains only edges coming from the loop. In the case of nested loops, the exit path for the inner loop might also be the back-edge of the outer loop. The new block which is inserted on this path, is now a latch for the outer loop, and it needs to hold the loop metadata for the outer loop. (The test case gives a more concrete view of the situation.)

Patch by Chang Lin (clin1)

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

llvm-svn: 346810

1 files changed, 32 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Utils/LoopUtils.cpp b/llvm/lib/Transforms/Utils/LoopUtils.cpp
index 249869e1bde..56b4c448c84 100644
--- a/llvm/lib/Transforms/Utils/LoopUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUtils.cpp
@@ -74,9 +74,41 @@ bool llvm::formDedicatedExitBlocks(Loop *L, DominatorTree *DT, LoopInfo *LI,
     if (IsDedicatedExit)
       return false;
 
+    // With nested loops, the inner loop might exit to the header of an
+    // enclosing loop, and the in-loop-predecessor is a latch for that
+    // enclosing loop. If we insert a block between the latch and the header,
+    // that block becomes the new latch. Any loop metadata from the old latch
+    // needs to be moved to the new one.
+    MDNode *OuterLoopMD = nullptr;
+
+    // If the exit block is a header of a different loop, get that loop's
+    // metadata before we split the block.
+    if (LI->isLoopHeader(BB))
+      OuterLoopMD = LI->getLoopFor(BB)->getLoopID();
+
     auto *NewExitBB = SplitBlockPredecessors(
         BB, InLoopPredecessors, ".loopexit", DT, LI, nullptr, PreserveLCSSA);
 
+    // If OuterLoopMD is non-null, we know that the exit block BB is a
+    // loop header for a different loop, with metadata on its back edges.
+    // If NewExitBB is a member of that loop, then NewExitBB is a latch,
+    // and the loop's metadata needs to be copied to NewExitBB.
+    if (NewExitBB && OuterLoopMD &&
+        LI->getLoopFor(NewExitBB) == LI->getLoopFor(BB)) {
+      // The preds of NewExitBB are all former latches of the outer loop.
+      // Remove their metadata.
+      for (auto *PredLoopBB : InLoopPredecessors) {
+        Instruction *TI = PredLoopBB->getTerminator();
+        // All the latches should have the same metadata (ensured by
+        // getLoopID()).
+        assert(TI->getMetadata(LLVMContext::MD_loop) == OuterLoopMD &&
+               "exit edge to other loop doesn't contain expected metadata");
+        TI->setMetadata(LLVMContext::MD_loop, nullptr);
+      }
+      NewExitBB->getTerminator()->setMetadata(LLVMContext::MD_loop,
+                                              OuterLoopMD);
+    }
+
     if (!NewExitBB)
       LLVM_DEBUG(
           dbgs() << "WARNING: Can't create a dedicated exit block for loop: "