Merge pull request #60478 from nate-chandler/lexical_lifetimes/optimize_more_alloc_stack_destroys

[SILOpt] Hoist destroys for non-lexical alloc_stacks more aggressively.

Author: nate-chandler

lib/SIL/Verifier/SILVerifier.cpp

@@ -2227,6 +2227,31 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
             "Found load borrow that is invalidated by a local write?!");
   }
 
+  void checkBeginBorrowInst(BeginBorrowInst *bbi) {
+    if (!bbi->isLexical())
+      return;
+    // Lexical begin_borrows of instances of some SILBoxType must derive from
+    // alloc_boxes or captures.
+    auto value = bbi->getOperand();
+    if (!value->getType().is<SILBoxType>())
+      return;
+    while (true) {
+      // Inlining may introduce additional begin_borrow instructions.
+      if (auto bbi = dyn_cast<BeginBorrowInst>(value))
+        value = bbi->getOperand();
+      // SILGen introduces copy_value instructions.
+      else if (auto cvi = dyn_cast<CopyValueInst>(value))
+        value = cvi->getOperand();
+      // SILGen inserts mark_uninitialized instructions of alloc_boxes.
+      else if (auto *mui = dyn_cast<MarkUninitializedInst>(value))
+        value = mui->getOperand();
+      else
+        break;
+    }
+    require(isa<AllocBoxInst>(value) || isa<SILFunctionArgument>(value),
+            "Lexical borrows of SILBoxTypes must be of vars or captures.");
+  }
+
   void checkEndBorrowInst(EndBorrowInst *EBI) {
     require(
         F.hasOwnership(),

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -13,12 +13,13 @@
 #define DEBUG_TYPE "allocbox-to-stack"
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/Basic/BlotMapVector.h"
+#include "swift/Basic/GraphNodeWorklist.h"
 #include "swift/SIL/ApplySite.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILCloner.h"
-#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -27,6 +28,7 @@
 #include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "swift/SILOptimizer/Utils/ValueLifetime.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -538,13 +540,35 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
   SILBuilderWithScope Builder(ABI);
   assert(ABI->getBoxType()->getLayout()->getFields().size() == 1
          && "rewriting multi-field box not implemented");
-  auto &mod = ABI->getFunction()->getModule();
-  bool isLexical = mod.getASTContext().SILOpts.supportsLexicalLifetimes(mod);
-  auto *ASI = Builder.createAllocStack(
-      ABI->getLoc(),
-      getSILBoxFieldType(TypeExpansionContext(*ABI->getFunction()),
-                         ABI->getBoxType(), ABI->getModule().Types, 0),
-      ABI->getVarInfo(), ABI->hasDynamicLifetime(), isLexical);
+  auto ty = getSILBoxFieldType(TypeExpansionContext(*ABI->getFunction()),
+                               ABI->getBoxType(), ABI->getModule().Types, 0);
+  auto isLexical = [&]() -> bool {
+    auto &mod = ABI->getFunction()->getModule();
+    bool lexicalLifetimesEnabled =
+        mod.getASTContext().SILOpts.supportsLexicalLifetimes(mod);
+    if (!lexicalLifetimesEnabled)
+      return false;
+    // Look for lexical borrows of the alloc_box.
+    GraphNodeWorklist<Operand *, 4> worklist;
+    worklist.initializeRange(ABI->getUses());
+    while (auto *use = worklist.pop()) {
+      // See through mark_uninitialized and non-lexical begin_borrow
+      // instructions.  It's verified that lexical begin_borrows of SILBoxType
+      // values originate either from AllocBoxInsts or SILFunctionArguments.
+      if (auto *mui = dyn_cast<MarkUninitializedInst>(use->getUser())) {
+        for (auto *use : mui->getUses())
+          worklist.insert(use);
+      } else if (auto *bbi = dyn_cast<BeginBorrowInst>(use->getUser())) {
+        if (bbi->isLexical())
+          return true;
+        for (auto *use : bbi->getUses())
+          worklist.insert(use);
+      }
+    }
+    return false;
+  };
+  auto *ASI = Builder.createAllocStack(ABI->getLoc(), ty, ABI->getVarInfo(),
+                                       ABI->hasDynamicLifetime(), isLexical());
 
   // Transfer a mark_uninitialized if we have one.
   SILValue StackBox = ASI;

lib/SILOptimizer/Transforms/SSADestroyHoisting.cpp

@@ -962,7 +962,7 @@ void SSADestroyHoisting::run() {
   }
   // Alloc stacks always enclose their accesses.
   for (auto *asi : asis) {
-    changed |= hoistDestroys(asi, /*ignoreDeinitBarriers=*/false,
+    changed |= hoistDestroys(asi, /*ignoreDeinitBarriers=*/!asi->isLexical(),
                              remainingDestroyAddrs, deleter);
   }
   // Arguments enclose everything.

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -501,11 +501,29 @@ void TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
   if (!tempObj)
     return;
 
-  // If the storage corresponds to a source-level var, do not optimize here.
-  // Mem2Reg will transform the lexical alloc_stack into a lexical begin_borrow
-  // which will ensure that the value's lifetime isn't observably shortened.
-  if (tempObj->isLexical())
-    return;
+  // If the temporary storage is lexical, it either came from a source-level var
+  // or was marked lexical because it was passed to a function that has been
+  // inlined.
+  // TODO: [begin_borrow_addr] Once we can mark addresses as being borrowed, we
+  //       won't need to mark alloc_stacks lexical during inlining.  At that
+  //       point, the above comment should change, but the implementation
+  //       remains the same.
+  //
+  // In either case, we can eliminate the temporary if the source of the copy is
+  // lexical and it is live for longer than the temporary.
+  if (tempObj->isLexical()) {
+    // TODO: Determine whether the base of the copy_addr's source is lexical and
+    //       its live range contains the range in which the alloc_stack
+    //       contains the value copied into it via the copy_addr.
+    //
+    // For now, only look for guaranteed arguments.
+    auto storage = AccessStorageWithBase::compute(copyInst->getSrc());
+    if (!storage.base)
+      return;
+    if (auto *arg = dyn_cast<SILFunctionArgument>(storage.base))
+      if (arg->getOwnershipKind() != OwnershipKind::Guaranteed)
+        return;
+  }
 
   bool isOSSA = copyInst->getFunction()->hasOwnership();
 
@@ -645,10 +663,20 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
     return std::next(si->getIterator());
   }
 
-  // If the storage corresponds to a source-level var, do not optimize here.
-  // Mem2Reg will transform the lexical alloc_stack into a lexical begin_borrow
-  // which will ensure that the value's lifetime isn't observably shortened.
+  // If the temporary storage is lexical, it either came from a source-level var
+  // or was marked lexical because it was passed to a function that has been
+  // inlined.
+  // TODO: [begin_borrow_addr] Once we can mark addresses as being borrowed, we
+  //       won't need to mark alloc_stacks lexical during inlining.  At that
+  //       point, the above comment should change, but the implementation
+  //       remains the same.
+  //
+  // In either case, we can eliminate the temporary if the source of the store
+  // is lexical and it is live for longer than the temporary.
   if (tempObj->isLexical()) {
+    // TODO: Find the lexical root of the source, if any, and allow optimization
+    //       if its live range contains the range in which the alloc_stack
+    //       contains the value stored into it.
     return std::next(si->getIterator());
   }
 

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -15,6 +15,8 @@
 #include "swift/SILOptimizer/Utils/SILInliner.h"
 #include "swift/AST/Builtins.h"
 #include "swift/AST/DiagnosticsSIL.h"
+#include "swift/Basic/Defer.h"
+#include "swift/SIL/MemAccessUtils.h"
 #include "swift/SIL/PrettyStackTrace.h"
 #include "swift/SIL/SILDebugScope.h"
 #include "swift/SIL/SILInstruction.h"
@@ -433,17 +435,44 @@ void SILInlineCloner::cloneInline(ArrayRef<SILValue> AppliedArgs) {
   auto calleeConv = getCalleeFunction()->getConventions();
   for (auto p : llvm::enumerate(AppliedArgs)) {
     SILValue callArg = p.value();
+    SWIFT_DEFER { entryArgs.push_back(callArg); };
     unsigned idx = p.index();
     if (idx >= calleeConv.getSILArgIndexOfFirstParam()) {
-      // Insert begin/end borrow for guaranteed arguments.
-      if (calleeConv.getParamInfoForSILArg(idx).isGuaranteed()) {
-        if (SILValue newValue = borrowFunctionArgument(callArg, Apply)) {
-          callArg = newValue;
-          borrowedArgs[idx] = true;
+      auto paramInfo = calleeConv.getParamInfoForSILArg(idx);
+      if (callArg->getType().isAddress()) {
+        // If lexical lifetimes are enabled, any alloc_stacks in the caller that
+        // are passed to the callee being inlined (except mutating exclusive
+        // accesses) need to be promoted to be lexical.  Otherwise,
+        // destroy_addrs could be hoisted through the body of the newly inlined
+        // function without regard to the deinit barriers it contains.
+        //
+        // TODO: [begin_borrow_addr] Instead of marking the alloc_stack as a
+        //       whole lexical, just mark the inlined range lexical via
+        //       begin_borrow_addr [lexical]/end_borrow_addr just as is done
+        //       with values.
+        auto &module = Apply.getFunction()->getModule();
+        auto enableLexicalLifetimes =
+            module.getASTContext().SILOpts.supportsLexicalLifetimes(module);
+        if (!enableLexicalLifetimes)
+          continue;
+
+        // Exclusive mutating accesses don't entail a lexical scope.
+        if (paramInfo.getConvention() == ParameterConvention::Indirect_Inout)
+          continue;
+
+        auto storage = AccessStorageWithBase::compute(callArg);
+        if (auto *asi = dyn_cast<AllocStackInst>(storage.base))
+          asi->setIsLexical();
+      } else {
+        // Insert begin/end borrow for guaranteed arguments.
+        if (paramInfo.isGuaranteed()) {
+          if (SILValue newValue = borrowFunctionArgument(callArg, Apply)) {
+            callArg = newValue;
+            borrowedArgs[idx] = true;
+          }
         }
       }
     }
-    entryArgs.push_back(callArg);
   }
 
   // Create the return block and set ReturnToBB for use in visitTerminator

test/SILGen/copy_operator.swift

@@ -37,7 +37,7 @@ class Klass {}
 // CHECK-SIL-LABEL: sil @$s8moveonly7useCopyyAA5KlassCADF : $@convention(thin) (@guaranteed Klass) -> @owned Klass {
 // CHECK-SIL: bb0([[ARG:%.*]] : $Klass):
 // CHECK-SIL-NEXT: debug_value
-// CHECK-SIL-NEXT: [[INPUT:%.*]] = alloc_stack $Klass
+// CHECK-SIL-NEXT: [[INPUT:%.*]] = alloc_stack [lexical] $Klass
 // CHECK-SIL-NEXT: store [[ARG]] to [[INPUT]]
 // CHECK-SIL-NEXT: [[VALUE:%[0-9][0-9]*]] = load [[INPUT]]{{.*}}
 // CHECK-SIL-NEXT: strong_retain [[VALUE]]
@@ -78,7 +78,7 @@ public func useCopy(_ k: Klass) -> Klass {
 // CHECK-SIL-LABEL: sil @$s8moveonly7useCopyyxxRlzClF : $@convention(thin) <T where T : AnyObject> (@guaranteed T) -> @owned T {
 // CHECK-SIL: bb0([[ARG:%.*]] :
 // CHECK-SIL-NEXT: debug_value
-// CHECK-SIL-NEXT: [[INPUT_TO_BE_ELIMINATED:%.*]] = alloc_stack $T
+// CHECK-SIL-NEXT: [[INPUT_TO_BE_ELIMINATED:%.*]] = alloc_stack [lexical] $T
 // CHECK-SIL-NEXT: store [[ARG]] to [[INPUT]] : $*T
 // CHECK-SIL-NEXT: [[VALUE:%.*]] = load [[INPUT]] : $*T
 // CHECK-SIL-NEXT: strong_retain [[VALUE]]

test/SILOptimizer/allocbox_to_stack_lifetime.sil

@@ -14,8 +14,13 @@ struct Bool {
 
 protocol Error {}
 
-// CHECK-LABEL: sil [ossa] @simple_promotion
-sil [ossa] @simple_promotion : $@convention(thin) (Int) -> Int {
+// CHECK-LABEL: sil [ossa] @promote_nonlexical
+// CHECK:         alloc_stack $
+// CHECK-NOT:     alloc_box
+// CHECK-NOT:     destroy_value
+// CHECK:         return
+// CHECK-LABEL: } // end sil function 'promote_nonlexical'
+sil [ossa] @promote_nonlexical : $@convention(thin) (Int) -> Int {
 bb0(%0 : $Int):
   %1 = alloc_box ${ var Int }
   %1a = project_box %1 : ${ var Int }, 0
@@ -24,9 +29,24 @@ bb0(%0 : $Int):
   %3 = load [trivial] %1a : $*Int
   destroy_value %1 : ${ var Int }
   return %3 : $Int
-// CHECK: alloc_stack [lexical]
-// CHECK-NOT: alloc_box
-// CHECK-NOT: destroy_value
-// CHECK: return
+}
+
+// CHECK-LABEL: sil [ossa] @promote_lexical
+// CHECK:         alloc_stack [lexical]
+// CHECK-NOT:     alloc_box
+// CHECK-NOT:     destroy_value
+// CHECK:         return
+// CHECK-LABEL: } // end sil function 'promote_lexical'
+sil [ossa] @promote_lexical : $@convention(thin) (Int) -> Int {
+bb0(%0 : $Int):
+  %1 = alloc_box ${ var Int }
+  %b = begin_borrow [lexical] %1 : ${ var Int }
+  %1a = project_box %b : ${ var Int }, 0
+  store %0 to [trivial] %1a : $*Int
+
+  %3 = load [trivial] %1a : $*Int
+  end_borrow %b : ${ var Int }
+  destroy_value %1 : ${ var Int }
+  return %3 : $Int
 }