Prevent LCSSA from creating phi operands for lifetime intrinsics.

[clin111]

If an alloca has a lifetime intrinsic User outside of a loop, preserve LCSSA by removing all the lifetime intrinsics for the alloca.

This prevents a problematic phi operand from being created for the lifetime intrinsic, and avoids potential stack coloring issues.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Chang Lin (clin111)

Changes

If an alloca has a lifetime intrinsic User outside of a loop, preserve LCSSA by removing all the lifetime intrinsics for the alloca.

This prevents a problematic phi operand from being created for the lifetime intrinsic, and avoids potential stack coloring issues.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/LCSSA.cpp (+33)
• (added) llvm/test/Transforms/LCSSA/lifetime-intrinsic.ll (+30)

diff --git a/llvm/lib/Transforms/Utils/LCSSA.cpp b/llvm/lib/Transforms/Utils/LCSSA.cpp
index a9e08ada82ca0..7a6c1b6b2c1f4 100644
--- a/llvm/lib/Transforms/Utils/LCSSA.cpp
+++ b/llvm/lib/Transforms/Utils/LCSSA.cpp
@@ -41,6 +41,7 @@
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PredIteratorCache.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
@@ -75,6 +76,36 @@ static bool isExitBlock(BasicBlock *BB,
 // expensive, and we're not mutating the loop structure.
 using LoopExitBlocksTy = SmallDenseMap<Loop *, SmallVector<BasicBlock *, 1>>;
 
+// If I is an alloca with lifetime intrinsics that are live out of the loop,
+// remove all the lifetime intrinsics for I.
+// This ensures we don't create a lifetime intrinsic based on an LCSSA phi,
+// and avoids potential lifetime inconsistencies.
+static void fixLifetimeIntrinsics(Instruction *I,
+                                  SmallVectorImpl<Use *> &UsesToRewrite) {
+  if (!isa<AllocaInst>(I))
+    return;
+  bool RemovedAny = false;
+  // First, remove lifetime intrinsics from UsesToRewrite.
+  llvm::erase_if(UsesToRewrite, [&](Use *U) {
+    if (auto *II = dyn_cast<IntrinsicInst>(U->getUser())) {
+      if (II->isLifetimeStartOrEnd()) {
+        RemovedAny = true;
+        return true;
+      }
+    }
+    return false;
+  });
+
+  // Ensure consistency in the simplest way, by removing all lifetime uses
+  // of I.
+  if (RemovedAny) {
+    for (auto *U : make_early_inc_range(I->users()))
+      if (auto *II = dyn_cast<IntrinsicInst>(U))
+        if (II->isLifetimeStartOrEnd())
+          II->eraseFromParent();
+  }
+}
+
 /// For every instruction from the worklist, check to see if it has any uses
 /// that are outside the current loop.  If so, insert LCSSA PHI nodes and
 /// rewrite the uses.
@@ -126,6 +157,8 @@ formLCSSAForInstructionsImpl(SmallVectorImpl<Instruction *> &Worklist,
         UsesToRewrite.push_back(&U);
     }
 
+    fixLifetimeIntrinsics(I, UsesToRewrite);
+
     // If there are no uses outside the loop, exit with no change.
     if (UsesToRewrite.empty())
       continue;
diff --git a/llvm/test/Transforms/LCSSA/lifetime-intrinsic.ll b/llvm/test/Transforms/LCSSA/lifetime-intrinsic.ll
new file mode 100644
index 0000000000000..fbdd4f4afa26f
--- /dev/null
+++ b/llvm/test/Transforms/LCSSA/lifetime-intrinsic.ll
@@ -0,0 +1,30 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes=lcssa -S < %s | FileCheck %s
+
+; Previously crashing in the verifier, due to LCSSA inserting a phi between
+; the alloca and the lifetime intrinsic.
+; We can instead remove the problematic intrinsic and its corresponding start.
+
+declare void @llvm.lifetime.end.p0(ptr captures(none))
+declare void @llvm.lifetime.start.p0(i64, ptr captures(none))
+
+define fastcc i32 @test() {
+; CHECK-LABEL: define fastcc i32 @test() {
+; CHECK-NEXT:    br label %[[BB1:.*]]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    br i1 false, label %[[BB2:.*]], label %[[BB1]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    ret i32 0
+;
+  br label %1
+
+1:                                                ; preds = %1, %0
+  %a = alloca i32, align 4
+  call void @llvm.lifetime.start.p0(i64 4, ptr %a)
+  br i1 false, label %3, label %1
+
+3:                                                ; preds = %1
+  call void @llvm.lifetime.end.p0(ptr %a)
+  ret i32 0
+}

[nikic]

This is an interesting problem, and another motivation to change our lifetime intrinsic design...

I'm a bit apprehensive about this approach of dropping the lifetime intrinsic users, because calling code may not expect LCSSA to modify IR in this way. Though from a cursory look, I think that the only place that would care is SCEVExpander, and that one shouldn't be tracking lifetime intrinsics in its internal structures.

An alternative here would be to instead treat alloca lifetime uses similar to how we treat instructions returning tokens, and exclude them from LCSSA formation. Of course, this means that we may have to adjust code that assumes only tokens are excluded from LCSSA.