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

[sidbav]

Summary

The function SplitBlockPredecessors 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. In this situation, the metadata from the old latch needs to added to the new latch, and removed from the old one.

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

[bmahjour]

llvm/test/Transforms/LoopSimplify/update_latch_md.ll

16

The function and label names should be renamed to make the test more readable.

[sidbav]

Address Bardia's comments

[asbirlea]

Thank you.

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

811

nit: it's enough to check if (OldLatch).

[sidbav]

Address review comments

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

BasicBlock *llvm::SplitBlockPredecessors(BasicBlock *BB,
   }
 
   // Create new basic block, insert right before the original block.
   BasicBlock *NewBB = BasicBlock::Create(
       BB->getContext(), BB->getName() + Suffix, BB->getParent(), 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.
   for (unsigned i = 0, e = Preds.size(); i != e; ++i) {
     // This is slightly more strict than necessary; the minimum requirement
     // is that there be no more than one indirectbr branching to BB. And
     // all BlockAddress uses would need to be updated.
     assert(!isa<IndirectBrInst>(Preds[i]->getTerminator()) &&
BasicBlock *llvm::SplitBlockPredecessors(BasicBlock *BB,
   UpdateAnalysisInformation(BB, NewBB, Preds, DT, LI, MSSAU, PreserveLCSSA,
                             HasLoopExit);
 
   if (!Preds.empty()) {
     // Update the PHI nodes in BB with the values coming from NewBB.
     UpdatePHINodes(BB, NewBB, Preds, BI, HasLoopExit);
   }
 
+  if (L && 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;
 }
 
 void llvm::SplitLandingPadPredecessors(BasicBlock *OrigBB,
                                        ArrayRef<BasicBlock *> Preds,
                                        const char *Suffix1, const char *Suffix2,
                                        SmallVectorImpl<BasicBlock *> &NewBBs,
                                        DominatorTree *DT, LoopInfo *LI,

llvm/test/Transforms/LoopSimplify/update_latch_md.ll

+; 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 @mat_times_vec(i32 %A) {
+mat_times_vec_entry:
+  %B = add i32 %A, 45
+  %C = add i32 %A, 22
+  br label %_loop_1_loopHeader_
+
+_loop_1_loopHeader_:                              ; preds = %_loop_2_loopHeader_, %mat_times_vec_entry
+  %k.0 = phi i32 [ 1, %mat_times_vec_entry ], [ %_loop_1_update_loop_idx, %_loop_2_loopHeader_ ]
+  %_grt_C = icmp slt i32 %k.0, %C
+  br i1 %_grt_C, label %_return_bb, label %_loop_1_do_
+
+_loop_1_do_:                                      ; preds = %_loop_1_loopHeader_
+  %_loop_1_update_loop_idx = add nuw nsw i32 %k.0, 1
+  br label %_loop_2_loopHeader_
+
+_loop_2_loopHeader_:                              ; preds = %_loop_2_do, %_loop_1_do_
+  %j.0 = phi i32 [ 1, %_loop_1_do_ ], [ %_loop_2_update_loop_ix, %_loop_2_do_ ]
+  %_grt_B = icmp slt i32 %j.0, %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_loop_ix = add nuw nsw i32 %j.0, 1
+  br label %_loop_2_loopHeader_, !llvm.loop !2
+
+_return_bb:                                       ; preds = %_loop_1_loopHeader_
+  ret void
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.unroll.disable"}
+!2 = distinct !{!2, !1}
+