Merge pull request #62189 from nate-chandler/opaque-values/1/20221118

[AddressLowering] Shorten stack lifetimes.

Author: swift-ci

include/swift/SIL/SILBuilder.h

@@ -2805,7 +2805,7 @@ class SILBuilder {
   }
 };
 
-/// An wrapper on top of SILBuilder's constructor that automatically sets the
+/// A wrapper on top of SILBuilder's constructor that automatically sets the
 /// current SILDebugScope based on the specified insertion point. This is useful
 /// for situations where a single SIL instruction is lowered into a sequence of
 /// SIL instructions.

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -1002,16 +1002,28 @@ namespace {
 /// function in postorder. If the definition is an argument of this function,
 /// simply replace the function argument with an address representing the
 /// caller's storage.
-///
-/// TODO: shrink lifetimes by inserting alloc_stack at the dominance LCA and
-/// finding the lifetime boundary with a simple backward walk from uses.
 class OpaqueStorageAllocation {
   AddressLoweringState &pass;
+  /// The alloc_stacks that have been created which eventually need to have
+  /// corresponding dealloc_stacks created.
+  ///
+  /// Supports erasure because created alloc_stacks may be erased when block
+  /// arguments are coalesced.
+  SmallBlotSetVector<AllocStackInst *, 16> allocs;
+  /// The alloc_stacks that have been created which eventually need to be
+  /// positioned appropriately.
+  ///
+  /// The subset of allocs which aren't for opened existentials.
+  InstructionSet allocsToReposition;
 
 public:
-  explicit OpaqueStorageAllocation(AddressLoweringState &pass) : pass(pass) {}
+  explicit OpaqueStorageAllocation(AddressLoweringState &pass)
+      : pass(pass), allocsToReposition(pass.function) {}
 
   void allocateOpaqueStorage();
+  /// Position alloc_stacks according to uses and create dealloc_stacks that
+  /// jointly postdominate.
+  void finalizeOpaqueStorage();
 
 protected:
   void allocateValue(SILValue value);
@@ -1037,6 +1049,8 @@ class OpaqueStorageAllocation {
 
   AllocStackInst *createStackAllocation(SILValue value);
 
+  SILBasicBlock *getLeastCommonAncestorOfUses(SILValue value);
+
   void createStackAllocationStorage(SILValue value) {
     pass.valueStorageMap.getStorage(value).storageAddress =
         createStackAllocation(value);
@@ -1226,9 +1240,20 @@ void OpaqueStorageAllocation::removeAllocation(SILValue value) {
   for (Operand *use : uses) {
     pass.deleter.forceDelete(cast<DeallocStackInst>(use->getUser()));
   }
+  allocs.erase(allocInst);
   pass.deleter.forceDelete(allocInst);
 }
 
+SILBasicBlock *
+OpaqueStorageAllocation::getLeastCommonAncestorOfUses(SILValue value) {
+  SILBasicBlock *lca = nullptr;
+  for (auto *use : value->getUses()) {
+    auto *block = use->getParentBlock();
+    lca = lca ? pass.domInfo->findNearestCommonDominator(lca, block) : block;
+  }
+  return lca;
+}
+
 // Create alloc_stack that dominates an owned value \p value. Create
 // jointly-postdominating dealloc_stack instructions.  Nesting will be fixed
 // later.
@@ -1250,8 +1275,8 @@ AllocStackInst *OpaqueStorageAllocation::createStackAllocation(SILValue value) {
 
   // For opened existential types, allocate stack space at the type
   // definition. Allocating as early as possible provides more opportunity for
-  // creating use projections into value. But allocation must be no earlier then
-  // the latest type definition.
+  // creating use projections into value. But allocation must be no earlier
+  // then the latest type definition.
   SILInstruction *latestOpeningInst = nullptr;
   allocTy.getASTType().visit([&](CanType type) {
     auto archetype = dyn_cast<ArchetypeType>(type);
@@ -1274,31 +1299,44 @@ AllocStackInst *OpaqueStorageAllocation::createStackAllocation(SILValue value) {
       latestOpeningInst = openingInst;
     }
   });
-  auto allocPt = latestOpeningInst ? std::next(latestOpeningInst->getIterator())
-                                   : pass.function->begin()->begin();
+
+  auto allocPt = latestOpeningInst
+                     ? latestOpeningInst->getNextInstruction()->getIterator()
+                     : pass.function->getEntryBlock()->front().getIterator();
   auto allocBuilder = pass.getBuilder(allocPt);
   AllocStackInst *alloc = allocBuilder.createAllocStack(pass.genLoc(), allocTy);
+  allocs.insert(alloc);
+  if (latestOpeningInst == nullptr) {
+    allocsToReposition.insert(alloc);
+  }
+  return alloc;
+}
+
+void OpaqueStorageAllocation::finalizeOpaqueStorage() {
+  SmallVector<SILBasicBlock *, 4> boundary;
+  for (auto maybeAlloc : allocs) {
+    // An allocation may be erased when coalescing block arguments.
+    if (!maybeAlloc.hasValue())
+      continue;
+
+    auto *alloc = maybeAlloc.value();
+
+    if (allocsToReposition.contains(alloc)) {
+      auto allocPt = &*getLeastCommonAncestorOfUses(alloc)->begin();
+      alloc->moveBefore(allocPt);
+    }
 
-  auto dealloc = [&](SILBasicBlock::iterator insertPt) {
-    auto deallocBuilder = pass.getBuilder(insertPt);
-    deallocBuilder.createDeallocStack(pass.genLoc(), alloc);
-  };
-  if (latestOpeningInst) {
     // Deallocate at the predecessors of dominance frontier blocks that are
     // dominated by the alloc to ensure that allocation encloses not only the
     // uses of the current value, but also of any values reusing this storage as
     // a use projection.
-    SmallVector<SILBasicBlock *, 4> boundary;
     computeDominatedBoundaryBlocks(alloc->getParent(), pass.domInfo, boundary);
     for (SILBasicBlock *deallocBlock : boundary) {
-      dealloc(deallocBlock->getTerminator()->getIterator());
-    }
-  } else {
-    for (SILInstruction *deallocPoint : pass.exitingInsts) {
-      dealloc(deallocPoint->getIterator());
+      auto deallocBuilder = pass.getBuilder(deallocBlock->back().getIterator());
+      deallocBuilder.createDeallocStack(pass.genLoc(), alloc);
     }
+    boundary.clear();
   }
-  return alloc;
 }
 
 //===----------------------------------------------------------------------===//
@@ -3850,6 +3888,8 @@ void AddressLowering::runOnFunction(SILFunction *function) {
   // forward order, setting 'storageAddress' for each projection as it goes.
   rewriteFunction(pass);
 
+  allocator.finalizeOpaqueStorage();
+
   deleteRewrittenInstructions(pass);
 
   StackNesting::fixNesting(function);

test/SILOptimizer/opaque_values_Onone.swift

@@ -15,9 +15,9 @@ enum Maybe1<T : Equatable> {
   case yep(T)
   // CHECK-LABEL: sil hidden @maybe1_compare {{.*}} {
   // CHECK:         [[LHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe1
-  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe1
   // CHECK:         switch_enum_addr [[LHS_ADDR]] : $*Maybe1<T>, case #Maybe1.yep!enumelt: [[L_YEP:bb[0-9]+]], case #Maybe1.nope!enumelt: {{bb[0-9]+}}
   // CHECK:       [[L_YEP]]:
+  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe1
   // CHECK:         unchecked_take_enum_data_addr [[LHS_ADDR]] : $*Maybe1<T>, #Maybe1.yep!enumelt
   // CHECK:         switch_enum_addr [[RHS_ADDR]] : $*Maybe1<T>, case #Maybe1.yep!enumelt: [[L_AND_R_YEP:bb[0-9]+]], default {{bb[0-9]+}}
   // CHECK:       [[L_AND_R_YEP]]:
@@ -42,9 +42,9 @@ enum Maybe2<T : Equatable> {
 
   // CHECK-LABEL: sil hidden @maybe2_compare {{.*}} {
   // CHECK:         [[LHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe2<T>
-  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe2<T>
   // CHECK:         switch_enum_addr [[LHS_ADDR]] : $*Maybe2<T>, case #Maybe2.yep!enumelt: [[L_YEP:bb[0-9]+]], case #Maybe2.nope!enumelt: {{bb[0-9]+}}
   // CHECK:       [[L_YEP]]:
+  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe2<T>
   // CHECK:         unchecked_take_enum_data_addr [[LHS_ADDR]] : $*Maybe2<T>, #Maybe2.yep!enumelt
   // CHECK:         switch_enum_addr [[RHS_ADDR]] : $*Maybe2<T>, case #Maybe2.yep!enumelt: [[R_YEP:bb[0-9]+]], default {{bb[0-9]+}}
   // CHECK:       [[L_AND_R_YEP]]: