[llvm] [ControlHeightReduction] Drop lifetime annotations where necessary (PR #159686)

[Aiden Grossman via llvm-commits]

https://github.com/boomanaiden154 created https://github.com/llvm/llvm-project/pull/159686

ControlHeightReduction will duplicate some blocks and insert phi nodes in exit blocks of regions that it operates on for any live values. This includes allocas. Having a lifetime annotation refer to a phi node was made illegal in 92c55a315eab455d5fed2625fe0f61f88cb25499, which causes the verifier to fail after CHR.

There are some cases where we might not need to drop lifetime annotations (usually because we do not need the phi to begin with), but drop all annotations for now to be conservative. I believe this also ends up being a correctness issue on top of the assertion because the lifetime annotation in some cases would get shifted to after the store, making the store UB when it was not before.

Fixes #159621.

>From 9062f71967b38560eacfdcb6e50ad35e43e03ad1 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Fri, 19 Sep 2025 01:36:25 +0000
Subject: [PATCH] [ControlHeightReduction] Drop lifetime annotations where
 necessary

ControlHeightReduction will duplicate some blocks and insert phi nodes
in exit blocks of regions that it operates on for any live values. This
includes allocas. Having a lifetime annotation refer to a phi node was
made illegal in 92c55a315eab455d5fed2625fe0f61f88cb25499, which causes
the verifier to fail after CHR.

There are some cases where we might not need to drop lifetime
annotations (usually because we do not need the phi to begin with), but
drop all annotations for now to be conservative. I believe this also
ends up being a correctness issue on top of the assertion because the
lifetime annotation in some cases would get shifted to after the store,
making the store UB when it was not before.

Fixes #159621.
---
 .../ControlHeightReduction.cpp                |  6 ++
 llvm/test/Transforms/PGOProfile/chr.ll        | 59 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp b/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp
index c14bbecf0d4e1..3e05b8539dcb8 100644
--- a/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp
+++ b/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp
@@ -1592,6 +1592,12 @@ static void insertTrivialPHIs(CHRScope *Scope,
         TrivialPHIs.insert(PN);
         CHR_DEBUG(dbgs() << "Insert phi " << *PN << "\n");
         for (Instruction *UI : Users) {
+          // Drop lifetime annotations as it is illegal for them to refer to a
+          // phi node.
+          if (UI->isLifetimeStartOrEnd()) {
+            UI->eraseFromParent();
+            continue;
+          }
           for (unsigned J = 0, NumOps = UI->getNumOperands(); J < NumOps; ++J) {
             if (UI->getOperand(J) == &I) {
               UI->setOperand(J, PN);
diff --git a/llvm/test/Transforms/PGOProfile/chr.ll b/llvm/test/Transforms/PGOProfile/chr.ll
index f0a1574c5f209..860e40917a77f 100644
--- a/llvm/test/Transforms/PGOProfile/chr.ll
+++ b/llvm/test/Transforms/PGOProfile/chr.ll
@@ -2678,6 +2678,65 @@ bb6:
   ret void
 }
 
+; Test that when we create phi nodes for allocas, we drop lifetime annotations
+; that would otherwise be illegally referring to a phi node.
+declare void @baz(i64)
+declare void @llvm.lifetime.start.p0(ptr captures(none))
+define void @test_chr_with_lifetimes(ptr %i) !prof !14 {
+; CHECK-LABEL: @test_chr_with_lifetimes(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[I:%.*]], align 4
+; CHECK-NEXT:    [[DOTFR:%.*]] = freeze i32 [[TMP0]]
+; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp eq i32 [[DOTFR]], 0
+; CHECK-NEXT:    br i1 [[DOTNOT]], label [[ENTRY_SPLIT_NONCHR:%.*]], label [[ENTRY_SPLIT:%.*]], !prof [[PROF21:![0-9]+]]
+; CHECK:       entry.split:
+; CHECK-NEXT:    [[TEST:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    call void @baz(i64 0)
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       entry.split.nonchr:
+; CHECK-NEXT:    [[TEST_NONCHR:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    call void @baz(i64 4)
+; CHECK-NEXT:    br label [[BB1]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[TMP1:%.*]] = phi ptr [ [[TEST]], [[ENTRY_SPLIT]] ], [ [[TEST_NONCHR]], [[ENTRY_SPLIT_NONCHR]] ]
+; CHECK-NEXT:    store ptr [[TMP1]], ptr [[I]], align 8
+; CHECK-NEXT:    br label [[BB2:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[TMP2:%.*]] = phi ptr [ [[TMP3:%.*]], [[BB2]] ], [ null, [[BB1]] ]
+; CHECK-NEXT:    [[TMP3]] = getelementptr i8, ptr [[TMP2]], i64 24
+; CHECK-NEXT:    [[TMP4:%.*]] = icmp eq ptr [[TMP2]], [[I]]
+; CHECK-NEXT:    br i1 [[TMP4]], label [[BB3:%.*]], label [[BB2]]
+; CHECK:       bb3:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %1 = load i32, ptr %i
+  %2 = icmp eq i32 %1, 0
+  %3 = select i1 %2, i64 4, i64 0, !prof !15
+  %test = alloca i32, align 8
+  call void @baz(i64 %3)
+  br i1 %2, label %bb1, label %bb0, !prof !15
+
+bb0:
+  call void @foo()
+  br label %bb1
+
+bb1:
+  call void @llvm.lifetime.start.p0(ptr %test)
+  store ptr %test, ptr %i, align 8
+  br label %bb2
+
+bb2:
+  %4 = phi ptr [ %5, %bb2 ], [ null, %bb1 ]
+  %5 = getelementptr i8, ptr %4, i64 24
+  %6 = icmp eq ptr %4, %i
+  br i1 %6, label %bb3, label %bb2
+
+bb3:
+  ret void
+}
+
 
 !llvm.module.flags = !{!0}
 !0 = !{i32 1, !"ProfileSummary", !1}