MandatoryInlining fixes for partial_apply [stack]

Author: aschwaighofer

lib/SILOptimizer/Mandatory/MandatoryInlining.cpp

@@ -373,6 +373,17 @@ static SILFunction *getCalleeFunction(
       return ThinToNoescapeCast->getOperand();
     }
 
+    // Ignore mark_dependence users. A partial_apply [stack] uses them to mark
+    // the dependence of the trivial closure context value on the captured
+    // arguments.
+    if (auto *MD = dyn_cast<MarkDependenceInst>(CalleeValue)) {
+      while (MD) {
+        CalleeValue = MD->getValue();
+        MD = dyn_cast<MarkDependenceInst>(CalleeValue);
+      }
+      return CalleeValue;
+    }
+
     auto *CFI = dyn_cast<ConvertEscapeToNoEscapeInst>(CalleeValue);
     if (!CFI)
       return CalleeValue;

lib/SILOptimizer/Utils/Local.cpp

@@ -1138,14 +1138,50 @@ static bool releaseCapturedArgsOfDeadPartialApply(PartialApplyInst *PAI,
   return true;
 }
 
+static bool
+deadMarkDependenceUser(SILInstruction *Inst,
+                       SmallVectorImpl<SILInstruction *> &DeleteInsts) {
+  if (!isa<MarkDependenceInst>(Inst))
+    return false;
+  DeleteInsts.push_back(Inst);
+  for (auto *Use : cast<SingleValueInstruction>(Inst)->getUses()) {
+    if (!deadMarkDependenceUser(Use->getUser(), DeleteInsts))
+      return false;
+  }
+  return true;
+}
+
 /// TODO: Generalize this to general objects.
 bool swift::tryDeleteDeadClosure(SingleValueInstruction *Closure,
                                  InstModCallbacks Callbacks) {
+  auto *PA = dyn_cast<PartialApplyInst>(Closure);
+
   // We currently only handle locally identified values that do not escape. We
   // also assume that the partial apply does not capture any addresses.
-  if (!isa<PartialApplyInst>(Closure) && !isa<ThinToThickFunctionInst>(Closure))
+  if (!PA && !isa<ThinToThickFunctionInst>(Closure))
     return false;
 
+  // A stack allocated partial apply does not have any release users. Delete it
+  // if the only users are the dealloc_stack and mark_dependence instructions.
+  if (PA && PA->isOnStack()) {
+    SmallVector<SILInstruction*, 8> DeleteInsts;
+    for (auto *Use : PA->getUses()) {
+      if (isa<DeallocStackInst>(Use->getUser()) ||
+          isa<DebugValueInst>(Use->getUser()))
+        DeleteInsts.push_back(Use->getUser());
+      else if (!deadMarkDependenceUser(Use->getUser(), DeleteInsts))
+        return false;
+    }
+    for (auto *Inst : reverse(DeleteInsts))
+      Callbacks.DeleteInst(Inst);
+    Callbacks.DeleteInst(PA);
+
+    // Note: the lifetime of the captured arguments is managed outside of the
+    // trivial closure value i.e: there will already be releases for the
+    // captured arguments. Releasing captured arguments is not necessary.
+    return true;
+  }
+
   // We only accept a user if it is an ARC object that can be removed if the
   // object is dead. This should be expanded in the future. This also ensures
   // that we are locally identified and non-escaping since we only allow for
@@ -1768,6 +1804,13 @@ swift::findLocalApplySites(FunctionRefBaseInst *FRI) {
            std::back_inserter(worklist));
       continue;
 
+    // A partial_apply [stack] marks its captured arguments with
+    // mark_dependence.
+    case SILInstructionKind::MarkDependenceInst:
+      copy(cast<SingleValueInstruction>(user)->getUses(),
+           std::back_inserter(worklist));
+      continue;
+
     // Look through any reference count instructions since these are not
     // escapes:
     case SILInstructionKind::CopyValueInst:
@@ -1778,6 +1821,8 @@ swift::findLocalApplySites(FunctionRefBaseInst *FRI) {
     case SILInstructionKind::RetainValueInst:
     case SILInstructionKind::ReleaseValueInst:
     case SILInstructionKind::DestroyValueInst:
+    // A partial_apply [stack] is deallocated with a dealloc_stack.
+    case SILInstructionKind::DeallocStackInst:
       continue;
     default:
       break;

test/SILOptimizer/mandatory_inlining.sil

@@ -1184,3 +1184,23 @@ bb0:
   %3 = apply %2() : $@noescape () -> Builtin.Int32
   return %3 : $Builtin.Int32
 }
+
+// CHECK-LABEL: sil @test_guaranteed_on_stack_closure
+// CHECK:  strong_retain %0 : $C
+// CHECK:  [[F:%.*]] = function_ref @use_c : $@convention(thin) (@guaranteed C) -> ()
+// CHECK:  apply [[F]](%0) : $@convention(thin) (@guaranteed C) -> ()
+// CHECK:  strong_release %0 : $C
+// CHECK-NOT:  strong_release
+// CHECK: return
+sil @test_guaranteed_on_stack_closure : $@convention(thin) (@guaranteed C) ->  () {
+bb0(%0: $C):
+  strong_retain %0 : $C
+  %closure_fun = function_ref @guaranteed_closure_func : $@convention(thin) (@guaranteed C) -> ()
+  %closure = partial_apply [on_stack] [callee_guaranteed] %closure_fun(%0) : $@convention(thin) (@guaranteed C) -> ()
+  %closure2 = mark_dependence %closure : $@noescape @callee_guaranteed () -> () on %0 : $C
+  apply %closure2() : $@noescape @callee_guaranteed () -> ()
+  strong_release %0: $C
+  dealloc_stack %closure : $@noescape @callee_guaranteed () -> ()
+  %tuple = tuple()
+  return %tuple : $()
+}