SILOptimizer: restructure the apply(partial_apply) peephole and the dead partial_apply elimination optimizations

include/swift/SILOptimizer/Analysis/ARCAnalysis.h

@@ -392,47 +392,6 @@ class ConsumedArgToEpilogueReleaseMatcher {
   void processMatchingReleases();
 };
 
-class ReleaseTracker {
-  llvm::SmallSetVector<SILInstruction *, 4> TrackedUsers;
-  llvm::SmallSetVector<SILInstruction *, 4> FinalReleases;
-  std::function<bool(SILInstruction *)> AcceptableUserQuery;
-  std::function<bool(SILInstruction *)> TransitiveUserQuery;
-
-public:
-  ReleaseTracker(std::function<bool(SILInstruction *)> AcceptableUserQuery,
-                 std::function<bool(SILInstruction *)> TransitiveUserQuery)
-      : TrackedUsers(), FinalReleases(),
-        AcceptableUserQuery(AcceptableUserQuery),
-        TransitiveUserQuery(TransitiveUserQuery) {}
-
-  void trackLastRelease(SILInstruction *Inst) { FinalReleases.insert(Inst); }
-
-  bool isUserAcceptable(SILInstruction *User) const {
-    return AcceptableUserQuery(User);
-  }
-  bool isUserTransitive(SILInstruction *User) const {
-    return TransitiveUserQuery(User);
-  }
-
-  bool isUser(SILInstruction *User) { return TrackedUsers.count(User); }
-
-  void trackUser(SILInstruction *User) { TrackedUsers.insert(User); }
-
-  using range = iterator_range<llvm::SmallSetVector<SILInstruction *, 4>::iterator>;
-
-  // An ordered list of users, with "casts" before their transitive uses.
-  range getTrackedUsers() { return {TrackedUsers.begin(), TrackedUsers.end()}; }
-
-  range getFinalReleases() {
-    return {FinalReleases.begin(), FinalReleases.end()};
-  }
-};
-
-/// Return true if we can find a set of post-dominating final releases. Returns
-/// false otherwise. The FinalRelease set is placed in the out parameter
-/// FinalRelease.
-bool getFinalReleasesForValue(SILValue Value, ReleaseTracker &Tracker);
-
 /// Match a call to a trap BB with no ARC relevant side effects.
 bool isARCInertTrapBB(const SILBasicBlock *BB);
 

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -304,16 +304,37 @@ struct InstModCallbacks {
   InstModCallbacks(InstModCallbacks &&) = default;
 };
 
+/// Get all consumed arguments of a partial_apply.
+///
+/// These are basically all arguments, except inout arguments and arguments
+/// of trivial type.
+/// If \p includeTrivialAddrArgs is true, also trivial address-type arguments
+/// are included.
+void getConsumedPartialApplyArgs(PartialApplyInst *pai,
+                                 SmallVectorImpl<Operand *> &argOperands,
+                                 bool includeTrivialAddrArgs);
+
+/// Collect all (transitive) users of \p inst which just copy or destroy \p
+/// inst.
+///
+/// In other words: all users which do not prevent \p inst from being considered
+/// as "dead".
+/// Returns true, if there are no other users beside those collected in \p
+/// destroys, i.e. if \p inst can be considered as "dead".
+bool collectDestroys(SingleValueInstruction *inst,
+                     SmallVectorImpl<SILInstruction *> &destroys);
 /// If Closure is a partial_apply or thin_to_thick_function with only local
 /// ref count users and a set of post-dominating releases:
 ///
 /// 1. Remove all ref count operations and the closure.
-/// 2. Add each one of the last release locations insert releases for the
-///    captured args if we have a partial_apply.
+/// 2. At each one of the last release locations insert releases for the
+///    captured args if we have a partial_apply (except \p needKeepArgsAlive is
+///    false).
 ///
 /// In the future this should be extended to be less conservative with users.
 bool tryDeleteDeadClosure(SingleValueInstruction *closure,
-                          InstModCallbacks callbacks = InstModCallbacks());
+                          InstModCallbacks callbacks = InstModCallbacks(),
+                          bool needKeepArgsAlive = true);
 
 /// Given a SILValue argument to a partial apply \p Arg and the associated
 /// parameter info for that argument, perform the necessary cleanups to Arg when
