[StackSafety] Add "Must Live" logic

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.

Reviewers: eugenis

Reviewed By: eugenis

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D82124

include/llvm/Analysis/StackLifetime.h

@@ -13,6 +13,7 @@
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/raw_ostream.h"
@@ -78,8 +79,16 @@ public:
     bool test(unsigned Idx) const { return Bits.test(Idx); }
   };
 
+  // Controls what is "alive" if control flow may reach the instruction
+  // with a different liveness of the alloca.
+  enum class LivenessType {
+    May,  // May be alive on some path.
+    Must, // Must be alive on every path.
+  };
+
 private:
   const Function &F;
+  LivenessType Type;
 
   /// Maps active slots (per bit) for each basic block.
   using LivenessMap = DenseMap<const BasicBlock *, BlockLifetimeInfo>;
@@ -124,7 +133,8 @@ private:
   void calculateLiveIntervals();
 
 public:
-  StackLifetime(const Function &F, ArrayRef<const AllocaInst *> Allocas);
+  StackLifetime(const Function &F, ArrayRef<const AllocaInst *> Allocas,
+                LivenessType Type);
 
   void run();
   std::vector<const IntrinsicInst *> getMarkers() const;
@@ -168,10 +178,12 @@ inline raw_ostream &operator<<(raw_ostream &OS,
 /// Printer pass for testing.
 class StackLifetimePrinterPass
     : public PassInfoMixin<StackLifetimePrinterPass> {
+  StackLifetime::LivenessType Type;
   raw_ostream &OS;
 
 public:
-  explicit StackLifetimePrinterPass(raw_ostream &OS) : OS(OS) {}
+  StackLifetimePrinterPass(raw_ostream &OS, StackLifetime::LivenessType Type)
+      : Type(Type), OS(OS) {}
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 

lib/Analysis/StackLifetime.cpp

@@ -166,7 +166,17 @@ void StackLifetime::calculateLocalLiveness() {
         // If a predecessor is unreachable, ignore it.
         if (I == BlockLiveness.end())
           continue;
-        LocalLiveIn |= I->second.LiveOut;
+        switch (Type) {
+        case LivenessType::May:
+          LocalLiveIn |= I->second.LiveOut;
+          break;
+        case LivenessType::Must:
+          if (LocalLiveIn.empty())
+            LocalLiveIn = I->second.LiveOut;
+          else
+            LocalLiveIn &= I->second.LiveOut;
+          break;
+        }
       }
 
       // Compute LiveOut by subtracting out lifetimes that end in this
@@ -272,8 +282,9 @@ LLVM_DUMP_METHOD void StackLifetime::dumpLiveRanges() const {
 #endif
 
 StackLifetime::StackLifetime(const Function &F,
-                             ArrayRef<const AllocaInst *> Allocas)
-    : F(F), Allocas(Allocas), NumAllocas(Allocas.size()) {
+                             ArrayRef<const AllocaInst *> Allocas,
+                             LivenessType Type)
+    : F(F), Type(Type), Allocas(Allocas), NumAllocas(Allocas.size()) {
   LLVM_DEBUG(dumpAllocas());
 
   for (unsigned I = 0; I < NumAllocas; ++I)
@@ -351,7 +362,7 @@ PreservedAnalyses StackLifetimePrinterPass::run(Function &F,
   for (auto &I : instructions(F))
     if (const AllocaInst *AI = dyn_cast<AllocaInst>(&I))
       Allocas.push_back(AI);
-  StackLifetime SL(F, Allocas);
+  StackLifetime SL(F, Allocas, Type);
   SL.run();
   SL.print(OS);
   return PreservedAnalyses::all();

lib/CodeGen/SafeStack.cpp

@@ -497,7 +497,7 @@ Value *SafeStack::moveStaticAllocasToUnsafeStack(
 
   DIBuilder DIB(*F.getParent());
 
-  StackLifetime SSC(F, StaticAllocas);
+  StackLifetime SSC(F, StaticAllocas, StackLifetime::LivenessType::May);
   static const StackLifetime::LiveRange NoColoringRange(1, true);
   if (ClColoring)
     SSC.run();

lib/Passes/PassBuilder.cpp

@@ -1851,6 +1851,26 @@ Expected<GVNOptions> parseGVNOptions(StringRef Params) {
   return Result;
 }
 
+Expected<StackLifetime::LivenessType>
+parseStackLifetimeOptions(StringRef Params) {
+  StackLifetime::LivenessType Result = StackLifetime::LivenessType::May;
+  while (!Params.empty()) {
+    StringRef ParamName;
+    std::tie(ParamName, Params) = Params.split(';');
+
+    if (ParamName == "may") {
+      Result = StackLifetime::LivenessType::May;
+    } else if (ParamName == "must") {
+      Result = StackLifetime::LivenessType::Must;
+    } else {
+      return make_error<StringError>(
+          formatv("invalid StackLifetime parameter '{0}' ", ParamName).str(),
+          inconvertibleErrorCode());
+    }
+  }
+  return Result;
+}
+
 } // namespace
 
 /// Tests whether a pass name starts with a valid prefix for a default pipeline

lib/Passes/PassRegistry.def

@@ -239,7 +239,6 @@ FUNCTION_PASS("print<phi-values>", PhiValuesPrinterPass(dbgs()))
 FUNCTION_PASS("print<regions>", RegionInfoPrinterPass(dbgs()))
 FUNCTION_PASS("print<scalar-evolution>", ScalarEvolutionPrinterPass(dbgs()))
 FUNCTION_PASS("print<stack-safety-local>", StackSafetyPrinterPass(dbgs()))
-FUNCTION_PASS("print<stack-lifetime>", StackLifetimePrinterPass(dbgs()))
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("scalarizer", ScalarizerPass())
 FUNCTION_PASS("sccp", SCCPPass())
@@ -302,6 +301,11 @@ FUNCTION_PASS_WITH_PARAMS("gvn",
                              return GVN(Opts);
                            },
                            parseGVNOptions)
+FUNCTION_PASS_WITH_PARAMS("print<stack-lifetime>",
+                           [](StackLifetime::LivenessType Type) {
+                             return StackLifetimePrinterPass(dbgs(), Type);
+                           },
+                           parseStackLifetimeOptions)
 #undef FUNCTION_PASS_WITH_PARAMS
 
 #ifndef LOOP_ANALYSIS

test/Analysis/StackSafetyAnalysis/lifetime.ll

@@ -1,4 +1,5 @@
-; RUN: opt -passes='print<stack-lifetime>' -disable-output %s 2>&1 | FileCheck %s
+; RUN: opt -passes='print<stack-lifetime><may>' -disable-output %s 2>&1 | FileCheck %s --check-prefixes=CHECK,MAY
+; RUN: opt -passes='print<stack-lifetime><must>' -disable-output %s 2>&1 | FileCheck %s --check-prefixes=CHECK,MUST
 
 define void @f() {
 ; CHECK-LABEL: define void @f()
@@ -710,7 +711,8 @@ entry:
 
 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)
@@ -758,6 +760,55 @@ l2:                                               ; preds = %l2, %entry
 ; 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*)