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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@346810 91177308-0d34-0410-b5e6-96231b3b80d8

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: "

test/Transforms/LoopSimplify/preserve-llvm-loop-metadata2.ll

@@ -0,0 +1,48 @@
+; RUN: opt -S -loop-simplify < %s | FileCheck %s
+
+; Two-loop nest with llvm.loop metadata on each loop.
+; inner.header exits to outer.header. inner.header is a latch for the outer
+; loop, and contains the outer loop's metadata.
+; After loop-simplify, a new block "outer.header.loopexit" is created between
+; inner.header and outer.header. The metadata from inner.header must be moved
+; to the new block, as the new block becomes the outer loop latch.
+; The metadata on the inner loop's latch should be untouched.
+
+; CHECK: outer.header.loopexit:
+; CHECK-NEXT: llvm.loop [[UJAMTAG:.*]]
+; CHECK-NOT: br i1 {{.*}}, label {{.*}}, label %outer.header.loopexit, !llvm.loop
+; CHECK: br label %inner.header, !llvm.loop [[UNROLLTAG:.*]]
+
+; CHECK: distinct !{[[UJAMTAG]], [[UJAM:.*]]}
+; CHECK: [[UJAM]] = !{!"llvm.loop.unroll_and_jam.count", i32 17}
+; CHECK: distinct !{[[UNROLLTAG]], [[UNROLL:.*]]}
+; CHECK: [[UNROLL]] = !{!"llvm.loop.unroll.count", i32 1}
+
+
+define dso_local void @loopnest() local_unnamed_addr #0 {
+entry:
+  br label %outer.header
+
+outer.header:                                         ; preds = %inner.header, %entry
+  %ii.0 = phi i64 [ 2, %entry ], [ %add, %inner.header ]
+  %cmp = icmp ult i64 %ii.0, 64
+  br i1 %cmp, label %inner.header, label %outer.header.cleanup
+
+outer.header.cleanup:                                 ; preds = %outer.header
+  ret void
+
+inner.header:                                        ; preds = %outer.header, %inner.body
+  %j.0 = phi i64 [ %add10, %inner.body ], [ %ii.0, %outer.header ]
+  %add = add nuw nsw i64 %ii.0, 16
+  %cmp2 = icmp ult i64 %j.0, %add
+  br i1 %cmp2, label %inner.body, label %outer.header, !llvm.loop !2
+
+inner.body:                                        ; preds = %inner.header
+  %add10 = add nuw nsw i64 %j.0, 1
+  br label %inner.header, !llvm.loop !4
+}
+
+!2 = distinct !{!2, !3}
+!3 = !{!"llvm.loop.unroll_and_jam.count", i32 17}
+!4 = distinct !{!4, !5}
+!5 = !{!"llvm.loop.unroll.count", i32 1}