@@ -522,9 +543,9 @@ findLocalApplySites(FunctionRefBaseInst *fri);
 /// Gets the base implementation of a method.
 AbstractFunctionDecl *getBaseMethod(AbstractFunctionDecl *FD);
 
-SILInstruction *
-tryOptimizeApplyOfPartialApply(PartialApplyInst *pai, SILBuilder &builder,
-                               InstModCallbacks callbacks = InstModCallbacks());
+bool tryOptimizeApplyOfPartialApply(
+    PartialApplyInst *pai, SILBuilderContext &builderCtxt,
+    InstModCallbacks callbacks = InstModCallbacks());
 
 } // end namespace swift
 

include/swift/SILOptimizer/Utils/ValueLifetime.h

@@ -79,6 +79,9 @@ class ValueLifetimeAnalysis {
   /// instructions of the frontier that are not in the critical edges. Note that
   /// the method getCriticalEdges can be used to retrieve the critical edges.
   ///
+  /// An edge is also considered as "critical" if it has a single precedessor
+  /// but the predecessor's terminal instruction is a user of the value.
+  ///
   /// If \p deBlocks is provided, all dead-end blocks are ignored. This
   /// prevents unreachable-blocks to be included in the frontier.
   bool computeFrontier(Frontier &frontier, Mode mode,
@@ -132,6 +135,15 @@ class ValueLifetimeAnalysis {
   SILInstruction *findLastUserInBlock(SILBasicBlock *bb);
 };
 
+/// Destroys \p valueOrStackLoc at \p frontier.
+///
+/// If  \p valueOrStackLoc is an alloc_stack, inserts destroy_addr and
+/// dealloc_stack at each instruction of the \p frontier.
+/// Otherwise \p valueOrStackLoc must be a value type and in this case, inserts
+/// destroy_value at each instruction of the \p frontier.
+void endLifetimeAtFrontier(SILValue valueOrStackLoc,
+                           const ValueLifetimeAnalysis::Frontier &frontier,
+                           SILBuilderContext &builderCtxt);
 
 } // end namespace swift
 

lib/SILOptimizer/Analysis/ARCAnalysis.cpp

@@ -987,131 +987,6 @@ findMatchingReleases(SILBasicBlock *BB) {
   processMatchingReleases();
 }
 
-//===----------------------------------------------------------------------===//
-//                    Code for Determining Final Releases
-//===----------------------------------------------------------------------===//
-
-// Propagate liveness backwards from an initial set of blocks in our
-// LiveIn set.
-static void propagateLiveness(llvm::SmallPtrSetImpl<SILBasicBlock *> &LiveIn,
-                              SILBasicBlock *DefBB) {
-  // First populate a worklist of predecessors.
-  llvm::SmallVector<SILBasicBlock *, 64> Worklist;
-  for (auto *BB : LiveIn)
-    for (auto Pred : BB->getPredecessorBlocks())
-      Worklist.push_back(Pred);
-
-  // Now propagate liveness backwards until we hit the alloc_box.
-  while (!Worklist.empty()) {
-    auto *BB = Worklist.pop_back_val();
-
-    // If it's already in the set, then we've already queued and/or
-    // processed the predecessors.
-    if (BB == DefBB || !LiveIn.insert(BB).second)
-      continue;
-
-    for (auto Pred : BB->getPredecessorBlocks())
-      Worklist.push_back(Pred);
-  }
-}
-
-// Is any successor of BB in the LiveIn set?
-static bool successorHasLiveIn(SILBasicBlock *BB,
-                               llvm::SmallPtrSetImpl<SILBasicBlock *> &LiveIn) {
-  for (auto &Succ : BB->getSuccessors())
-    if (LiveIn.count(Succ))
-      return true;
-
-  return false;
-}
-
-// Walk backwards in BB looking for the last use of a given
-// value, and add it to the set of release points.
-static bool addLastUse(SILValue V, SILBasicBlock *BB,
-                       ReleaseTracker &Tracker) {
-  for (auto I = BB->rbegin(); I != BB->rend(); ++I) {
-    if (Tracker.isUser(&*I)) {
-      Tracker.trackLastRelease(&*I);
-      return true;
-    }
-  }
-
-  llvm_unreachable("BB is expected to have a use of a closure");
-  return false;
-}
-
-/// TODO: Refactor this code so the decision on whether or not to accept an
-/// instruction.
-bool swift::getFinalReleasesForValue(SILValue V, ReleaseTracker &Tracker) {
-  llvm::SmallPtrSet<SILBasicBlock *, 16> LiveIn;
-  llvm::SmallPtrSet<SILBasicBlock *, 16> UseBlocks;
-
-  // First attempt to get the BB where this value resides.
-  auto *DefBB = V->getParentBlock();
-  if (!DefBB)
-    return false;
-
-  bool seenRelease = false;
-  SILInstruction *OneRelease = nullptr;
-
-  // We'll treat this like a liveness problem where the value is the def. Each
-  // block that has a use of the value has the value live-in unless it is the
-  // block with the value.
-  SmallVector<Operand *, 8> Uses(V->getUses());
-  while (!Uses.empty()) {
-    auto *Use = Uses.pop_back_val();
-    auto *User = Use->getUser();
-    auto *BB = User->getParent();
-
-    if (Tracker.isUserTransitive(User)) {
-      Tracker.trackUser(User);
-      auto *CastInst = cast<SingleValueInstruction>(User);
-      Uses.append(CastInst->getUses().begin(), CastInst->getUses().end());
-      continue;
-    }
-
-    if (!Tracker.isUserAcceptable(User))
-      return false;
-
-    Tracker.trackUser(User);
-
-    if (BB != DefBB)
-      LiveIn.insert(BB);
-
-    // Also keep track of the blocks with uses.
-    UseBlocks.insert(BB);
-
-    // Try to speed up the trivial case of single release/dealloc.
-    if (isa<StrongReleaseInst>(User) || isa<DeallocBoxInst>(User) ||
-        isa<DestroyValueInst>(User) || isa<ReleaseValueInst>(User)) {
-      if (!seenRelease)
-        OneRelease = User;
-      else
-        OneRelease = nullptr;
-
-      seenRelease = true;
-    }
-  }
-
-  // Only a single release/dealloc? We're done!
-  if (OneRelease) {
-    Tracker.trackLastRelease(OneRelease);
-    return true;
-  }
-
-  propagateLiveness(LiveIn, DefBB);
-
-  // Now examine each block we saw a use in. If it has no successors
-  // that are in LiveIn, then the last use in the block is the final
-  // release/dealloc.
-  for (auto *BB : UseBlocks)
-    if (!successorHasLiveIn(BB, LiveIn))
-      if (!addLastUse(V, BB, Tracker))
-        return false;
-
-  return true;
-}
-
 //===----------------------------------------------------------------------===//
 //                            Leaking BB Analysis
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/IPO/ClosureSpecializer.cpp

@@ -971,6 +971,8 @@ void SILClosureSpecializerTransform::run() {
     // specialized all of their uses.
     LLVM_DEBUG(llvm::dbgs() << "Trying to remove dead closures!\n");
     sortUnique(PropagatedClosures);
+    bool needUpdateStackNesting = false;
+
     for (auto *Closure : PropagatedClosures) {
       LLVM_DEBUG(llvm::dbgs() << "    Visiting: " << *Closure);
       if (!tryDeleteDeadClosure(Closure)) {
@@ -981,6 +983,11 @@ void SILClosureSpecializerTransform::run() {
 
       LLVM_DEBUG(llvm::dbgs() << "        Deleted closure!\n");
       ++NumPropagatedClosuresEliminated;
+      needUpdateStackNesting = true;
+    }
+
+    if (needUpdateStackNesting) {
+      StackNesting().correctStackNesting(F);
     }
   }
 

lib/SILOptimizer/Mandatory/MandatoryCombine.cpp

@@ -32,6 +32,7 @@
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/raw_ostream.h"
@@ -69,6 +70,10 @@ class MandatoryCombiner final
   /// Whether any changes have been made.
   bool madeChange;
 
+  /// Set to true if some alloc/dealloc_stack instruction are inserted and at
+  /// the end of the run stack nesting needs to be corrected.
+  bool needUpdateStackNesting;
+
   /// The number of times that the worklist has been processed.
   unsigned iteration;
 
@@ -114,6 +119,10 @@ class MandatoryCombiner final
       ++iteration;
     }
 
+    if (needUpdateStackNesting) {
+      StackNesting().correctStackNesting(&function);
+    }
+
     return changed;
   }
 
@@ -269,7 +278,9 @@ SILInstruction *MandatoryCombiner::visitApplyInst(ApplyInst *instruction) {
                                              /*instructionDescription=*/""
 #endif
   );
-  tryDeleteDeadClosure(partialApply, instModCallbacks);
+  if (tryDeleteDeadClosure(partialApply, instModCallbacks)) {
+    needUpdateStackNesting = true;
+  }
   return nullptr;
 }
 

