[IR] Allow poison argument to lifetime markers (#151148)

This slightly relaxes the invariant established in #149310, by also
allowing the lifetime argument to be poison. This is to support the
typical pattern of RAUWing with poison when removing an instruction.

It's worth noting that this does not require any conservative
assumptions, lifetimes with poison arguments can simply be skipped.

Fixes #151119.

Author: nikic

llvm/docs/LangRef.rst

@@ -26658,14 +26658,17 @@ Arguments:
 The first argument is a constant integer, which is ignored and will be removed
 in the future.
 
-The second argument is a pointer to an ``alloca`` instruction.
+The second argument is either a pointer to an ``alloca`` instruction or
+a ``poison`` value.
 
 Semantics:
 """"""""""
 
-The stack-allocated object that ``ptr`` points to is initially marked as dead.
-After '``llvm.lifetime.start``', the stack object is marked as alive and has an
-uninitialized value.
+If ``ptr`` is a ``poison`` value, the intrinsic has no effect.
+
+Otherwise, the stack-allocated object that ``ptr`` points to is initially
+marked as dead. After '``llvm.lifetime.start``', the stack object is marked as
+alive and has an uninitialized value.
 The stack object is marked as dead when either
 :ref:`llvm.lifetime.end <int_lifeend>` to the alloca is executed or the
 function returns.
@@ -26699,13 +26702,16 @@ Arguments:
 The first argument is a constant integer, which is ignored and will be removed
 in the future.
 
-The second argument is a pointer to an ``alloca`` instruction.
+The second argument is either a pointer to an ``alloca`` instruction or
+a ``poison`` value.
 
 Semantics:
 """"""""""
 
-The stack-allocated object that ``ptr`` points to becomes dead after the call
-to this intrinsic.
+If ``ptr`` is a ``poison`` value, the intrinsic has no effect.
+
+Otherwise, the stack-allocated object that ``ptr`` points to becomes dead after
+the call to this intrinsic.
 
 Calling ``llvm.lifetime.end`` on an already dead alloca is no-op.
 

llvm/lib/Analysis/StackLifetime.cpp

@@ -63,7 +63,10 @@ bool StackLifetime::isAliveAfter(const AllocaInst *AI,
 // markers has the same size and points to the alloca start.
 static const AllocaInst *findMatchingAlloca(const IntrinsicInst &II,
                                             const DataLayout &DL) {
-  const AllocaInst *AI = cast<AllocaInst>(II.getArgOperand(1));
+  const AllocaInst *AI = dyn_cast<AllocaInst>(II.getArgOperand(1));
+  if (!AI)
+    return nullptr;
+
   auto AllocaSize = AI->getAllocationSize(DL);
   if (!AllocaSize)
     return nullptr;

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -2189,8 +2189,8 @@ bool IRTranslator::translateKnownIntrinsic(const CallInst &CI, Intrinsic::ID ID,
     unsigned Op = ID == Intrinsic::lifetime_start ? TargetOpcode::LIFETIME_START
                                                   : TargetOpcode::LIFETIME_END;
 
-    const AllocaInst *AI = cast<AllocaInst>(CI.getArgOperand(1));
-    if (!AI->isStaticAlloca())
+    const AllocaInst *AI = dyn_cast<AllocaInst>(CI.getArgOperand(1));
+    if (!AI || !AI->isStaticAlloca())
       return true;
 
     MIRBuilder.buildInstr(Op).addFrameIndex(getOrCreateFrameIndex(*AI));

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -7598,7 +7598,9 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     if (TM.getOptLevel() == CodeGenOptLevel::None)
       return;
 
-    const AllocaInst *LifetimeObject = cast<AllocaInst>(I.getArgOperand(1));
+    const AllocaInst *LifetimeObject = dyn_cast<AllocaInst>(I.getArgOperand(1));
+    if (!LifetimeObject)
+      return;
 
     // First check that the Alloca is static, otherwise it won't have a
     // valid frame index.

llvm/lib/IR/Verifier.cpp

@@ -6769,10 +6769,13 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
     break;
   }
   case Intrinsic::lifetime_start:
-  case Intrinsic::lifetime_end:
-    Check(isa<AllocaInst>(Call.getArgOperand(1)),
-          "llvm.lifetime.start/end can only be used on alloca", &Call);
+  case Intrinsic::lifetime_end: {
+    Value *Ptr = Call.getArgOperand(1);
+    Check(isa<AllocaInst>(Ptr) || isa<PoisonValue>(Ptr),
+          "llvm.lifetime.start/end can only be used on alloca or poison",
+          &Call);
     break;
+  }
   };
 
   // Verify that there aren't any unmediated control transfers between funclets.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -1222,9 +1222,9 @@ struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
         !ConstantInt::isValueValidForType(IntptrTy, SizeValue))
       return;
     // Find alloca instruction that corresponds to llvm.lifetime argument.
-    AllocaInst *AI = cast<AllocaInst>(II.getArgOperand(1));
+    AllocaInst *AI = dyn_cast<AllocaInst>(II.getArgOperand(1));
     // We're interested only in allocas we can handle.
-    if (!ASan.isInterestingAlloca(*AI))
+    if (!AI || !ASan.isInterestingAlloca(*AI))
       return;
     bool DoPoison = (ID == Intrinsic::lifetime_end);
     AllocaPoisonCall APC = {&II, AI, SizeValue, DoPoison};

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -3301,8 +3301,9 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   void handleLifetimeStart(IntrinsicInst &I) {
     if (!PoisonStack)
       return;
-    AllocaInst *AI = cast<AllocaInst>(I.getArgOperand(1));
-    LifetimeStartList.push_back(std::make_pair(&I, AI));
+    AllocaInst *AI = dyn_cast<AllocaInst>(I.getArgOperand(1));
+    if (AI)
+      LifetimeStartList.push_back(std::make_pair(&I, AI));
   }
 
   void handleBswap(IntrinsicInst &I) {

llvm/lib/Transforms/Utils/Local.cpp

@@ -482,6 +482,9 @@ bool llvm::wouldInstructionBeTriviallyDead(const Instruction *I,
 
     if (II->isLifetimeStartOrEnd()) {
       auto *Arg = II->getArgOperand(1);
+      if (isa<PoisonValue>(Arg))
+        return true;
+
       // If the only uses of the alloca are lifetime intrinsics, then the
       // intrinsics are dead.
       return llvm::all_of(Arg->uses(), [](Use &Use) {

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

@@ -155,8 +155,9 @@ void StackInfoBuilder::visit(OptimizationRemarkEmitter &ORE,
     return;
   }
   if (auto *II = dyn_cast<LifetimeIntrinsic>(&Inst)) {
-    AllocaInst *AI = cast<AllocaInst>(II->getArgOperand(1));
-    if (getAllocaInterestingness(*AI) != AllocaInterestingness::kInteresting)
+    AllocaInst *AI = dyn_cast<AllocaInst>(II->getArgOperand(1));
+    if (!AI ||
+        getAllocaInterestingness(*AI) != AllocaInterestingness::kInteresting)
       return;
     if (II->getIntrinsicID() == Intrinsic::lifetime_start)
       Info.AllocasToInstrument[AI].LifetimeStart.push_back(II);

llvm/test/CodeGen/AArch64/lifetime-poison.ll

@@ -0,0 +1,14 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=aarch64 -global-isel=0 < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64 -global-isel=1 < %s | FileCheck %s
+
+; Check that lifetime.start/end with poison argument are ignored.
+
+define void @test() {
+; CHECK-LABEL: test:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ret
+  call void @llvm.lifetime.start.p0(i64 4, ptr poison)
+  call void @llvm.lifetime.end.p0(i64 4, ptr poison)
+  ret void
+}