[StackSafety] Add "Must Live" logic

[vitalybuka]

Summary

Extend StackLifetime with option to calculate liveliness
where alloca is only considered alive on basic block entry
if all non-dead predecessors had it alive at terminators.

Depends on D82043.

[eugenis]

llvm/test/Analysis/StackSafetyAnalysis/lifetime.ll

2

It would be nicer to make it a parametrized pass, something like this:

-passes=stack-lifetime<must-live> or stack-lifetime<may-live>

And turn the bool argument into an enum with named options.

[vitalybuka]

@eugenis comments

[eugenis]

LGTM

llvm/include/llvm/Analysis/StackLifetime.h

bool overlaps(const LiveRange &Other) const {
       return Bits.anyCommon(Other.Bits);
     }
 
     void join(const LiveRange &Other) { Bits |= Other.Bits; }
   };
 
 private:
   const Function &F;
+  const bool MustLive;
 
   /// Maps active slots (per bit) for each basic block.
   using LivenessMap = DenseMap<const BasicBlock *, BlockLifetimeInfo>;
   LivenessMap BlockLiveness;
 
   /// Number of interesting instructions.
   int NumInst = -1;
 
private:
   void dumpBlockLiveness() const;
   void dumpLiveRanges() const;
 
   void collectMarkers();
   void calculateLocalLiveness();
   void calculateLiveIntervals();
 
 public:
-  StackLifetime(const Function &F, ArrayRef<const AllocaInst *> Allocas);
+  StackLifetime(const Function &F, ArrayRef<const AllocaInst *> Allocas,
+                bool MustLive);
 
   void run();
   std::vector<const IntrinsicInst *> getMarkers() const;
 
   /// Returns a set of "interesting" instructions where the given alloca is
   /// live. Not all instructions in a function are interesting: we pick a set
   /// that is large enough for LiveRange::Overlaps to be correct.
   const LiveRange &getLiveRange(const AllocaInst *AI) const;

llvm/lib/Analysis/StackLifetime.cpp

 #include "llvm/Support/FormattedStream.h"
 #include <memory>
 #include <tuple>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "stack-lifetime"
 
+static cl::opt<bool> ClMustLive("stack-lifetime-must", cl::Hidden,
+                                cl::init(false));
+
 const StackLifetime::LiveRange &
 StackLifetime::getLiveRange(const AllocaInst *AI) const {
   const auto IT = AllocaNumbering.find(AI);
   assert(IT != AllocaNumbering.end());
   return LiveRanges[IT->second];
 }
 
 static bool readMarker(const Instruction *I, bool *IsStart) {
for (const BasicBlock *BB : depth_first(&F)) {
 
       // Compute LiveIn by unioning together the LiveOut sets of all preds.
       BitVector LocalLiveIn;
       for (auto *PredBB : predecessors(BB)) {
         LivenessMap::const_iterator I = BlockLiveness.find(PredBB);
         // If a predecessor is unreachable, ignore it.
         if (I == BlockLiveness.end())
           continue;
+        if (MustLive) {
+          if (LocalLiveIn.empty())
+            LocalLiveIn = I->second.LiveOut;
+          else
+            LocalLiveIn &= I->second.LiveOut;
+        } else {
           LocalLiveIn |= I->second.LiveOut;
         }
+      }
 
       // Compute LiveOut by subtracting out lifetimes that end in this
       // block, then adding in lifetimes that begin in this block.  If
       // we have both BEGIN and END markers in the same basic block
       // then we know that the BEGIN marker comes after the END,
       // because we already handle the case where the BEGIN comes
       // before the END when collecting the markers (and building the
       // BEGIN/END vectors).

 LLVM_DUMP_METHOD void StackLifetime::dumpLiveRanges() const {
   dbgs() << "Alloca liveness:\n";
   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
     dbgs() << "  " << AllocaNo << ": " << LiveRanges[AllocaNo] << "\n";
 }
 #endif
 
 StackLifetime::StackLifetime(const Function &F,
-                             ArrayRef<const AllocaInst *> Allocas)
-    : F(F), Allocas(Allocas), NumAllocas(Allocas.size()) {
+                             ArrayRef<const AllocaInst *> Allocas,
+                             bool MustLive)
+    : F(F), MustLive(MustLive), Allocas(Allocas), NumAllocas(Allocas.size()) {
   LLVM_DEBUG(dumpAllocas());
 
   for (unsigned I = 0; I < NumAllocas; ++I)
     AllocaNumbering[Allocas[I]] = I;
 
   collectMarkers();
 }
 
void StackLifetimePass::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
 }
 
 bool StackLifetimePass::runOnFunction(Function &F) {
   SmallVector<const AllocaInst *, 8> Allocas;
   for (auto &I : instructions(F))
     if (const AllocaInst *AI = dyn_cast<AllocaInst>(&I))
       Allocas.push_back(AI);
-  Result = std::make_unique<StackLifetime>(F, Allocas);
+  Result = std::make_unique<StackLifetime>(F, Allocas, ClMustLive);
   Result->run();
   return false;
 }
 
 } // namespace
 } // namespace llvm
 
 static const char LocalPassName[] = "Stack Lifetime Analysis";
 INITIALIZE_PASS(StackLifetimePass, DEBUG_TYPE, LocalPassName, false, true)
 No newline at end of file