lib/SILOptimizer/Mandatory/MandatoryInlining.cpp

@@ -342,7 +342,8 @@ static SILValue cleanupLoadedCalleeValue(SILValue calleeValue, LoadInst *li) {
 
 /// Removes instructions that create the callee value if they are no
 /// longer necessary after inlining.
-static void cleanupCalleeValue(SILValue calleeValue) {
+static void cleanupCalleeValue(SILValue calleeValue,
+                               bool &needUpdateStackNesting) {
   // Handle the case where the callee of the apply is a load instruction. If we
   // fail to optimize, return. Otherwise, see if we can look through other
   // abstractions on our callee.
@@ -382,6 +383,7 @@ static void cleanupCalleeValue(SILValue calleeValue) {
       return;
     calleeValue = callee;
   }
+  needUpdateStackNesting = true;
 
   calleeValue = stripCopiesAndBorrows(calleeValue);
 
@@ -446,6 +448,10 @@ class ClosureCleanup {
   SmallBlotSetVector<SILInstruction *, 4> deadFunctionVals;
 
 public:
+  /// Set to true if some alloc/dealloc_stack instruction are inserted and at
+  /// the end of the run stack nesting needs to be corrected.
+  bool needUpdateStackNesting = false;
+
   /// This regular instruction deletion callback checks for any function-type
   /// values that may be unused after deleting the given instruction.
   void recordDeadFunction(SILInstruction *deletedInst) {
@@ -475,7 +481,7 @@ class ClosureCleanup {
         continue;
 
       if (auto *SVI = dyn_cast<SingleValueInstruction>(I.getValue()))
-        cleanupCalleeValue(SVI);
+        cleanupCalleeValue(SVI, needUpdateStackNesting);
     }
   }
 };
@@ -938,6 +944,7 @@ runOnFunctionRecursively(SILOptFunctionBuilder &FuncBuilder,
       // we may be able to remove dead callee computations (e.g. dead
       // partial_apply closures).
       closureCleanup.cleanupDeadClosures(F);
+      needUpdateStackNesting |= closureCleanup.needUpdateStackNesting;
 
       // Resume inlining within nextBB, which contains only the inlined
       // instructions and possibly instructions in the original call block that

lib/SILOptimizer/SILCombiner/SILCombine.cpp

@@ -30,6 +30,7 @@
 #include "swift/SILOptimizer/Utils/CanonicalizeInstruction.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
@@ -229,6 +230,10 @@ bool SILCombiner::runOnFunction(SILFunction &F) {
     Iteration++;
   }
 
+  if (needUpdateStackNesting) {
+    StackNesting().correctStackNesting(&F);
+  }
+
   // Cleanup the builder and return whether or not we made any changes.
   return Changed;
 }

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -65,6 +65,10 @@ class SILCombiner :
   /// If set to true then the optimizer is free to erase cond_fail instructions.
   bool RemoveCondFails;
 
+  /// Set to true if some alloc/dealloc_stack instruction are inserted and at
+  /// the end of the run stack nesting needs to be corrected.
+  bool needUpdateStackNesting = false;
+
   /// The current iteration of the SILCombine.
   unsigned Iteration;