[LICM] update BlockColors after splitting predecessors

[junbuml]

Summary

Update BlockColors after splitting predecessors. Do not allow splitting
LandingPad for sinking when the BlockColors is not empty, so we can
simply assign predecessor's color to the new block.

Fixes PR36184

[junbuml]

Added a test case.

[majnemer]

lib/Transforms/Scalar/LICM.cpp

860–867

You don't need to worry about updating BlockColors if we have landingpads. If there is a landingpad, BlockColors should be entirely empty.

[junbuml]

lib/Transforms/Scalar/LICM.cpp

860–867

As far as I see, BlockColors is not empty when we have landingpads.
For example, in the IR below, I can see the color for %lpadBB is %lpadBB and the color of its successor (%lpadBBSucc1) is %lpadBB, which is just what I expected after colorEHFunclets(). If we split %lpadBB for sinking, SplitLandingPadPredecessors() will be called and the new blocks split (%lpadBB.split) become landpad blocks. Then, color for new blocks and its all successors should be changed accordingly based the rule in colorEHFunclets(). Please correct me if I miss something.

 define void @test(i1 %b, i32 %v1, i32 %v2) personality i32 (...)* @__CxxFrameHandler3 {
entry:
  br label %while.cond
while.cond:
  br i1 %b, label %try.cont, label %while.body

while.body:
  invoke void @may_throw()
          to label %while.body2 unwind label %lpadBB

while.body2:
  %v = call i32 @getv()
  %sinkable= mul i32 %v, %v2
  invoke void @may_throw2()
          to label %while.cond unwind label %lpadBB
lpadBB:
  %.lcssa1 = phi i32 [ 0, %while.body ], [ %sinkable, %while.body2 ]

  landingpad { i8*, i32 }
               catch i8* null
  br label %lpadBBSucc1

lpadBBSucc1:
  ret void

try.cont:
  ret void
}

[majnemer]

lib/Transforms/Scalar/LICM.cpp

860–867

Using _CxxFrameHandler3 with LandingPadInst is unsupported to the best of my knowledge.

Regardless, I think the best version of the predicate is:

if (!BlockColors.empty() && BB->getFirstNonPHI()->isEHPad())

This way we will still DTRT when using less interesting personality routines.

[junbuml]

Incorporated David comment.

lib/Transforms/Scalar/LICM.cpp

860–867

It make sense to me. I added a test case for this.
Thanks David.

[junbuml]

Minor changes in comments. Please let me know any further comment.

[junbuml]

Kindly ping.
Sorry for the rush, but this is one of release-6.0 blockers.

[junbuml]

Kindly ping.

[junbuml]

Since this is one of release-6.0 blockers, I want to kindly ping one more time. After fixing @majnemer's comment, I believe this is pretty good shape now.
Just in case @majnemer is unavailable for reviewing this soon, added @hfinkel and @mcrosier as reviewer.

[majnemer]

LGTM

lib/Transforms/Scalar/LICM.cpp

862–865

I'd remove the FIXME and replace LandingPad with EHPad

[junbuml]

Thanks David for the review.
Committed in r324916.

