[MustExec][LICM] Refactor the first iteration condition check

[jdoerfert]

Summary

This patch exposes components that allow to simplify conditions for the
first iteration of a loop. The logic was almost completely copied.
Slight modifications include:

• Use the loop predecessor instead of loop preheader.
• Try to simplify both sides of the ICmp.
• Do not give up if no PHI was found or the PHI was not in a header.

One test was improved, a proper negative example was added to keep
coverage high.

[jdoerfert]

@reames friendly ping

[reames]

LGTM w/required changes before commit.

llvm/lib/Analysis/MustExecute.cpp

123

All callers pass L. Simplify.

126

The LI but no L path is utterly untested in this patch. Remove.

llvm/lib/Analysis/MustExecute.cpp

void LoopSafetyInfo::computeBlockColors(const Loop *CurLoop) {
   // personality routine.
   Function *Fn = CurLoop->getHeader()->getParent();
   if (Fn->hasPersonalityFn())
     if (Constant *PersonalityFn = Fn->getPersonalityFn())
       if (isScopedEHPersonality(classifyEHPersonality(PersonalityFn)))
         BlockColors = colorEHFunclets(*Fn);
 }
 
+static Optional<uint64_t>
+getConstantIntegerValueInFirstIteration(const Value &V, const Loop *L,
+                                        const DominatorTree *DT,
+                                        const LoopInfo *LI) {
+  if (auto *Cond = dyn_cast<ConstantInt>(&V))
+    return Optional<uint64_t>(Cond->getZExtValue());
+
+  if (auto *Cond = dyn_cast<CmpInst>(&V)) {
+    // TODO: this would be a lot more powerful if we used scev, but all the
+    //       plumbing is currently missing to pass a pointer in from the pass
+    // Check for `cmp (phi [x, predecessor] ...), y` where `pred x, y` is known
+    auto SimplifyPHI = [&](Value *V) -> Value * {
+      auto *PHI = dyn_cast<PHINode>(V);
+      if (!PHI)
+        return V;
+      // TODO: Remove the handling of a special loop in favor of the loop info
+      //       solution once the user is gone.
+      if (L && PHI->getParent() == L->getHeader() && L->getLoopPredecessor())
+        return PHI->getIncomingValueForBlock(L->getLoopPredecessor());
+      const Loop *PL = LI ? LI->getLoopFor(PHI->getParent()) : nullptr;
+      if (PL && PL->getHeader() == PHI->getParent() && PL->getLoopPredecessor())
+        return PHI->getIncomingValueForBlock(PL->getLoopPredecessor());
+      return V;
+    };
+    auto *LHS = SimplifyPHI(Cond->getOperand(0));
+    auto *RHS = SimplifyPHI(Cond->getOperand(1));
+
+    const DataLayout &DL = Cond->getModule()->getDataLayout();
+    auto *SimpleValOrNull =
+        SimplifyCmpInst(Cond->getPredicate(), LHS, RHS,
+                        {DL, /*TLI*/ nullptr, DT, /*AC*/ nullptr, Cond});
+    if (auto *Cst = dyn_cast_or_null<ConstantInt>(SimpleValOrNull))
+      return Optional<uint64_t>(Cst->getZExtValue());
+  }
+
+  return None;
+}
+
+static const BasicBlock *
+getSuccessorInFirstIteration(const Instruction &TI, const Loop *L,
+                             const DominatorTree *DT, const LoopInfo *LI) {
+  assert(TI.isTerminator() && "Expected a terminator");
+  if (auto *BI = dyn_cast<BranchInst>(&TI)) {
+    if (BI->isUnconditional())
+      return BI->getSuccessor(0);
+
+    if (L || LI) {
+      Optional<uint64_t> CV = getConstantIntegerValueInFirstIteration(
+          *BI->getCondition(), L, DT, LI);
+      if (CV.hasValue()) {
+        assert(CV.getValue() < 2 && "Expected boolean value!");
+        return BI->getSuccessor(1 - CV.getValue());
+      }
+    }
+  }
+  return nullptr;
+}
+
 /// Return true if we can prove that the given ExitBlock is not reached on the
 /// first iteration of the given loop.  That is, the backedge of the loop must
 /// be executed before the ExitBlock is executed in any dynamic execution trace.
 static bool CanProveNotTakenFirstIteration(const BasicBlock *ExitBlock,
                                            const DominatorTree *DT,
                                            const Loop *CurLoop) {
   auto *CondExitBlock = ExitBlock->getSinglePredecessor();
   if (!CondExitBlock)
     // expect unique exits
     return false;
   assert(CurLoop->contains(CondExitBlock) && "meaning of exit block");
-  auto *BI = dyn_cast<BranchInst>(CondExitBlock->getTerminator());
-  if (!BI || !BI->isConditional())
-    return false;
-  // If condition is constant and false leads to ExitBlock then we always
-  // execute the true branch.
-  if (auto *Cond = dyn_cast<ConstantInt>(BI->getCondition()))
-    return BI->getSuccessor(Cond->getZExtValue() ? 1 : 0) == ExitBlock;
-  auto *Cond = dyn_cast<CmpInst>(BI->getCondition());
-  if (!Cond)
-    return false;
-  // todo: this would be a lot more powerful if we used scev, but all the
-  // plumbing is currently missing to pass a pointer in from the pass
-  // Check for cmp (phi [x, preheader] ...), y where (pred x, y is known
-  auto *LHS = dyn_cast<PHINode>(Cond->getOperand(0));
-  auto *RHS = Cond->getOperand(1);
-  if (!LHS || LHS->getParent() != CurLoop->getHeader())
-    return false;
-  auto DL = ExitBlock->getModule()->getDataLayout();
-  auto *IVStart = LHS->getIncomingValueForBlock(CurLoop->getLoopPreheader());
-  auto *SimpleValOrNull = SimplifyCmpInst(Cond->getPredicate(),
-                                          IVStart, RHS,
-                                          {DL, /*TLI*/ nullptr,
-                                              DT, /*AC*/ nullptr, BI});
-  auto *SimpleCst = dyn_cast_or_null<Constant>(SimpleValOrNull);
-  if (!SimpleCst)
-    return false;
-  if (ExitBlock == BI->getSuccessor(0))
-    return SimpleCst->isZeroValue();
-  assert(ExitBlock == BI->getSuccessor(1) && "implied by above");
-  return SimpleCst->isAllOnesValue();
+  const BasicBlock * SuccInFirstIteration =
+      getSuccessorInFirstIteration(*CondExitBlock->getTerminator(), CurLoop, DT,
+                                   /* LoopInfo */ nullptr);
+  if (SuccInFirstIteration)
+    return SuccInFirstIteration != ExitBlock;
 }
 
 /// Collect all blocks from \p CurLoop which lie on all possible paths from
 /// the header of \p CurLoop (inclusive) to BB (exclusive) into the set
 /// \p Predecessors. If \p BB is the header, \p Predecessors will be empty.
 static void collectTransitivePredecessors(
     const Loop *CurLoop, const BasicBlock *BB,
     SmallPtrSetImpl<const BasicBlock *> &Predecessors) {

llvm/test/Transforms/LICM/hoist-mustexec.ll

for.body:
   %acc = phi i32 [ 0, %entry ], [ %add, %continue ]
   %cond = icmp ult i32 %iv, 500
   br i1 %cond, label %dummy_block1, label %dummy_block2
 
 dummy_block1:
   br label %dummy_block2
 
 dummy_block2:
-  %wrongphi = phi i32 [11, %for.body], [12, %dummy_block1]
+  %wrongphi = phi i32 [11, %for.body], [2001, %dummy_block1]
   %r.chk = icmp ugt i32 %wrongphi, 2000
   br i1 %r.chk, label %fail, label %continue
 continue:
 ; CHECK-LABEL: continue
 ; CHECK: %i1 = load i32, i32* %a, align 4
   %i1 = load i32, i32* %a, align 4
   %add = add nsw i32 %i1, %acc
   %inc = add nuw nsw i32 %iv, 1
   %exitcond = icmp eq i32 %inc, 1000
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 
 for.cond.cleanup:
   ret i32 %add
 
 fail:
   call void @f()
   ret i32 -1
 }
 
