[Loop Simplify] Resolve an issue where metadata is not applied to a loop latch.

Summary:
This patch resolves an issue where the metadata of a loop is not added to the
new loop latch, and not removed from the old loop latch. This issue occurs in
the SplitBlockPredecessors function, which  adds a new block in a loop, and
in the case that the block passed into this function is the header of the loop,
the loop can be modified such that the latch of the loop is replaced.
This patch applies to the Loop Simplify pass since it ensures that each loop
has exit blocks which only have predecessors that are inside of the loop. In
the case that this is not true, the pass will create a new exit block for the
loop. This guarantees that the loop preheader/header will dominate the exit blocks.

Author: sidbav (Sidharth Baveja)

Reviewers: asbirlea (Alina Sbirlea), chandlerc (Chandler Carruth), Whitney (Whitney Tsang), bmahjour (Bardia Mahjour)

Reviewed By:  asbirlea (Alina Sbirlea)

Subscribers: hiraditya (Aditya Kumar), llvm-commits

Tag: LLVM

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

diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index 085d910..86b2eb0 100644 (file)
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -758,12 +758,22 @@ BasicBlock *llvm::SplitBlockPredecessors(BasicBlock *BB,
 
   // The new block unconditionally branches to the old block.
   BranchInst *BI = BranchInst::Create(BB, NewBB);
+
+  Loop *L = nullptr;
+  BasicBlock *OldLatch = nullptr;
   // Splitting the predecessors of a loop header creates a preheader block.
-  if (LI && LI->isLoopHeader(BB))
+  if (LI && LI->isLoopHeader(BB)) {
+    L = LI->getLoopFor(BB);
     // Using the loop start line number prevents debuggers stepping into the
     // loop body for this instruction.
-    BI->setDebugLoc(LI->getLoopFor(BB)->getStartLoc());
-  else
+    BI->setDebugLoc(L->getStartLoc());
+
+    // If BB is the header of the Loop, it is possible that the loop is
+    // modified, such that the current latch does not remain the latch of the
+    // loop. If that is the case, the loop metadata from the current latch needs
+    // to be applied to the new latch.
+    OldLatch = L->getLoopLatch();
+  } else
     BI->setDebugLoc(BB->getFirstNonPHIOrDbg()->getDebugLoc());
 
   // Move the edges from Preds to point to NewBB instead of BB.
@@ -798,6 +808,15 @@ BasicBlock *llvm::SplitBlockPredecessors(BasicBlock *BB,
     UpdatePHINodes(BB, NewBB, Preds, BI, HasLoopExit);
   }
 
+  if (OldLatch) {
+    BasicBlock *NewLatch = L->getLoopLatch();
+    if (NewLatch != OldLatch) {
+      MDNode *MD = OldLatch->getTerminator()->getMetadata("llvm.loop");
+      NewLatch->getTerminator()->setMetadata("llvm.loop", MD);
+      OldLatch->getTerminator()->setMetadata("llvm.loop", nullptr);
+    }
+  }
+
   return NewBB;
 }
 

diff --git a/llvm/test/Transforms/LoopSimplify/update_latch_md.ll b/llvm/test/Transforms/LoopSimplify/update_latch_md.ll
new file mode 100644 (file)
index 0000000..17d4b68
--- /dev/null
+++ b/llvm/test/Transforms/LoopSimplify/update_latch_md.ll
@@ -0,0 +1,47 @@
+; RUN: opt -loop-simplify -S < %s | FileCheck %s
+
+; This will test whether or not the metadata from the current loop 1 latch
+; is removed, and applied to the new latch after running the loop-simplify
+; pass on this function. The loop simplify pass ensures that each loop has exit
+; blocks which only have predecessors that are inside of the loop. This
+; guarantees that the loop preheader/header will dominate the exit blocks. For
+; this function currently loop 2 does not have a dedicated exit block.
+
+; CHECK: loop_1_loopHeader.loopexit:
+; CHECK: br label %loop_1_loopHeader, !llvm.loop [[LOOP_1_LATCH_MD:![0-9]+]]
+; CHECK: loop_2_loopHeader
+; CHECK: br i1 %grt_B, label %loop_1_loopHeader.loopexit, label %loop_2_do
+; CHECK-NOT:  br i1 %grt_B, label %loop_1_loopHeader, label %loop_2_do, !llvm.loop{{.*}}
+
+define void @function(i32 %A) {
+entry:
+  %B = add i32 %A, 45
+  %C = add i32 %A, 22
+  br label %loop_1_loopHeader
+
+loop_1_loopHeader:                              ; preds = %loop_2_loopHeader, %entry
+  %loop_1_idx = phi i32 [ 1, %entry], [ %loop_1_update_idx, %loop_2_loopHeader ]
+  %grt_C = icmp slt i32 %loop_1_idx, %C
+  br i1 %grt_C, label %exit, label %loop_1_do
+
+loop_1_do:                                      ; preds = %loop_1_loopHeader
+  %loop_1_update_idx = add nuw nsw i32 %loop_1_idx, 1
+  br label %loop_2_loopHeader
+
+loop_2_loopHeader:                              ; preds = %loop_2_do, %_loop_1_do
+  %loop_2_idx = phi i32 [ 1, %loop_1_do ], [ %loop_2_update_idx, %loop_2_do ]
+  %grt_B = icmp slt i32 %loop_2_idx, %B
+  br i1 %grt_B, label %loop_1_loopHeader, label %loop_2_do, !llvm.loop !0
+
+loop_2_do:                                      ; preds = %loop_2_loopHeader
+  %loop_2_update_idx = add nuw nsw i32 %loop_2_idx, 1
+  br label %loop_2_loopHeader, !llvm.loop !2
+
+exit:                                       ; preds = %loop_1_loopHeader
+  ret void
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.unroll.disable"}
+!2 = distinct !{!2, !1}
+