Fix capture promotion for partial_apply [stack]

Author: aschwaighofer

lib/SILOptimizer/IPO/CapturePromotion.cpp

@@ -1034,6 +1034,17 @@ static bool scanUsesForEscapesAndMutations(Operand *Op,
     return isPartialApplyNonEscapingUser(Op, PAI, State);
   }
 
+  // A mark_dependence user on a partial_apply is safe.
+  if (auto *MD = dyn_cast<MarkDependenceInst>(User)) {
+    if (MD->getBase() == Op->get()) {
+      auto parent = MD->getValue();
+      while ((MD = dyn_cast<MarkDependenceInst>(parent))) {
+        parent = MD->getValue();
+      }
+      return isa<PartialApplyInst>(parent);
+    }
+  }
+
   if (auto *PBI = dyn_cast<ProjectBoxInst>(User)) {
     // It is assumed in later code that we will only have 1 project_box. This
     // can be seen since there is no code for reasoning about multiple
@@ -1198,6 +1209,26 @@ static SILValue getOrCreateProjectBoxHelper(SILValue PartialOperand) {
   return B.createProjectBox(Box->getLoc(), Box, 0);
 }
 
+/// Change the base in mark_dependence.
+static void
+mapMarkDependenceArguments(SingleValueInstruction *root,
+                           llvm::DenseMap<SILValue, SILValue> &map) {
+  for (auto *Use : root->getUses()) {
+    if (auto *MD = dyn_cast<MarkDependenceInst>(Use->getUser())) {
+      mapMarkDependenceArguments(MD, map);
+      auto iter = map.find(MD->getBase());
+      if (iter != map.end()) {
+        MD->setBase(iter->second);
+      }
+      // Remove mark_dependence on trivial values.
+      if (MD->getBase()->getType().isTrivial(MD->getModule())) {
+        MD->replaceAllUsesWith(MD->getValue());
+        MD->eraseFromParent();
+      }
+    }
+  }
+}
+
 /// Given a partial_apply instruction and a set of promotable indices,
 /// clone the closure with the promoted captures and replace the partial_apply
 /// with a partial_apply of the new closure, fixing up reference counting as
@@ -1236,6 +1267,8 @@ processPartialApplyInst(SILOptFunctionBuilder &FuncBuilder,
   unsigned OpCount = PAI->getNumOperands() - PAI->getNumTypeDependentOperands();
   SmallVector<SILValue, 16> Args;
   auto NumIndirectResults = calleeConv.getNumIndirectSILResults();
+  llvm::DenseMap<SILValue, SILValue> capturedMap;
+  llvm::SmallSet<SILValue, 16> newCaptures;
   for (; OpNo != OpCount; ++OpNo) {
     unsigned Index = OpNo - 1 + FirstIndex;
     if (!PromotableIndices.count(Index)) {
@@ -1250,8 +1283,18 @@ processPartialApplyInst(SILOptFunctionBuilder &FuncBuilder,
     SILValue Addr = getOrCreateProjectBoxHelper(Box);
 
     auto &typeLowering = M.getTypeLowering(Addr->getType());
-    Args.push_back(
-        typeLowering.emitLoadOfCopy(B, PAI->getLoc(), Addr, IsNotTake));
+    auto newCaptured =
+        typeLowering.emitLoadOfCopy(B, PAI->getLoc(), Addr, IsNotTake);
+    Args.push_back(newCaptured);
+
+    capturedMap[Box] = newCaptured;
+    newCaptures.insert(newCaptured);
+
+    // A partial_apply [stack] does not own the captured argument but we must
+    // destroy the projected object. We will do so after having created the new
+    // partial_apply below.
+    if (PAI->isOnStack())
+      continue;
 
     // Cleanup the captured argument.
     //
@@ -1268,14 +1311,30 @@ processPartialApplyInst(SILOptFunctionBuilder &FuncBuilder,
   // Create a new partial apply with the new arguments.
   auto *NewPAI = B.createPartialApply(
       PAI->getLoc(), FnVal, PAI->getSubstitutionMap(), Args,
-      PAI->getType().getAs<SILFunctionType>()->getCalleeConvention());
+      PAI->getType().getAs<SILFunctionType>()->getCalleeConvention(),
+      PAI->isOnStack());
   PAI->replaceAllUsesWith(NewPAI);
   PAI->eraseFromParent();
   if (FRI->use_empty()) {
     FRI->eraseFromParent();
     // TODO: If this is the last use of the closure, and if it has internal
     // linkage, we should remove it from the SILModule now.
   }
+
+  if (NewPAI->isOnStack()) {
+    // Insert destroy's of new captured arguments.
+    for (auto *Use : NewPAI->getUses()) {
+      if (auto *DS = dyn_cast<DeallocStackInst>(Use->getUser())) {
+        B.setInsertionPoint(std::next(SILBasicBlock::iterator(DS)));
+        insertDestroyOfCapturedArguments(NewPAI, B, [&](SILValue arg) -> bool {
+          return newCaptures.count(arg);
+        });
+      }
+    }
+    // Map the mark dependence arguments.
+    mapMarkDependenceArguments(NewPAI, capturedMap);
+  }
+
   return ClonedFn;
 }
 

test/SILOptimizer/capture_promotion.sil

@@ -407,3 +407,70 @@ bb0:
 
   return %21 : $@callee_guaranteed () -> Int
 }
+
+// CHECK-LABEL: sil @test_capture_promotion_on_stack
+sil @test_capture_promotion_on_stack : $@convention(thin) () -> () {
+bb0:
+  %0 = tuple ()
+  %1 = alloc_box $<τ_0_0> { var τ_0_0 } <Foo>
+  %1a = project_box %1 : $<τ_0_0> { var τ_0_0 } <Foo>, 0
+  %2 = function_ref @foo_allocating_init : $@convention(thin) (@thick Foo.Type) -> @owned Foo
+  %3 = metatype $@thick Foo.Type
+  %4 = apply %2(%3) : $@convention(thin) (@thick Foo.Type) -> @owned Foo
+  store %4 to %1a : $*Foo
+  %6 = alloc_box $<τ_0_0> { var τ_0_0 } <Baz>
+  %6a = project_box %6 : $<τ_0_0> { var τ_0_0 } <Baz>, 0
+  %7 = function_ref @baz_init : $@convention(thin) (@thin Baz.Type) -> @owned Baz
+  %8 = metatype $@thin Baz.Type
+  %9 = apply %7(%8) : $@convention(thin) (@thin Baz.Type) -> @owned Baz
+  store %9 to %6a : $*Baz
+  %11 = alloc_box $<τ_0_0> { var τ_0_0 } <Int>
+  %11a = project_box %11 : $<τ_0_0> { var τ_0_0 } <Int>, 0
+  %12 = function_ref @convert_from_integer_literal : $@convention(thin) (Builtin.Word, @thin Int.Type) -> Int
+  %13 = metatype $@thin Int.Type
+  %14 = integer_literal $Builtin.Word, 3
+  %15 = apply %12(%14, %13) : $@convention(thin) (Builtin.Word, @thin Int.Type) -> Int
+  store %15 to %11a : $*Int
+
+// CHECK-NOT: function_ref @closure0_guaranteed
+// CHECK: [[CLOSURE_PROMOTE:%.*]] = function_ref @$s19closure0_guaranteedTf2iii_n
+// CHECK-NOT: function_ref @closure0_guaranteed
+
+// The Foo variable is loaded from and retained, because it is a reference type
+// CHECK-NEXT: [[LOADFOO:%.*]] = load {{.*}} : $*Foo
+// CHECK-NEXT: strong_retain [[LOADFOO]] : $Foo
+
+// The Baz variable is loaded and retain_value'd, because it is a non-trivial
+// aggregate type
+// CHECK-NEXT: [[LOADBAZ:%.*]] = load {{.*}} : $*Baz
+// CHECK-NEXT: retain_value [[LOADBAZ:%.*]] : $Baz
+
+// The Int variable is loaded only, because it is trivial
+// CHECK-NEXT: [[LOADINT:%.*]] = load {{.*}} : $*Int
+
+// The partial apply has one value argument for each pair of arguments that was
+// previously used to capture and pass the variable by reference
+// CHECK-NEXT: [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[CLOSURE_PROMOTE]]([[LOADFOO]], [[LOADBAZ]], [[LOADINT]])
+// CHECK-NEXT: [[MD1:%.*]] = mark_dependence [[PA]] : $@noescape @callee_guaranteed () -> Int on [[LOADFOO]] : $Foo
+// CHECK-NEXT: [[MD2:%.*]] = mark_dependence [[MD1]] : $@noescape @callee_guaranteed () -> Int on [[LOADBAZ]] : $Baz
+// CHECK-NEXT: apply [[MD2]]() : $@noescape @callee_guaranteed () -> Int
+// CHECK-NEXT:  dealloc_stack [[PA]] : $@noescape @callee_guaranteed () -> Int
+// CHECK-NEXT:  strong_release [[LOADFOO]] : $Foo
+// CHECK-NEXT:  release_value [[LOADBAZ]] : $Baz
+// CHECK-NEXT:  strong_release {{.*}} : $<τ_0_0> { var τ_0_0 } <Int>
+// CHECK-NEXT:  strong_release {{.*}} : $<τ_0_0> { var τ_0_0 } <Baz>
+// CHECK-NEXT:  strong_release {{.*}} : $<τ_0_0> { var τ_0_0 } <Foo>
+
+  %17 = function_ref @closure0_guaranteed : $@convention(thin) (@guaranteed <τ_0_0> { var τ_0_0 } <Foo>, @guaranteed <τ_0_0> { var τ_0_0 } <Baz>, @guaranteed <τ_0_0> { var τ_0_0 } <Int>) -> Int
+  %21 = partial_apply [callee_guaranteed] [on_stack] %17(%1, %6, %11) : $@convention(thin) (@guaranteed <τ_0_0> { var τ_0_0 } <Foo>, @guaranteed <τ_0_0> { var τ_0_0 } <Baz>, @guaranteed <τ_0_0> { var τ_0_0 } <Int>) -> Int
+  %md = mark_dependence %21 :$@noescape @callee_guaranteed () -> Int on %1 : $<τ_0_0> { var τ_0_0 } <Foo>
+  %md2 = mark_dependence %md :$@noescape @callee_guaranteed () -> Int on %6 : $<τ_0_0> { var τ_0_0 } <Baz>
+  %md3 = mark_dependence %md2 :$@noescape @callee_guaranteed () -> Int on %11 : $<τ_0_0> { var τ_0_0 } <Int>
+  apply %md3() : $@noescape @callee_guaranteed () -> Int
+  dealloc_stack %21 : $@noescape @callee_guaranteed () -> Int
+  strong_release %11 : $<τ_0_0> { var τ_0_0 } <Int>
+  strong_release %6 : $<τ_0_0> { var τ_0_0 } <Baz>
+  strong_release %1 : $<τ_0_0> { var τ_0_0 } <Foo>
+  %tuple = tuple()
+  return %tuple : $()
+}