[llvm] [ConstraintElim] Generalize IV logic to chain of exiting blocks. (PR #108031)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

<details>
<summary>Changes</summary>

Generalize the logic to add facts for inductions in the header block to any chain of exiting blocks starting at the header.

Such IR is for generated by Swift for a number of loops with runtime checks.

The logic could probably generalized further.

---
Full diff: https://github.com/llvm/llvm-project/pull/108031.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/ConstraintElimination.cpp (+36-11) 
- (modified) llvm/test/Transforms/ConstraintElimination/monotonic-pointer-phis-chain-of-exits.ll (+3-9) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 6565aed4bc390c..08fe5b2aa1ff2c 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -193,10 +193,14 @@ struct State {
   /// Process block \p BB and add known facts to work-list.
   void addInfoFor(BasicBlock &BB);
 
-  /// Try to add facts for loop inductions (AddRecs) in EQ/NE compares
-  /// controlling the loop header.
+  /// Try to add facts for loop inductions (AddRecs) from EQ/NE compares
+  /// controlling loop exits.
   void addInfoForInductions(BasicBlock &BB);
 
+  /// Try to add facts for loop inductions (AddRecs) from \p Cond which controls
+  /// an exit from \p L.
+  void addInfoForInduction(ICmpInst *Cond, Loop *L);
+
   /// Returns true if we can add a known condition from BB to its successor
   /// block Succ.
   bool canAddSuccessor(BasicBlock &BB, BasicBlock *Succ) const {
@@ -900,18 +904,14 @@ static void dumpConstraint(ArrayRef<int64_t> C,
 }
 #endif
 
-void State::addInfoForInductions(BasicBlock &BB) {
-  auto *L = LI.getLoopFor(&BB);
-  if (!L || L->getHeader() != &BB)
-    return;
-
+void State::addInfoForInduction(ICmpInst *Cond, Loop *L) {
   Value *A;
   Value *B;
   CmpInst::Predicate Pred;
 
-  if (!match(BB.getTerminator(),
-             m_Br(m_ICmp(Pred, m_Value(A), m_Value(B)), m_Value(), m_Value())))
+  if (!match(Cond, m_ICmp(Pred, m_Value(A), m_Value(B))))
     return;
+
   PHINode *PN = dyn_cast<PHINode>(A);
   if (!PN) {
     Pred = CmpInst::getSwappedPredicate(Pred);
@@ -919,10 +919,11 @@ void State::addInfoForInductions(BasicBlock &BB) {
     PN = dyn_cast<PHINode>(A);
   }
 
-  if (!PN || PN->getParent() != &BB || PN->getNumIncomingValues() != 2 ||
-      !SE.isSCEVable(PN->getType()))
+  if (!PN || PN->getParent() != L->getHeader() ||
+      PN->getNumIncomingValues() != 2 || !SE.isSCEVable(PN->getType()))
     return;
 
+  BasicBlock &BB = *Cond->getParent();
   BasicBlock *InLoopSucc = nullptr;
   if (Pred == CmpInst::ICMP_NE)
     InLoopSucc = cast<BranchInst>(BB.getTerminator())->getSuccessor(0);
@@ -931,6 +932,8 @@ void State::addInfoForInductions(BasicBlock &BB) {
   else
     return;
 
+  assert(InLoopSucc != L->getHeader() &&
+         "Cannot inject condition back to loop header");
   if (!L->contains(InLoopSucc) || !L->isLoopExiting(&BB) || InLoopSucc == &BB)
     return;
 
@@ -1048,6 +1051,28 @@ void State::addInfoForInductions(BasicBlock &BB) {
   }
 }
 
+void State::addInfoForInductions(BasicBlock &BB) {
+  auto *L = LI.getLoopFor(&BB);
+  if (!L)
+    return;
+  if (L->getHeader() != &BB)
+    return;
+
+  BasicBlock *Curr = &BB;
+  while (L->isLoopExiting(Curr)) {
+    // Don't try to add condition from latch to loop header.
+    if (L->isLoopLatch(Curr))
+      break;
+    auto *Term = dyn_cast<BranchInst>(Curr->getTerminator());
+    if (!Term)
+      break;
+    if (isa<ICmpInst>(Term->getCondition()))
+      addInfoForInduction(cast<ICmpInst>(Term->getCondition()), L);
+    Curr = L->contains(Term->getSuccessor(0)) ? Term->getSuccessor(0)
+                                              : Term->getSuccessor(1);
+  }
+}
+
 void State::addInfoFor(BasicBlock &BB) {
   addInfoForInductions(BB);
 
diff --git a/llvm/test/Transforms/ConstraintElimination/monotonic-pointer-phis-chain-of-exits.ll b/llvm/test/Transforms/ConstraintElimination/monotonic-pointer-phis-chain-of-exits.ll
index 6b128d9e525ca3..a2673bd8e73efb 100644
--- a/llvm/test/Transforms/ConstraintElimination/monotonic-pointer-phis-chain-of-exits.ll
+++ b/llvm/test/Transforms/ConstraintElimination/monotonic-pointer-phis-chain-of-exits.ll
@@ -25,9 +25,7 @@ define void @test_monotonic_ptr_iv_inc_1_different_element_types_1_chain_of_2_ex
 ; CHECK-NEXT:    [[C_2:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[C_2]], label [[LOOP_NEXT:%.*]], label [[EXIT]]
 ; CHECK:       loop.next:
-; CHECK-NEXT:    [[T_1:%.*]] = icmp uge ptr [[PTR_IV]], [[START]]
-; CHECK-NEXT:    [[T_2:%.*]] = icmp ult ptr [[PTR_IV]], [[UPPER]]
-; CHECK-NEXT:    [[AND:%.*]] = and i1 [[T_1]], [[T_2]]
+; CHECK-NEXT:    [[AND:%.*]] = and i1 true, true
 ; CHECK-NEXT:    br i1 [[AND]], label [[LOOP_LATCH]], label [[EXIT]]
 ; CHECK:       loop.latch:
 ; CHECK-NEXT:    call void @use(ptr [[PTR_IV]])
@@ -94,9 +92,7 @@ define void @test_monotonic_ptr_iv_inc_1_different_element_types_1_chain_of_3_ex
 ; CHECK-NEXT:    [[C_2:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[C_2]], label [[LOOP_NEXT:%.*]], label [[EXIT]]
 ; CHECK:       loop.next:
-; CHECK-NEXT:    [[T_1:%.*]] = icmp uge ptr [[PTR_IV]], [[START]]
-; CHECK-NEXT:    [[T_2:%.*]] = icmp ult ptr [[PTR_IV]], [[UPPER]]
-; CHECK-NEXT:    [[AND:%.*]] = and i1 [[T_1]], [[T_2]]
+; CHECK-NEXT:    [[AND:%.*]] = and i1 true, true
 ; CHECK-NEXT:    br i1 [[AND]], label [[LOOP_LATCH]], label [[EXIT]]
 ; CHECK:       loop.latch:
 ; CHECK-NEXT:    call void @use(ptr [[PTR_IV]])
@@ -163,9 +159,7 @@ define void @test_monotonic_ptr_iv_inc_1_different_element_types_1_chain_of_2_ex
 ; CHECK-NEXT:    [[C_2:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[C_2]], label [[LOOP_NEXT:%.*]], label [[EXIT]]
 ; CHECK:       loop.next:
-; CHECK-NEXT:    [[T_1:%.*]] = icmp uge ptr [[PTR_IV]], [[START]]
-; CHECK-NEXT:    [[T_2:%.*]] = icmp ult ptr [[PTR_IV]], [[UPPER]]
-; CHECK-NEXT:    [[AND:%.*]] = and i1 [[T_1]], [[T_2]]
+; CHECK-NEXT:    [[AND:%.*]] = and i1 true, true
 ; CHECK-NEXT:    br i1 [[AND]], label [[LOOP_LATCH]], label [[EXIT]]
 ; CHECK:       loop.latch:
 ; CHECK-NEXT:    call void @use(ptr [[PTR_IV]])

``````````

</details>


https://github.com/llvm/llvm-project/pull/108031