[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

 //===- StackLifetime.h - Alloca Lifetime Analysis --------------*- C++ -*--===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_STACKLIFETIME_H
 #define LLVM_ANALYSIS_STACKLIFETIME_H
 
 #include "llvm/ADT/ArrayRef.h"
 #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"
 #include <cassert>
 #include <utility>
 
 namespace llvm {
 
bool overlaps(const LiveRange &Other) const {
       return Bits.anyCommon(Other.Bits);
     }
 
     void join(const LiveRange &Other) { Bits |= Other.Bits; }
 
     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>;
   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,
+                LivenessType Type);
 
   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;

 inline raw_ostream &operator<<(raw_ostream &OS,
                                const StackLifetime::LiveRange &R) {
   return OS << R.Bits;
 }
 
 /// 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);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_STACKLIFETIME_H

llvm/lib/Analysis/StackLifetime.cpp

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;
+        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
       // 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

 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,
+                             LivenessType Type)
+    : F(F), Type(Type), Allocas(Allocas), NumAllocas(Allocas.size()) {
   LLVM_DEBUG(dumpAllocas());
 
   for (unsigned I = 0; I < NumAllocas; ++I)
     AllocaNumbering[Allocas[I]] = I;
 
   collectMarkers();
 }
 

 }
 
 PreservedAnalyses StackLifetimePrinterPass::run(Function &F,
                                                 FunctionAnalysisManager &AM) {
   SmallVector<const AllocaInst *, 8> Allocas;
   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();
 }

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, StackLifetime::LivenessType::May);
   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/lib/Passes/PassBuilder.cpp

if (ParamName == "pre") {
       return make_error<StringError>(
           formatv("invalid GVN pass parameter '{0}' ", ParamName).str(),
           inconvertibleErrorCode());
     }
   }
   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
 /// alias.
 static bool startsWithDefaultPipelineAliasPrefix(StringRef Name) {
   return Name.startswith("default") || Name.startswith("thinlto") ||
          Name.startswith("lto");
 }

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("print<demanded-bits>", DemandedBitsPrinterPass(dbgs()))
 FUNCTION_PASS("print<domfrontier>", DominanceFrontierPrinterPass(dbgs()))
 FUNCTION_PASS("print<loops>", LoopPrinterPass(dbgs()))
 FUNCTION_PASS("print<memoryssa>", MemorySSAPrinterPass(dbgs()))
 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())
 FUNCTION_PASS("sink", SinkingPass())
 FUNCTION_PASS("slp-vectorizer", SLPVectorizerPass())
 FUNCTION_PASS("speculative-execution", SpeculativeExecutionPass())
 FUNCTION_PASS("spec-phis", SpeculateAroundPHIsPass())
 FUNCTION_PASS("sroa", SROA())
FUNCTION_PASS_WITH_PARAMS("mldst-motion",
                              return MergedLoadStoreMotionPass(Opts);
                            },
                            parseMergedLoadStoreMotionOptions)
 FUNCTION_PASS_WITH_PARAMS("gvn",
                            [](GVNOptions Opts) {
                              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
 #define LOOP_ANALYSIS(NAME, CREATE_PASS)
 #endif
 LOOP_ANALYSIS("no-op-loop", NoOpLoopAnalysis())
 LOOP_ANALYSIS("access-info", LoopAccessAnalysis())
 LOOP_ANALYSIS("ddg", DDGAnalysis())

llvm/test/Analysis/StackSafetyAnalysis/lifetime.ll

-; 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()
 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]*)