[ClosureLifetimeFixup] Dealloc args on frontier.

Previously, the dealloc_stacks created for the alloc_stacks used to pass
@in_guaranteed arguments to on_stack closures were created after the
users of the closure.  When SILGen created these alloc_stacks in the
same block as the users, this happened to work.  Now that
AddressLowering creates such alloc_stacks elsewhere, this approach
results in invalid SIL.

Here, the dealloc_stacks are instead at the end of each block in the
dominance frontier of the alloc_stack.

Author: nate-chandler

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -271,18 +271,14 @@ void releasePartialApplyCapturedArg(
     SILParameterInfo paramInfo,
     InstModCallbacks callbacks = InstModCallbacks());
 
-void deallocPartialApplyCapturedArg(
-    SILBuilder &builder, SILLocation loc, SILValue arg,
-    SILParameterInfo paramInfo);
-
 /// Insert destroys of captured arguments of partial_apply [stack].
 void insertDestroyOfCapturedArguments(
     PartialApplyInst *pai, SILBuilder &builder,
     llvm::function_ref<bool(SILValue)> shouldInsertDestroy =
         [](SILValue arg) -> bool { return true; });
 
-void insertDeallocOfCapturedArguments(
-    PartialApplyInst *pai, SILBuilder &builder);
+void insertDeallocOfCapturedArguments(PartialApplyInst *pai,
+                                      DominanceInfo *domInfo);
 
 /// This iterator 'looks through' one level of builtin expect users exposing all
 /// users of the looked through builtin expect instruction i.e it presents a

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -444,6 +444,12 @@ static SILValue skipConvert(SILValue v) {
   return pa;
 }
 
+static SILAnalysis::InvalidationKind
+analysisInvalidationKind(const bool &modifiedCFG) {
+  return modifiedCFG ? SILAnalysis::InvalidationKind::FunctionBody
+                     : SILAnalysis::InvalidationKind::CallsAndInstructions;
+}
+
 /// Rewrite a partial_apply convert_escape_to_noescape sequence with a single
 /// apply/try_apply user to a partial_apply [stack] terminated with a
 /// dealloc_stack placed after the apply.
