[llvm] [CodeGenPrepare] Do not undo GEP unmerging in AddrSinkUsingGEPs (PR #197685)

[via llvm-commits]

https://github.com/MuellerMP created https://github.com/llvm/llvm-project/pull/197685

Fixes #197684.

Do not sink gep calculation after it has been unmerged to prevent an infinite loop.

>From 661e8057900368268a92cc5fda0f7956ac5b069d Mon Sep 17 00:00:00 2001
From: MuellerMP <mirkomueller97 at live.de>
Date: Thu, 14 May 2026 15:51:24 +0200
Subject: [PATCH] [CodeGenPrepare] Do not undo GEP unmerging in
 AddrSinkUsingGEPs

Fixes #197684.

Do not sink gep calculation after it has been unmerged to prevent an infinite loop.
---
 llvm/lib/CodeGen/CodeGenPrepare.cpp           |  9 +++++
 .../X86/indirect-br-gep-unmerge-sink.ll       | 37 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/indirect-br-gep-unmerge-sink.ll

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 74a0502d8cb7c..d5e1957ce0a84 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -6140,6 +6140,13 @@ bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
       } else {
         if (ResultPtr->getType() != I8PtrTy)
           ResultPtr = Builder.CreatePointerCast(ResultPtr, I8PtrTy);
+
+        if (auto *GEP = dyn_cast<GetElementPtrInst>(Addr)) {
+          // Do not revert the gep unmerging optimization.
+          if (GEP->getMetadata("cgp.unmerged_gep"))
+            return Modified;
+        }
+
         SunkAddr = Builder.CreatePtrAdd(ResultPtr, ResultIndex, "sunkaddr",
                                         AddrMode.InBounds);
       }
@@ -8824,6 +8831,8 @@ static bool tryUnmergingGEPsAcrossIndirectBr(GetElementPtrInst *GEPI,
     if (NewIdx.isNegative() && TargetFlags.hasNoUnsignedWrap())
       TargetFlags = TargetFlags.withoutNoUnsignedWrap();
     UGEPI->setNoWrapFlags(TargetFlags);
+    UGEPI->setMetadata("cgp.unmerged_gep",
+                       MDNode::get(UGEPI->getContext(), {}));
   }
   // After unmerging, verify that GEPIOp is actually only used in SrcBlock (not
   // alive on IndirectBr edges).
diff --git a/llvm/test/CodeGen/X86/indirect-br-gep-unmerge-sink.ll b/llvm/test/CodeGen/X86/indirect-br-gep-unmerge-sink.ll
new file mode 100644
index 0000000000000..8728fa51b29d3
--- /dev/null
+++ b/llvm/test/CodeGen/X86/indirect-br-gep-unmerge-sink.ll
@@ -0,0 +1,37 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes="require<profile-summary>,function(codegenprepare)" -mtriple=x86_64-linux-gnu %s -o - | FileCheck %s
+ at addrTable = global [2 x ptr] [ptr blockaddress(@testGepUnmergingDoesNotGetUndone, %exit1), ptr blockaddress(@testGepUnmergingDoesNotGetUndone, %exit2)]
+
+define void @testGepUnmergingDoesNotGetUndone() {
+; CHECK-LABEL: define void @testGepUnmergingDoesNotGetUndone() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[PTR:%.*]] = alloca i8, align 1
+; CHECK-NEXT:    [[BASEGEP:%.*]] = getelementptr i8, ptr [[PTR]], i64 16
+; CHECK-NEXT:    indirectbr ptr null, [label %[[EXIT1:.*]], label %[[EXIT2:.*]]]
+; CHECK:       [[EXIT1]]:
+; CHECK-NEXT:    call void @useOfGEP(ptr [[BASEGEP]])
+; CHECK-NEXT:    ret void
+; CHECK:       [[EXIT2]]:
+; CHECK-NEXT:    [[UNMERGEDGEP:%.*]] = getelementptr i8, ptr [[BASEGEP]], i64 24, !cgp.unmerged_gep [[META0:![0-9]+]]
+; CHECK-NEXT:    store i32 0, ptr [[UNMERGEDGEP]], align 8
+; CHECK-NEXT:    ret void
+;
+entry:
+  %ptr = alloca i8
+  %baseGEP = getelementptr i8, ptr %ptr, i64 16
+  indirectbr ptr null, [label %exit1, label %exit2]
+
+exit1:                                             ; preds = %entry
+  call void @useOfGEP(ptr %baseGEP)
+  ret void
+
+exit2:                                             ; preds = %entry
+  %unmergedGEP = getelementptr i8, ptr %ptr, i64 40
+  store i32 0, ptr %unmergedGEP, align 8
+  ret void
+}
+
+declare void @useOfGEP(ptr)
+;.
+; CHECK: [[META0]] = !{}
+;.