Merge pull request #67031 from jckarter/allocbox-to-stack-remove-unused-original-bodies

AllocBoxToStack: Remove bodies of closure functions left unused after specialization.

Author: jckarter

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -1213,9 +1213,51 @@ static void rewriteApplySites(AllocBoxToStackState &pass) {
     auto *FRI = cast<FunctionRefInst>(Apply.getCallee());
     Apply.getInstruction()->eraseFromParent();
 
-    // TODO: Erase from module if there are no more uses.
-    if (FRI->use_empty())
+    if (FRI->use_empty()) {
+      auto referencedFn = FRI->getReferencedFunction();
       FRI->eraseFromParent();
+
+      // TODO: Erase from module if there are no more uses.
+      // If the function has no remaining references, it should eventually
+      // be deleted. We can't do that from a function pass, since the function
+      // is still queued up for other passes to run after this one, but we
+      // can at least gut the implementation, since subsequent passes that
+      // rely on stack promotion to occur (particularly closure lifetime
+      // fixup and move-only checking) may not be able to proceed in a
+      // sensible way on the now non-canonical original implementation.
+      if (referencedFn->getRefCount() == 0
+          && !isPossiblyUsedExternally(referencedFn->getLinkage(),
+                                   referencedFn->getModule().isWholeModule())) {
+        LLVM_DEBUG(llvm::dbgs() << "*** Deleting original function " << referencedFn->getName() << "'s body since it is unused");
+        // Remove all non-entry blocks.
+        auto entryBB = referencedFn->begin();
+        auto nextBB = std::next(entryBB);
+        
+        while (nextBB != referencedFn->end()) {
+          auto thisBB = nextBB;
+          ++nextBB;
+          thisBB->eraseFromParent();
+        }
+        
+        // Rewrite the entry block to only contain an unreachable.
+        auto loc = entryBB->begin()->getLoc();
+        entryBB->eraseAllInstructions(referencedFn->getModule());
+        {
+          SILBuilder b(&*entryBB);
+          b.createUnreachable(loc);
+        }
+        
+        // Refresh the CFG in case we removed any function calls.
+        pass.CFGChanged = true;
+        
+        // If the function has shared linkage, reduce this version to private
+        // linkage, because we don't want the deleted-body form to win in any
+        // ODR shootouts.
+        if (referencedFn->getLinkage() == SILLinkage::Shared) {
+          referencedFn->setLinkage(SILLinkage::Private);
+        }
+      }
+    }
   }
 }
 

test/SILOptimizer/allocbox_to_stack.sil

@@ -651,9 +651,13 @@ bb0(%0 : $Int, %1 : $*S<T>):
   return %9 : $Bool
 }
 
-// CHECK-LABEL: sil shared [_semantics "sil.optimizer.moveonly.diagnostic.ignore"] @closure1
-sil shared @closure1 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
+// This closure body gets specialized with unboxed capture parameters, so
+// the original function body is stubbed out once the call sites are all
+// specialized.
+// CHECK-LABEL: sil private [_semantics "sil.optimizer.moveonly.diagnostic.ignore"] @closure1
 // CHECK: bb0
+// CHECK-NEXT: unreachable
+sil shared @closure1 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
 bb0(%0 : $Int, %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):
   %2 = project_box %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>, 0
   %3 = function_ref @inner : $@convention(thin) (@owned @callee_owned () -> Bool) -> Bool
@@ -664,7 +668,6 @@ bb0(%0 : $Int, %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):
   %7 = partial_apply %4<T>(%0, %5) : $@convention(thin) <τ_0_0 where τ_0_0 : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <τ_0_0>) -> Bool
   %8 = apply %3(%7) : $@convention(thin) (@owned @callee_owned () -> Bool) -> Bool
   strong_release %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>
-// CHECK: return
   return %8 : $Bool
 }
 

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -751,9 +751,13 @@ bb0(%0 : $Int, %1 : $*S<T>):
   return %9 : $Bool
 }
 
-// CHECK-LABEL: sil shared [_semantics "sil.optimizer.moveonly.diagnostic.ignore"] [ossa] @closure1
-sil shared [ossa] @closure1 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
+// This closure body gets specialized with unboxed capture parameters, so
+// the original function body is stubbed out once the call sites are all
+// specialized.
+// CHECK-LABEL: sil private [_semantics "sil.optimizer.moveonly.diagnostic.ignore"] [ossa] @closure1
 // CHECK: bb0
+// CHECK: unreachable
+sil shared [ossa] @closure1 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
 bb0(%0 : $Int, %1 : @owned $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):
   %2 = project_box %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>, 0
   %3 = function_ref @inner : $@convention(thin) (@owned @callee_owned () -> Bool) -> Bool
@@ -764,7 +768,6 @@ bb0(%0 : $Int, %1 : @owned $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>)
   %7 = partial_apply %4<T>(%0, %5) : $@convention(thin) <τ_0_0 where τ_0_0 : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <τ_0_0>) -> Bool
   %8 = apply %3(%7) : $@convention(thin) (@owned @callee_owned () -> Bool) -> Bool
   destroy_value %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>
-// CHECK: return
   return %8 : $Bool
 }