lib/Transforms/Scalar/LICM.cpp

 static bool inSubLoop(BasicBlock *BB, Loop *CurLoop, LoopInfo *LI);
 static bool isNotUsedOrFreeInLoop(const Instruction &I, const Loop *CurLoop,
                                   const LoopSafetyInfo *SafetyInfo,
                                   TargetTransformInfo *TTI, bool &FreeInLoop);
 static bool hoist(Instruction &I, const DominatorTree *DT, const Loop *CurLoop,
                   const LoopSafetyInfo *SafetyInfo,
                   OptimizationRemarkEmitter *ORE);
 static bool sink(Instruction &I, LoopInfo *LI, DominatorTree *DT,
-                 const Loop *CurLoop, const LoopSafetyInfo *SafetyInfo,
+                 const Loop *CurLoop, LoopSafetyInfo *SafetyInfo,
                  OptimizationRemarkEmitter *ORE, bool FreeInLoop);
 static bool isSafeToExecuteUnconditionally(Instruction &Inst,
                                            const DominatorTree *DT,
                                            const Loop *CurLoop,
                                            const LoopSafetyInfo *SafetyInfo,
                                            OptimizationRemarkEmitter *ORE,
                                            const Instruction *CtxI = nullptr);
 static bool pointerInvalidatedByLoop(Value *V, uint64_t Size,
static Instruction *sinkThroughTriviallyReplacablePHI(
   if (It != SunkCopies.end())
     New = It->second;
   else
     New = SunkCopies[ExitBlock] =
         CloneInstructionInExitBlock(*I, *ExitBlock, *TPN, LI, SafetyInfo);
   return New;
 }
 
-static bool canSplitPredecessors(PHINode *PN) {
+static bool canSplitPredecessors(PHINode *PN, LoopSafetyInfo *SafetyInfo) {
   BasicBlock *BB = PN->getParent();
   if (!BB->canSplitPredecessors())
     return false;
+  // FIXME: it's not impossible to split LandingPad blocks, but if BlockColors
+  // already exist it require updating BlockColors for all offspring blocks
+  // accordingly. By skipping such corner case, we can make updating BlockColors
+  // after splitting predecessor fairly simple.
+  if (!SafetyInfo->BlockColors.empty() && BB->getFirstNonPHI()->isEHPad())
+    return false;
   for (pred_iterator PI = pred_begin(BB), E = pred_end(BB); PI != E; ++PI) {
     BasicBlock *BBPred = *PI;
     if (isa<IndirectBrInst>(BBPred->getTerminator()))
       return false;
   }
   return true;
 }
 
 static void splitPredecessorsOfLoopExit(PHINode *PN, DominatorTree *DT,
-                                        LoopInfo *LI, const Loop *CurLoop) {
+                                        LoopInfo *LI, const Loop *CurLoop,
+                                        LoopSafetyInfo *SafetyInfo) {
 #ifndef NDEBUG
   SmallVector<BasicBlock *, 32> ExitBlocks;
   CurLoop->getUniqueExitBlocks(ExitBlocks);
   SmallPtrSet<BasicBlock *, 32> ExitBlockSet(ExitBlocks.begin(),
                                              ExitBlocks.end());
 #endif
   BasicBlock *ExitBB = PN->getParent();
   assert(ExitBlockSet.count(ExitBB) && "Expect the PHI is in an exit block.");
#endif
   //   %p1 = phi [%v1, %LB1]  <-- trivially replacable
   //   br %LE
   // LE.split2:
   //   %p2 = phi [%v2, %LB2]  <-- trivially replacable
   //   br %LE
   // LE:
   //   %p = phi [%p1, %LE.split], [%p2, %LE.split2]
   //
+
+  auto &BlockColors = SafetyInfo->BlockColors;
   SmallSetVector<BasicBlock *, 8> PredBBs(pred_begin(ExitBB), pred_end(ExitBB));
+
   while (!PredBBs.empty()) {
     BasicBlock *PredBB = *PredBBs.begin();
     assert(CurLoop->contains(PredBB) &&
            "Expect all predecessors are in the loop");
-    if (PN->getBasicBlockIndex(PredBB) >= 0)
-      SplitBlockPredecessors(ExitBB, PredBB, ".split.loop.exit", DT, LI, true);
+    if (PN->getBasicBlockIndex(PredBB) >= 0) {
+      BasicBlock *NewPred = SplitBlockPredecessors(
+          ExitBB, PredBB, ".split.loop.exit", DT, LI, true);
+
+      // Sine we don't allow splitting EH-blocks, we can simply assign
+      // predecessor's color to the new block.
+      if (!BlockColors.empty())
+        BlockColors[NewPred] = BlockColors[PredBB];
+    }
+
     PredBBs.remove(PredBB);
   }
 }
 
 /// When an instruction is found to only be used outside of the loop, this
 /// function moves it to the exit blocks and patches up SSA form as needed.
 /// This method is guaranteed to remove the original instruction from its
 /// position, and may either delete it or move it to outside of the loop.
 ///
 static bool sink(Instruction &I, LoopInfo *LI, DominatorTree *DT,
-                 const Loop *CurLoop, const LoopSafetyInfo *SafetyInfo,
+                 const Loop *CurLoop, LoopSafetyInfo *SafetyInfo,
                  OptimizationRemarkEmitter *ORE, bool FreeInLoop) {
   DEBUG(dbgs() << "LICM sinking instruction: " << I << "\n");
   ORE->emit([&]() {
     return OptimizationRemark(DEBUG_TYPE, "InstSunk", &I)
            << "sinking " << ore::NV("Inst", &I);
   });
   bool Changed = false;
   if (isa<LoadInst>(I))
if (!DT->isReachableFromEntry(BB)) {
       Changed = true;
       continue;
     }
 
     VisitedUsers.insert(PN);
     if (isTriviallyReplacablePHI(*PN, I))
       continue;
 
-    if (!canSplitPredecessors(PN))
+    if (!canSplitPredecessors(PN, SafetyInfo))
       return Changed;
 
     // Split predecessors of the PHI so that we can make users trivially
     // replacable.
-    splitPredecessorsOfLoopExit(PN, DT, LI, CurLoop);
+    splitPredecessorsOfLoopExit(PN, DT, LI, CurLoop, SafetyInfo);
 
     // Should rebuild the iterators, as they may be invalidated by
     // splitPredecessorsOfLoopExit().
     UI = I.user_begin();
     UE = I.user_end();
   }
 
   if (VisitedUsers.empty())

test/Transforms/LICM/sinking.ll

catch.dispatch:
   %.lcssa1 = phi i32 [ %sinkable, %while.body ], [ %sinkable2, %while.body2 ]
   %cp = cleanuppad within none []
   store i32 %.lcssa1, i32* %s
   cleanupret from %cp unwind to caller
 try.cont:
   ret void
 }
 
+; The sinkable call should be sunk into an exit block split. After splitting
+; the exit block, BlockColor for new blocks should be added properly so
+; that we should be able to access valid ColorVector.
+;
+; CHECK-LABEL:@test21_pr36184
+; CHECK-LABEL: Loop
+; CHECK-NOT: %sinkableCall
+; CHECK-LABEL:Out.split.loop.exit
+; CHECK: %sinkableCall
+define i32 @test21_pr36184(i8* %P) personality i32 (...)* @__CxxFrameHandler3 {
+entry:
+  br label %loop.ph
+
+loop.ph:
+  br label %Loop
+
+Loop:
+  %sinkableCall = call i32 @strlen( i8* %P ) readonly
+  br i1 undef, label %ContLoop, label %Out
+
+ContLoop:
+  br i1 undef, label %Loop, label %Out
+
+Out:
+  %idx = phi i32 [ %sinkableCall, %Loop ], [0, %ContLoop ]
+  ret i32 %idx
+}
+
+; We do not support splitting a landingpad block if BlockColors is not empty.
+; CHECK-LABEL: @test22
+; CHECK-LABEL: while.body2
+; CHECK-LABEL: %mul
+; CHECK-NOT: lpadBB.split{{.*}}
+define void @test22(i1 %b, i32 %v1, i32 %v2) personality i32 (...)* @__CxxFrameHandler3 {
+entry:
+  br label %while.cond
+while.cond:
+  br i1 %b, label %try.cont, label %while.body
+
+while.body:
+  invoke void @may_throw()
+          to label %while.body2 unwind label %lpadBB
+
+while.body2:
+  %v = call i32 @getv()
+  %mul = mul i32 %v, %v2
+  invoke void @may_throw2()
+          to label %while.cond unwind label %lpadBB
+lpadBB:
+  %.lcssa1 = phi i32 [ 0, %while.body ], [ %mul, %while.body2 ]
+  landingpad { i8*, i32 }
+               catch i8* null
+  br label %lpadBBSucc1
+
+lpadBBSucc1:
+  ret void
+
+try.cont:
+  ret void
+}
+
 declare void @may_throw()
 declare void @may_throw2()
 declare i32 @__CxxFrameHandler3(...)
 declare i32 @getv()
 declare i1 @getc()
 declare void @f(i32*)
 declare void @g()