[LoopInfo] Return conditional branch of preheader's predecessor if loop's exit has exactly two predecessors which are latch and guard

[vdsered]

Summary

While working on an optimization for DSE I've noticed that for an example like the one below we return getLoopGuardBranch == nullptr and isGuarded == false for both loops even if they are in rotated form. See IR in an example on godbolt.

Technically, entry block decides if both loops execute, but there is a block with br i1 %3, label %7, label %14 where %3 is a icmp in entry block, so it can be a guard for the second loop.

void func(int *P, int N) {
    for (int i = 0; i < N; ++i) {
        P[i] = 1;
    }
    for (int i = 0; i < N; ++i) {
        P[i] = 2;
    }
}

This happens because loops are not in canonical form (specifically, because of no dedicated exits). I suggest to relax a little this requirement to have dedicated exits and add a check for a case when there is an exit that has exactly two predecessors and both of them are latch from loop and the block whose branch instruction we return in getLoopGuardBranch. So, we find a guard in a case as in the example above.

[Whitney]

Why is the input to DSE not in simplified form? By running -loop-simplify after the IR you provided resolves this problem.

llvm/lib/Analysis/LoopInfo.cpp

396

no need braces for single statement block.

llvm/unittests/Analysis/LoopInfoTest.cpp

1587

isn't it better to make sure is FI->getTerminator().

[vdsered]

Loops lose simplify form in JumpThreading in default pipeline before DSE, for example. I don't know the details and if this is correct or not. There is currently no need for a loop to be in simplify form in DSE in order to work with it (except that it should be a natural loop). It became only required for my optimization.

I think now that it is probably be better to transform a loop to simplify form for this optimization in DSE than generalizing getLoopGuardBranch.

llvm/lib/Analysis/LoopInfo.cpp

if (IndDesc.getInductionOpcode() != Instruction::Add &&
       IndDesc.getInductionOpcode() != Instruction::Sub)
     return false;
 
   // Incremented by a loop invariant step for each loop iteration
   return SE.isLoopInvariant(IndDesc.getStep(), this);
 }
 
 BranchInst *Loop::getLoopGuardBranch() const {
-  if (!isLoopSimplifyForm())
-    return nullptr;
-
   BasicBlock *Preheader = getLoopPreheader();
-  assert(Preheader && getLoopLatch() &&
-         "Expecting a loop with valid preheader and latch");
-
+  if (!Preheader)
+    return nullptr;
+  if (!getLoopLatch())
+    return nullptr;
   // Loop should be in rotate form.
   if (!isRotatedForm())
     return nullptr;
-
   // Disallow loops with more than one unique exit block, as we do not verify
   // that GuardOtherSucc post dominates all exit blocks.
   BasicBlock *ExitFromLatch = getUniqueExitBlock();
   if (!ExitFromLatch)
     return nullptr;
 
   BasicBlock *GuardBB = Preheader->getUniquePredecessor();
   if (!GuardBB)
     return nullptr;
-
   assert(GuardBB->getTerminator() && "Expecting valid guard terminator");
 
   BranchInst *GuardBI = dyn_cast<BranchInst>(GuardBB->getTerminator());
   if (!GuardBI || GuardBI->isUnconditional())
     return nullptr;
 
   BasicBlock *GuardOtherSucc = (GuardBI->getSuccessor(0) == Preheader)
                                    ? GuardBI->getSuccessor(1)
                                    : GuardBI->getSuccessor(0);
+  if (!hasDedicatedExits()) {
+    if (ExitFromLatch->hasNPredecessors(2) && ExitFromLatch == GuardOtherSucc) {
+      return GuardBI;
+    }
+    return nullptr;
+  }
 
   // Check if ExitFromLatch (or any BasicBlock which is an empty unique
   // successor of ExitFromLatch) is equal to GuardOtherSucc. If
   // skipEmptyBlockUntil returns GuardOtherSucc, then the guard branch for the
   // loop is GuardBI (return GuardBI), otherwise return nullptr.
   if (&LoopNest::skipEmptyBlockUntil(ExitFromLatch, GuardOtherSucc,
                                      /*CheckUniquePred=*/true) ==
       GuardOtherSucc)

llvm/unittests/Analysis/LoopInfoTest.cpp

runWithLoopInfo(*M, "foo", [&](Function &F, LoopInfo &LI) {
 
     Loop *L = LI.getLoopFor(Header);
     EXPECT_NE(L, nullptr);
 
     // L should not have a guard branch
     EXPECT_EQ(L->getLoopGuardBranch(), nullptr);
   });
 }
+
+TEST(LoopInfoTest, NoDedicatedExitsButTheSameBBWithGuard) {
+  const char *ModuleStr =
+      "target datalayout = \"e-m:o-i64:64-f80:128-n8:16:32:64-S128\"\n"
+      "define void @foo(i32* %B, i64 signext %nx, i1 %cond) {\n"
+      "entry:\n"
+      "  %cmp.guard = icmp slt i64 0, %nx\n"
+      "  br i1 %cmp.guard, label %for.i.preheader, label %for.i.exit\n"
+      "for.i.preheader:\n"
+      "  br label %for.i\n"
+      "for.i:\n"
+      "  %i = phi i64 [ 0, %for.i.preheader ], [ %inc13, %for.i ]\n"
+      "  %Bi = getelementptr inbounds i32, i32* %B, i64 %i\n"
+      "  store i32 0, i32* %Bi, align 4\n"
+      "  %inc13 = add nsw i64 %i, 1\n"
+      "  %cmp = icmp slt i64 %inc13, %nx\n"
+      "  br i1 %cmp, label %for.i, label %for.i.exit\n"
+      "for.i.exit:\n"
+      "  ret void\n"
+      "}\n";
+
+  // Parse the module.
+  LLVMContext Context;
+  std::unique_ptr<Module> M = makeLLVMModule(Context, ModuleStr);
+
+  runWithLoopInfo(*M, "foo", [&](Function &F, LoopInfo &LI) {
+    Function::iterator FI = F.begin();
+    assert(FI->getName() == "entry");
+
+    FI = ++FI;
+    BasicBlock *Header = &*(++FI);
+    assert(Header->getName() == "for.i");
+
+    Loop *L = LI.getLoopFor(Header);
+    EXPECT_NE(L, nullptr);
+
+    EXPECT_EQ(L->isGuarded(), true);
+    EXPECT_NE(L->getLoopGuardBranch(), nullptr);
+  });
+}