[llvm] [SCEVExpander] Do not insert after callbr (PR #201322)

[Le Gia Bao via llvm-commits]

https://github.com/L3GiaBao updated https://github.com/llvm/llvm-project/pull/201322

>From 4ae5444c202819cb37211e94f4969c2aee137828 Mon Sep 17 00:00:00 2001
From: Le Gia Bao <legiabao7300 at gmail.com>
Date: Wed, 3 Jun 2026 07:08:18 -0400
Subject: [PATCH 1/2] [SCEVExpander] Do not insert after callbr

---
 .../Transforms/Utils/ScalarEvolutionExpander.cpp    | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
index 39da38e491817..d21fb4339b664 100644
--- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
@@ -147,9 +147,18 @@ Value *SCEVExpander::ReuseOrCreateCast(Value *V, Type *Ty,
 BasicBlock::iterator
 SCEVExpander::findInsertPointAfter(Instruction *I,
                                    Instruction *MustDominate) const {
-  BasicBlock::iterator IP = ++I->getIterator();
-  if (auto *II = dyn_cast<InvokeInst>(I))
+  BasicBlock::iterator IP;
+  if (auto *II = dyn_cast<InvokeInst>(I)) {
     IP = II->getNormalDest()->begin();
+  } else if (isa<CallBrInst>(I)) {
+    // callbr is a terminator, so there is no valid same-block point after it.
+    // Use the point we need the result to dominate.
+    assert(SE.DT.dominates(I, MustDominate) &&
+           "callbr result must dominate the insertion point");
+    IP = MustDominate->getIterator();
+  } else {
+    IP = ++I->getIterator();
+  }
 
   while (isa<PHINode>(IP))
     ++IP;

>From bcdcf9a8a94acab1645a07335b5f85e286f1e79e Mon Sep 17 00:00:00 2001
From: Le Gia Bao <legiabao7300 at gmail.com>
Date: Sat, 6 Jun 2026 17:24:18 -0400
Subject: [PATCH 2/2] Add test file and cleanup

---
 .../Utils/ScalarEvolutionExpander.cpp         | 32 +++++------
 llvm/test/Transforms/IndVarSimplify/callbr.ll | 53 +++++++++++++++++++
 2 files changed, 65 insertions(+), 20 deletions(-)
 create mode 100644 llvm/test/Transforms/IndVarSimplify/callbr.ll

diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
index d21fb4339b664..97d0bedddfce8 100644
--- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
@@ -148,27 +148,19 @@ BasicBlock::iterator
 SCEVExpander::findInsertPointAfter(Instruction *I,
                                    Instruction *MustDominate) const {
   BasicBlock::iterator IP;
-  if (auto *II = dyn_cast<InvokeInst>(I)) {
-    IP = II->getNormalDest()->begin();
-  } else if (isa<CallBrInst>(I)) {
-    // callbr is a terminator, so there is no valid same-block point after it.
-    // Use the point we need the result to dominate.
-    assert(SE.DT.dominates(I, MustDominate) &&
-           "callbr result must dominate the insertion point");
-    IP = MustDominate->getIterator();
-  } else {
-    IP = ++I->getIterator();
-  }
-
-  while (isa<PHINode>(IP))
-    ++IP;
-
-  if (isa<FuncletPadInst>(IP) || isa<LandingPadInst>(IP)) {
-    ++IP;
-  } else if (isa<CatchSwitchInst>(IP)) {
-    IP = MustDominate->getParent()->getFirstInsertionPt();
+  if (auto MaybeIP = I->getInsertionPointAfterDef()) {
+    IP = *MaybeIP;
   } else {
-    assert(!IP->isEHPad() && "unexpected eh pad!");
+    assert(SE.DT.dominates(I, MustDominate) &&
+           "instruction must dominate the insertion point");
+    if (isa<CallBrInst>(I)) {
+      // Def is available in multiple successors, so use the point that needs
+      // to be dominated.
+      IP = MustDominate->getIterator();
+    } else {
+      // catchswitch blocks don't have any legal insertion point.
+      IP = MustDominate->getParent()->getFirstInsertionPt();
+    }
   }
 
   // Adjust insert point to be after instructions inserted by the expander, so
diff --git a/llvm/test/Transforms/IndVarSimplify/callbr.ll b/llvm/test/Transforms/IndVarSimplify/callbr.ll
new file mode 100644
index 0000000000000..eb545fa3be236
--- /dev/null
+++ b/llvm/test/Transforms/IndVarSimplify/callbr.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=indvars < %s | FileCheck %s
+
+target datalayout = "e-p:64:64"
+
+define ptr @callbr_result_exit_value(ptr %a0) {
+; CHECK-LABEL: define ptr @callbr_result_exit_value(
+; CHECK-SAME: ptr [[A0:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[A02:%.*]] = ptrtoint ptr [[A0]] to i64
+; CHECK-NEXT:    [[OUT:%.*]] = callbr ptr asm "", "=r,!i"()
+; CHECK-NEXT:            to label %[[D:.*]] [label %indirect]
+; CHECK:       [[INDIRECT:.*:]]
+; CHECK-NEXT:    ret ptr null
+; CHECK:       [[D]]:
+; CHECK-NEXT:    [[CMP0:%.*]] = icmp eq ptr [[A0]], [[OUT]]
+; CHECK-NEXT:    br i1 [[CMP0]], label %[[END:.*]], label %[[LOOP_PREHEADER:.*]]
+; CHECK:       [[LOOP_PREHEADER]]:
+; CHECK-NEXT:    [[OUT1:%.*]] = ptrtoint ptr [[OUT]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = sub i64 [[OUT1]], [[A02]]
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    br i1 true, label %[[EXIT:.*]], label %[[LOOP]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[SCEVGEP:%.*]] = getelementptr i8, ptr [[A0]], i64 [[TMP4]]
+; CHECK-NEXT:    ret ptr [[SCEVGEP]]
+; CHECK:       [[END]]:
+; CHECK-NEXT:    ret ptr [[A0]]
+;
+entry:
+  %out = callbr ptr asm "", "=r,!i"()
+  to label %d [label %indirect]
+
+indirect:
+  ret ptr null
+
+d:
+  %cmp0 = icmp eq ptr %a0, %out
+  br i1 %cmp0, label %end, label %loop
+
+loop:
+  %iv = phi ptr [ %a0, %d ], [ %inc, %loop ]
+  %inc = getelementptr inbounds i8, ptr %iv, i64 1
+  %cmp = icmp eq ptr %inc, %out
+  br i1 %cmp, label %exit, label %loop
+
+exit:
+  %lcssa = phi ptr [ %inc, %loop ]
+  ret ptr %lcssa
+
+end:
+  ret ptr %a0
+}