@@ -468,9 +474,10 @@ static SILValue skipConvert(SILValue v) {
 /// caller needs to use the StackNesting utility to update the dealloc_stack
 /// nesting.
 static SILValue tryRewriteToPartialApplyStack(
-    ConvertEscapeToNoEscapeInst *cvt,
-    SILInstruction *closureUser, InstructionDeleter &deleter,
-    llvm::DenseMap<SILInstruction *, SILInstruction *> &memoized) {
+    ConvertEscapeToNoEscapeInst *cvt, SILInstruction *closureUser,
+    DominanceAnalysis *dominanceAnalysis, InstructionDeleter &deleter,
+    llvm::DenseMap<SILInstruction *, SILInstruction *> &memoized,
+    const bool &modifiedCFG) {
 
   auto *origPA = dyn_cast<PartialApplyInst>(skipConvert(cvt->getOperand()));
   if (!origPA)
@@ -567,16 +574,24 @@ static SILValue tryRewriteToPartialApplyStack(
     auto loc = RegularLocation(builder.getInsertionPointLoc());
     builder.createDeallocStack(loc, newPA);
     insertDestroyOfCapturedArguments(newPA, builder);
-    // dealloc_stack of the in_guaranteed capture is inserted
-    insertDeallocOfCapturedArguments(newPA, builder);
   });
+  // The CFG may have been modified during this run.  If it was, the dominance
+  // analysis would no longer be valid.  Invalidate it now if necessary,
+  // according to the kinds of changes that may have been made.  Note that if
+  // the CFG hasn't been modified, this is a noop thanks to
+  // DominanceAnalysis::shouldInvalidate's definition.
+  dominanceAnalysis->invalidate(closureUser->getFunction(),
+                                analysisInvalidationKind(modifiedCFG));
+  // Insert dealloc_stacks of any in_guaranteed captures.
+  insertDeallocOfCapturedArguments(
+      newPA, dominanceAnalysis->get(closureUser->getFunction()));
   return closure;
 }
 
 static bool tryExtendLifetimeToLastUse(
-    ConvertEscapeToNoEscapeInst *cvt,
+    ConvertEscapeToNoEscapeInst *cvt, DominanceAnalysis *dominanceAnalysis,
     llvm::DenseMap<SILInstruction *, SILInstruction *> &memoized,
-    InstructionDeleter &deleter) {
+    InstructionDeleter &deleter, const bool &modifiedCFG) {
   // If there is a single user that is an apply this is simple: extend the
   // lifetime of the operand until after the apply.
   auto *singleUser = lookThroughRebastractionUsers(cvt, memoized);
@@ -597,8 +612,9 @@ static bool tryExtendLifetimeToLastUse(
     return false;
   }
 
-  if (SILValue closure = tryRewriteToPartialApplyStack(cvt, singleUser,
-                                                       deleter, memoized)) {
+  if (SILValue closure = tryRewriteToPartialApplyStack(
+          cvt, singleUser, dominanceAnalysis, deleter, memoized,
+          /*const*/ modifiedCFG)) {
     if (auto *cfi = dyn_cast<ConvertFunctionInst>(closure))
       closure = cfi->getOperand();
     if (endAsyncLet && isa<MarkDependenceInst>(closure)) {
@@ -993,8 +1009,9 @@ static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *cb,
   return true;
 }
 
-static bool fixupClosureLifetimes(SILFunction &fn, bool &checkStackNesting,
-                                  bool &modifiedCFG) {
+static bool fixupClosureLifetimes(SILFunction &fn,
+                                  DominanceAnalysis *dominanceAnalysis,
+                                  bool &checkStackNesting, bool &modifiedCFG) {
   bool changed = false;
 
   // tryExtendLifetimeToLastUse uses a cache of recursive instruction use
@@ -1027,7 +1044,9 @@ static bool fixupClosureLifetimes(SILFunction &fn, bool &checkStackNesting,
         }
       }
 
-      if (tryExtendLifetimeToLastUse(cvt, memoizedQueries, updater.getDeleter())) {
+      if (tryExtendLifetimeToLastUse(cvt, dominanceAnalysis, memoizedQueries,
+                                     updater.getDeleter(),
+                                     /*const*/ modifiedCFG)) {
         changed = true;
         checkStackNesting = true;
         continue;
@@ -1064,15 +1083,15 @@ class ClosureLifetimeFixup : public SILFunctionTransform {
     bool checkStackNesting = false;
     bool modifiedCFG = false;
 
-    if (fixupClosureLifetimes(*getFunction(), checkStackNesting, modifiedCFG)) {
+    auto *dominanceAnalysis = PM->getAnalysis<DominanceAnalysis>();
+
+    if (fixupClosureLifetimes(*getFunction(), dominanceAnalysis,
+                              checkStackNesting, modifiedCFG)) {
       if (checkStackNesting){
         modifiedCFG |=
           StackNesting::fixNesting(getFunction()) == StackNesting::Changes::CFG;
       }
-      if (modifiedCFG)
-        invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);
-      else
-        invalidateAnalysis(SILAnalysis::InvalidationKind::CallsAndInstructions);
+      invalidateAnalysis(analysisInvalidationKind(modifiedCFG));
     }
     LLVM_DEBUG(getFunction()->verify());
 

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -910,15 +910,6 @@ void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
   emitDestroyOperation(builder, loc, arg, callbacks);
 }
 
-void swift::deallocPartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
-                                           SILValue arg,
-                                           SILParameterInfo paramInfo) {
-  if (!paramInfo.isIndirectInGuaranteed())
-    return;
-
-  builder.createDeallocStack(loc, arg);
-}
-
 static bool
 deadMarkDependenceUser(SILInstruction *inst,
                        SmallVectorImpl<SILInstruction *> &deleteInsts) {
@@ -1473,19 +1464,32 @@ void swift::insertDestroyOfCapturedArguments(
   }
 }
 
-void swift::insertDeallocOfCapturedArguments(
-    PartialApplyInst *pai, SILBuilder &builder) {
+void swift::insertDeallocOfCapturedArguments(PartialApplyInst *pai,
+                                             DominanceInfo *domInfo) {
   assert(pai->isOnStack());
 
   ApplySite site(pai);
   SILFunctionConventions calleeConv(site.getSubstCalleeType(),
                                     pai->getModule());
-  auto loc = CleanupLocation(RegularLocation::getAutoGeneratedLocation());
   for (auto &arg : pai->getArgumentOperands()) {
     unsigned calleeArgumentIndex = site.getCalleeArgIndex(arg);
     assert(calleeArgumentIndex >= calleeConv.getSILArgIndexOfFirstParam());
     auto paramInfo = calleeConv.getParamInfoForSILArg(calleeArgumentIndex);
-    deallocPartialApplyCapturedArg(builder, loc, arg.get(), paramInfo);
+    if (!paramInfo.isIndirectInGuaranteed())
+      continue;
+
+    SmallVector<SILBasicBlock *, 4> boundary;
+    auto *asi = cast<AllocStackInst>(arg.get());
+    computeDominatedBoundaryBlocks(asi->getParent(), domInfo, boundary);
+
+    for (auto *block : boundary) {
+      auto *terminator = block->getTerminator();
+      auto builder = SILBuilder(terminator);
+      if (isa<UnreachableInst>(terminator))
+        continue;
+      builder.createDeallocStack(CleanupLocation(terminator->getLoc()),
+                                 arg.get());
+    }
   }
 }
 

test/SILOptimizer/closure-lifetime-fixup-debuginfo.sil

@@ -18,10 +18,10 @@ associatedtype Index
 // CHECK:         [[CAPTURED_ADDR:%[^,]+]] = alloc_stack $Self
 // CHECK:         try_apply {{.*}}, normal [[SUCCESS:bb[0-9]+]]
 // CHECK:       [[SUCCESS]]
-// CHECK:         destroy_addr [[CAPTURED_ADDR]]{{.*}}  // id: {{%[^,]+}}; <invalid loc>:cleanup:auto_gen
-//                                                                                       ^^^^^^^
-// CHECK:         dealloc_stack [[CAPTURED_ADDR]]{{.*}} // id: {{%[^,]+}}; <invalid loc>:cleanup:auto_gen
-//                                                                                       ^^^^^^^
+// CHECK:         destroy_addr [[CAPTURED_ADDR]]{{.*}}:cleanup:
+//                                                     ^^^^^^^
+// CHECK:         dealloc_stack [[CAPTURED_ADDR]]{{.*}}:cleanup:
+//                                                      ^^^^^^^
 // CHECK-LABEL: } // end sil function 'destroy_and_dealloc_of_capture_are_cleanups'
 sil [ossa] @destroy_and_dealloc_of_capture_are_cleanups : $@convention(method) <Self where Self : Indexable> (@in Self) -> () {
 bb0(%self : $*Self):

test/SILOptimizer/closure-lifetime-fixup.sil

@@ -256,3 +256,47 @@ bb3:
   return %25 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @test_alloc_stack_arg_with_frontier_outside_try_apply_successors : $@convention(thin) <Self> (@in_guaranteed Self) -> @error any Error {
+// CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] :
+// CHECK:         [[STACK:%[^,]+]] = alloc_stack $Self
+// CHECK:         try_apply undef() {{.*}}, normal [[SUCCESS_1:bb[0-9]+]], error [[FAILURE_1:bb[0-9]+]]
+// CHECK:       [[SUCCESS_1]]
+// CHECK:         copy_addr [[INSTANCE]] to [init] [[STACK]] : $*Self
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] undef<Self>([[STACK]])
+// CHECK:         [[DEPENDENCY:%[^,]+]] = mark_dependence [[CLOSURE]] {{.*}} on [[STACK]]
+// CHECK:         try_apply undef([[DEPENDENCY]]) {{.*}}, normal [[SUCCESS_2:bb[0-9]+]], error [[FAILURE_2:bb[0-9]+]]
+// CHECK:       [[SUCCESS_2]]
+// CHECK:         dealloc_stack [[CLOSURE]]
+// CHECK:         destroy_addr [[STACK]]
+// CHECK:         dealloc_stack [[STACK]] : $*Self
+// CHECK:       [[FAILURE_1]]
+// CHECK:         dealloc_stack [[STACK]]
+// CHECK:         throw
+// CHECK:       [[FAILURE_2]]
+// CHECK-NOT:     dealloc_stack
+// CHECK:         unreachable
+// CHECK-LABEL: } // end sil function 'test_alloc_stack_arg_with_frontier_outside_try_apply_successors'
+sil [ossa] @test_alloc_stack_arg_with_frontier_outside_try_apply_successors : $@convention(thin) <Self> (@in_guaranteed Self) -> (@error any Error) {
+bb0(%2 : $*Self):
+  %4 = alloc_stack $Self
+  try_apply undef() : $@convention(thin) () -> (@error any Error), normal bb1, error bb3
+
+bb1(%10 : $()):
+  copy_addr %2 to [init] %4 : $*Self
+  %13 = partial_apply [callee_guaranteed] undef<Self>(%4) : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
+  %15 = convert_escape_to_noescape [not_guaranteed] %13 : $@callee_guaranteed () -> () to $@noescape @callee_guaranteed () -> ()
+  try_apply undef(%15) : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> (@error any Error), normal bb2, error bb4
+
+bb2(%18 : $()):
+  destroy_value %13 : $@callee_guaranteed () -> ()
+  %21 = tuple ()
+  dealloc_stack %4 : $*Self
+  return %21 : $()
+
+bb3(%25 : @owned $any Error):
+  dealloc_stack %4 : $*Self
+  throw %25 : $any Error
+
+bb4(%29 : @owned $any Error):
+  unreachable
+}