llvm/lib/CodeGen/SafeStack.cpp

Value *SafeStack::moveStaticAllocasToUnsafeStack(
     IRBuilder<> &IRB, Function &F, ArrayRef<AllocaInst *> StaticAllocas,
     ArrayRef<Argument *> ByValArguments, ArrayRef<ReturnInst *> Returns,
     Instruction *BasePointer, AllocaInst *StackGuardSlot) {
   if (StaticAllocas.empty() && ByValArguments.empty())
     return BasePointer;
 
   DIBuilder DIB(*F.getParent());
 
-  StackLifetime SSC(F, StaticAllocas);
+  StackLifetime SSC(F, StaticAllocas, false);
   static const StackLifetime::LiveRange NoColoringRange(1, true);
   if (ClColoring)
     SSC.run();
 
   for (auto *I : SSC.getMarkers()) {
     auto *Op = dyn_cast<Instruction>(I->getOperand(1));
     const_cast<IntrinsicInst *>(I)->eraseFromParent();
     // Remove the operand bitcast, too, if it has no more uses left.

llvm/test/Analysis/StackSafetyAnalysis/lifetime.ll

-; RUN: opt -analyze -stack-lifetime -S < %s -o - | FileCheck %s
+; RUN: opt -analyze -stack-lifetime -stack-lifetime-must=0 -S < %s -o - | FileCheck %s --check-prefixes=CHECK,MAY
+; RUN: opt -analyze -stack-lifetime -stack-lifetime-must=1 -S < %s -o - | FileCheck %s --check-prefixes=CHECK,MUST
 
 define void @f() {
 ; CHECK-LABEL: define void @f()
 entry:
 ; CHECK: entry:
 ; CHECK-NEXT: Alive: <>
   %x = alloca i32, align 4
   %y = alloca i32, align 4

 ; CHECK-NEXT: Alive: <x>
 
   br label %l2
 ; CHECK: br label %l2
 ; CHECK-NEXT: Alive: <x>
 
 l2:                                               ; preds = %l2, %entry
 ; CHECK: l2:
-; CHECK-NEXT: Alive: <x>
+; MAY-NEXT: Alive: <x>
+; MUST-NEXT: Alive: <>
   call void @capture8(i8* %x)
   call void @llvm.lifetime.end.p0i8(i64 4, i8* %x)
 ; CHECK: call void @llvm.lifetime.end.p0i8(i64 4, i8* %x)
 ; CHECK-NEXT: Alive: <>
 
   br label %l2
 ; CHECK: br label %l2
 ; CHECK-NEXT: Alive: <>

 ; CHECK-NEXT: Alive: <x>
 
   call void @capture8(i8* %x)
   br label %l2
 ; CHECK: br label %l2
 ; CHECK-NEXT: Alive: <x>
 }
 
+define void @if_must(i1 %a) {
+; CHECK-LABEL: define void @if_must
+entry:
+; CHECK: entry:
+; CHECK-NEXT: Alive: <>
+  %x = alloca i8, align 4
+  %y = alloca i8, align 4
+
+  br i1 %a, label %if.then, label %if.else
+; CHECK: br i1 %a
+; CHECK-NEXT: Alive: <>
+
+if.then:
+; CHECK: if.then:
+; CHECK-NEXT: Alive: <>
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %y)
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 4, i8* %y)
+; CHECK-NEXT: Alive: <y>
+
+  br label %if.end
+; CHECK: br label %if.end
+; CHECK-NEXT: Alive: <y>
+
+if.else:
+; CHECK: if.else:
+; CHECK-NEXT: Alive: <>
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %y)
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 4, i8* %y)
+; CHECK-NEXT: Alive: <y>
+
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %x)
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 4, i8* %x)
+; CHECK-NEXT: Alive: <x y>
+
+  br label %if.end
+; CHECK: br label %if.end
+; CHECK-NEXT: Alive: <x y>
+
+if.end:
+; CHECK: if.end:
+; MAY-NEXT: Alive: <x y>
+; MUST-NEXT: Alive: <y>
+
+ret void
+; CHECK: ret void
+; MAY-NEXT: Alive: <x y>
+; MUST-NEXT: Alive: <y>
+}
+
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
 declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)
 declare void @capture8(i8*)
 declare void @capture32(i32*)
 declare void @capture64(i64*)
 declare void @capture100x32([100 x i32]*)