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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Aiden Grossman (boomanaiden154)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/159686.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp (+6) 
- (modified) llvm/test/Transforms/PGOProfile/chr.ll (+59) 


``````````diff
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}

``````````

</details>


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