[llvm] [CodeGenPrepare] Handle address sinking obtained from invoke (PR #143566)

[Evgenii Kudriashov via llvm-commits]

https://github.com/e-kud created https://github.com/llvm/llvm-project/pull/143566

None

>From b7f4b17c69c0fd021aea986c3b802d3645c9a1ba Mon Sep 17 00:00:00 2001
From: Evgenii Kudriashov <evgenii.kudriashov at intel.com>
Date: Tue, 10 Jun 2025 09:00:35 -0700
Subject: [PATCH] [CodeGenPrepare] Handle address sinking obtained from invoke

---
 llvm/lib/CodeGen/CodeGenPrepare.cpp           |  3 +-
 .../CodeGenPrepare/X86/sink-addr-reuse.ll     | 35 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 32348a899683d..7529ead8af202 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -5790,7 +5790,8 @@ static BasicBlock::iterator findInsertPos(Value *Addr, Instruction *MemoryInst,
   // instruction after it.
   if (SunkAddr) {
     if (Instruction *AddrInst = dyn_cast<Instruction>(SunkAddr))
-      return std::next(AddrInst->getIterator());
+      return AddrInst->isTerminator() ? MemoryInst->getIterator()
+                                      : std::next(AddrInst->getIterator());
   }
 
   // Find the first user of Addr in current BB.
diff --git a/llvm/test/Transforms/CodeGenPrepare/X86/sink-addr-reuse.ll b/llvm/test/Transforms/CodeGenPrepare/X86/sink-addr-reuse.ll
index 019f311406550..5f084fc7ea2b5 100644
--- a/llvm/test/Transforms/CodeGenPrepare/X86/sink-addr-reuse.ll
+++ b/llvm/test/Transforms/CodeGenPrepare/X86/sink-addr-reuse.ll
@@ -42,3 +42,38 @@ bb3:
 bb7:
   ret void
 }
+
+; Test a scenario when the address is the last instruction in the basic block.
+
+define void @adress_terminator() personality ptr null {
+; CHECK-LABEL: define void @adress_terminator() personality ptr null {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[PTR:%.*]] = invoke ptr null(i64 0)
+; CHECK-NEXT:            to label %[[BODY_1:.*]] unwind label %[[EHCLEANUP:.*]]
+; CHECK:       [[EHCLEANUP]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = cleanuppad within none []
+; CHECK-NEXT:    cleanupret from [[TMP0]] unwind to caller
+; CHECK:       [[BODY_1]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast ptr [[PTR]] to ptr
+; CHECK-NEXT:    store <4 x i32> zeroinitializer, ptr [[TMP1]], align 4
+; CHECK-NEXT:    [[V0:%.*]] = load <4 x i32>, ptr [[PTR]], align 4
+; CHECK-NEXT:    store <4 x i32> zeroinitializer, ptr [[PTR]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %ptr = invoke ptr null(i64 0) to label %body.1 unwind label %ehcleanup
+
+body.2:
+  %v0 = load <4 x i32>, ptr %2, align 4
+  store <4 x i32> zeroinitializer, ptr %2, align 4
+  ret void
+
+ehcleanup:
+  %1 = cleanuppad within none []
+  cleanupret from %1 unwind to caller
+
+body.1:
+  %2 = getelementptr { i32 }, ptr %ptr, i64 0, i32 0
+  store <4 x i32> zeroinitializer, ptr %2, align 4
+  br label %body.2
+}