+define i32 @test-wrongphi2(i32* noalias nocapture readonly %a) nounwind uwtable {
+; CHECK-LABEL: @test-wrongphi2(
+; CHECK-LABEL: entry
+; CHECK: %i1 = load i32, i32* %a, align 4
+entry:
+  br label %for.body
+
+for.body:
+  %iv = phi i32 [ 0, %entry ], [ %inc, %continue ]
+  %acc = phi i32 [ 0, %entry ], [ %add, %continue ]
+  %cond = icmp ult i32 %iv, 500
+  br i1 %cond, label %dummy_block1, label %dummy_block2
+
+dummy_block1:
+  br label %dummy_block2
+
+dummy_block2:
+  %wrongphi = phi i32 [11, %for.body], [12, %dummy_block1]
+  %r.chk = icmp ugt i32 %wrongphi, 2000
+  br i1 %r.chk, label %fail, label %continue
+continue:
+  %i1 = load i32, i32* %a, align 4
+  %add = add nsw i32 %i1, %acc
+  %inc = add nuw nsw i32 %iv, 1
+  %exitcond = icmp eq i32 %inc, 1000
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:
+  ret i32 %add
+
+fail:
+  call void @f()
+  ret i32 -1
+}
+
 ; This works because loop-simplify is run implicitly, but test for it anyways
 define i32 @test-multiple-latch(i32* noalias nocapture readonly %a) nounwind uwtable {
 ; CHECK-LABEL: @test-multiple-latch(
 entry:
 ; CHECK: %i1 = load i32, i32* %a, align 4
 ; CHECK-NEXT: br label %for.body
   br label %for.body