[llvm] [ConstraintElim] Use cond from header as upper bound on IV in exit BB. (PR #94610)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/94610

For loops, we can use the condition in the loop header as upper bound on the compared induction in the unique exit block, if it exists. This can be done even if there are multiple in-loop edges to the unique exit block, as any other exit may only exit earlier.

More generally, we could add the OR of all exit conditions to the exit, but that's a possible future extension.

Fixes https://github.com/llvm/llvm-project/issues/90417.

>From dc6c9f5ae3906c0909cf3c6cb9c724e867781187 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Thu, 6 Jun 2024 13:47:20 +0100
Subject: [PATCH] [ConstraintElim] Use cond from header as upper bound on IV in
 exit BB.

For loops, we can use the condition in the loop header as upper bound on
the compared induction in the unique exit block, if it exists. This can
be done even if there are multiple in-loop edges to the unique exit
block, as any other exit may only exit earlier.

More generally, we could add the OR of all exit conditions to the exit,
but that's a possible future extension.

Fixes https://github.com/llvm/llvm-project/issues/90417.
---
 .../Transforms/Scalar/ConstraintElimination.cpp | 17 +++++++++++++++++
 .../induction-condition-in-loop-exit.ll         |  9 +++------
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 70bfa469193bf..4b3ef4d4c222c 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -1031,6 +1031,23 @@ void State::addInfoForInductions(BasicBlock &BB) {
   WorkList.push_back(FactOrCheck::getConditionFact(
       DTN, CmpInst::ICMP_SLT, PN, B,
       ConditionTy(CmpInst::ICMP_SLE, StartValue, B)));
+
+  assert(!StepOffset.isNegative() && "induction must be increasing");
+  // Try to add condition from header to the unique exit block, if there is one.
+  // When exiting either with EQ or NE, we know that the induction value must be
+  // u<= B, as a different exit may exit earlier.
+  if (Pred == CmpInst::ICMP_EQ) {
+    BasicBlock *EB = cast<BranchInst>(BB.getTerminator())->getSuccessor(0);
+    if (L->getUniqueExitBlock() == EB)
+      WorkList.emplace_back(FactOrCheck::getConditionFact(
+          DT.getNode(EB), CmpInst::ICMP_ULE, A, B));
+  }
+  if (Pred == CmpInst::ICMP_NE) {
+    BasicBlock *EB = cast<BranchInst>(BB.getTerminator())->getSuccessor(1);
+    if (L->getUniqueExitBlock() == EB)
+      WorkList.emplace_back(FactOrCheck::getConditionFact(
+          DT.getNode(EB), CmpInst::ICMP_ULE, A, B));
+  }
 }
 
 void State::addInfoFor(BasicBlock &BB) {
diff --git a/llvm/test/Transforms/ConstraintElimination/induction-condition-in-loop-exit.ll b/llvm/test/Transforms/ConstraintElimination/induction-condition-in-loop-exit.ll
index 44ce82b51d707..86828b5e7f369 100644
--- a/llvm/test/Transforms/ConstraintElimination/induction-condition-in-loop-exit.ll
+++ b/llvm/test/Transforms/ConstraintElimination/induction-condition-in-loop-exit.ll
@@ -17,8 +17,7 @@ define i1 @multi_exiting_loop_eq_same_unique_exit_const_compare_known(ptr %s) {
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 1
 ; CHECK-NEXT:    br i1 [[LATCH_C]], label %[[LOOP_HEADER]], label %[[EXIT]]
 ; CHECK:       [[EXIT]]:
-; CHECK-NEXT:    [[T:%.*]] = icmp ult i32 [[IV]], 1235
-; CHECK-NEXT:    ret i1 [[T]]
+; CHECK-NEXT:    ret i1 true
 ;
 entry:
   br label %loop.header
@@ -175,8 +174,7 @@ define i1 @multi_exiting_loop_eq_same_unique_exit_var_compare_known(ptr %s, i32
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 1
 ; CHECK-NEXT:    br i1 [[LATCH_C]], label %[[LOOP_HEADER]], label %[[EXIT]]
 ; CHECK:       [[EXIT]]:
-; CHECK-NEXT:    [[T:%.*]] = icmp ule i32 [[IV]], [[N]]
-; CHECK-NEXT:    ret i1 [[T]]
+; CHECK-NEXT:    ret i1 true
 ;
 entry:
   br label %loop.header
@@ -214,8 +212,7 @@ define i1 @multi_exiting_loop_ne_same_unique_exit_const_compare_known(ptr %s) {
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 1
 ; CHECK-NEXT:    br i1 [[LATCH_C]], label %[[LOOP_HEADER]], label %[[EXIT]]
 ; CHECK:       [[EXIT]]:
-; CHECK-NEXT:    [[T:%.*]] = icmp ult i32 [[IV]], 1235
-; CHECK-NEXT:    ret i1 [[T]]
+; CHECK-NEXT:    ret i1 true
 ;
 entry:
   br label %loop.header