Merge pull request #28653 from gottesmm/pr-48051636c430a08d11def5e474d526a5e8e1eed1

[semantic-arc-opts] Convert load [copy] -> load_borrow given single init alloc_stack.

Author: gottesmm

include/swift/SIL/MemAccessUtils.h

@@ -488,6 +488,17 @@ AccessedStorage findAccessedStorageNonNested(SILValue sourceAddr);
 /// uninitialized.
 bool memInstMustInitialize(Operand *memOper);
 
+/// Is this an alloc_stack instruction that is:
+///
+/// 1. Only initialized once in its own def block.
+/// 2. Never written to again except by destroy_addr.
+///
+/// On return, destroyingUsers contains the list of users that destroy the
+/// alloc_stack. If the alloc_stack is destroyed in pieces, we do not guarantee
+/// that the list of destroying users is a minimal jointly post-dominating set.
+bool isSingleInitAllocStack(AllocStackInst *asi,
+                            SmallVectorImpl<SILInstruction *> &destroyingUsers);
+
 /// Return true if the given address producer may be the source of a formal
 /// access (a read or write of a potentially aliased, user visible variable).
 ///

include/swift/SIL/OwnershipUtils.h

@@ -183,6 +183,21 @@ class LinearLifetimeChecker {
                        nullptr /*leakingBlocks*/)
                 .getFoundError();
   }
+
+  bool validateLifetime(SILValue value,
+                        ArrayRef<SILInstruction *> consumingUses,
+                        ArrayRef<SILInstruction *> nonConsumingUses) {
+    auto *consumingUsesCast =
+        reinterpret_cast<const BranchPropagatedUser *>(consumingUses.data());
+    auto *nonConsumingUsesCast =
+        reinterpret_cast<const BranchPropagatedUser *>(nonConsumingUses.data());
+    ArrayRef<BranchPropagatedUser> consumingUsesCastArray(consumingUsesCast,
+                                                          consumingUses.size());
+    ArrayRef<BranchPropagatedUser> nonConsumingUsesCastArray(
+        nonConsumingUsesCast, nonConsumingUses.size());
+    return validateLifetime(value, consumingUsesCastArray,
+                            nonConsumingUsesCastArray);
+  }
 };
 
 /// Returns true if v is an address or trivial.

lib/SIL/MemAccessUtils.cpp

@@ -14,6 +14,7 @@
 
 #include "swift/SIL/MemAccessUtils.h"
 #include "swift/SIL/ApplySite.h"
+#include "swift/SIL/Projection.h"
 #include "swift/SIL/SILGlobalVariable.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/SILUndef.h"
@@ -660,3 +661,64 @@ SILBasicBlock::iterator swift::removeBeginAccess(BeginAccessInst *beginAccess) {
   }
   return beginAccess->getParent()->erase(beginAccess);
 }
+
+bool swift::isSingleInitAllocStack(
+    AllocStackInst *asi, SmallVectorImpl<SILInstruction *> &destroyingUsers) {
+  // For now, we just look through projections and rely on memInstMustInitialize
+  // to classify all other uses as init or not.
+  SmallVector<Operand *, 32> worklist(asi->getUses());
+  bool foundInit = false;
+
+  while (!worklist.empty()) {
+    auto *use = worklist.pop_back_val();
+    auto *user = use->getUser();
+
+    if (Projection::isAddressProjection(user) ||
+        isa<OpenExistentialAddrInst>(user)) {
+      // Look through address projections.
+      for (SILValue r : user->getResults()) {
+        llvm::copy(r->getUses(), std::back_inserter(worklist));
+      }
+      continue;
+    }
+
+    if (auto *li = dyn_cast<LoadInst>(user)) {
+      // If we are not taking,
+      if (li->getOwnershipQualifier() != LoadOwnershipQualifier::Take) {
+        continue;
+      }
+      // Treat load [take] as a write.
+      return false;
+    }
+
+    switch (user->getKind()) {
+    default:
+      break;
+    case SILInstructionKind::DestroyAddrInst:
+      destroyingUsers.push_back(user);
+      continue;
+    case SILInstructionKind::DeallocStackInst:
+    case SILInstructionKind::LoadBorrowInst:
+    case SILInstructionKind::DebugValueAddrInst:
+      continue;
+    }
+
+    // See if we have an initializer and that such initializer is in the same
+    // block.
+    if (memInstMustInitialize(use)) {
+      if (user->getParent() != asi->getParent() || foundInit) {
+        return false;
+      }
+
+      foundInit = true;
+      continue;
+    }
+
+    // Otherwise, if we have found something not in our whitelist, return false.
+    return false;
+  }
+
+  // We did not find any users that we did not understand. So we can
+  // conservatively return true here.
+  return true;
+}

lib/SILOptimizer/Mandatory/SemanticARCOpts.cpp

@@ -280,7 +280,7 @@ struct SemanticARCOptVisitor
   FORWARDING_TERM(CondBranch)
 #undef FORWARDING_TERM
 
-  bool isWrittenTo(LoadInst *li);
+  bool isWrittenTo(LoadInst *li, ArrayRef<SILInstruction *> destroys);
 
   bool processWorklist();
 
@@ -664,6 +664,7 @@ bool mayFunctionMutateArgument(const AccessedStorage &storage, SILFunction &f) {
 // Then find our accessed storage to determine whether it provides a guarantee
 // for the loaded value.
 namespace {
+
 class StorageGuaranteesLoadVisitor
   : public AccessUseDefChainVisitor<StorageGuaranteesLoadVisitor>
 {
@@ -677,12 +678,15 @@ class StorageGuaranteesLoadVisitor
   SILValue currentAddress;
 
   Optional<bool> isWritten;
-  
+
+  ArrayRef<SILInstruction *> destroyValues;
+
 public:
-  StorageGuaranteesLoadVisitor(SemanticARCOptVisitor &arcOpt, LoadInst *load)
-    : ARCOpt(arcOpt), Load(load), currentAddress(load->getOperand())
-  {}
-  
+  StorageGuaranteesLoadVisitor(SemanticARCOptVisitor &arcOpt, LoadInst *load,
+                               ArrayRef<SILInstruction *> destroyValues)
+      : ARCOpt(arcOpt), Load(load), currentAddress(load->getOperand()),
+        destroyValues(destroyValues) {}
+
   void answer(bool written) {
     currentAddress = nullptr;
     isWritten = written;
@@ -790,17 +794,38 @@ class StorageGuaranteesLoadVisitor
     return answer(true);
   }
 
+  /// See if we have an alloc_stack that is only written to once by an
+  /// initializing instruction.
+  void visitStackAccess(AllocStackInst *stack) {
+    SmallVector<SILInstruction *, 8> destroyAddrs;
+    bool initialAnswer = isSingleInitAllocStack(stack, destroyAddrs);
+    if (!initialAnswer)
+      return answer(true);
+
+    // Then make sure that all of our load [copy] uses are within the
+    // destroy_addr.
+    SmallPtrSet<SILBasicBlock *, 4> visitedBlocks;
+    LinearLifetimeChecker checker(visitedBlocks, ARCOpt.getDeadEndBlocks());
+    // Returns true on success. So we invert.
+    bool foundError =
+        !checker.validateLifetime(stack, destroyAddrs /*consuming users*/,
+                                  destroyValues /*non consuming users*/);
+    return answer(foundError);
+  }
+
   bool doIt() {
     while (currentAddress) {
       visit(currentAddress);
     }
     return *isWritten;
   }
 };
+
 } // namespace
 
-bool SemanticARCOptVisitor::isWrittenTo(LoadInst *load) {
-  StorageGuaranteesLoadVisitor visitor(*this, load);
+bool SemanticARCOptVisitor::isWrittenTo(LoadInst *load,
+                                        ArrayRef<SILInstruction *> destroys) {
+  StorageGuaranteesLoadVisitor visitor(*this, load, destroys);
   return visitor.doIt();
 }
 
@@ -825,7 +850,7 @@ bool SemanticARCOptVisitor::visitLoadInst(LoadInst *li) {
   // Then check if our address is ever written to. If it is, then we cannot use
   // the load_borrow because the stored value may be released during the loaded
   // value's live range.
-  if (isWrittenTo(li))
+  if (isWrittenTo(li, destroyValues))
     return false;
 
   // Ok, we can perform our optimization. Convert the load [copy] into a

test/SILOptimizer/semantic-arc-opts.sil

@@ -947,4 +947,99 @@ bb0(%0 : @guaranteed $ClassLet):
   destroy_value %4 : $Klass
   %9999 = tuple()
   return %9999 : $()
-}
+}
+
+// CHECK-LABEL: sil [ossa] @single_init_allocstack : $@convention(thin) (@owned Klass) -> () {
+// CHECK-NOT: load [copy]
+// CHECK: } // end sil function 'single_init_allocstack'
+sil [ossa] @single_init_allocstack : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  %1 = alloc_stack $Klass
+  store %0 to [init] %1 : $*Klass
+  %2 = load [copy] %1 : $*Klass
+
+  %3 = function_ref @guaranteed_klass_user : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %3(%2) : $@convention(thin) (@guaranteed Klass) -> ()
+
+  destroy_value %2 : $Klass
+  destroy_addr %1 : $*Klass
+  dealloc_stack %1 : $*Klass
+
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @multiple_init_allocstack : $@convention(thin) (@owned Klass) -> () {
+// CHECK: load [copy]
+// CHECK: } // end sil function 'multiple_init_allocstack'
+sil [ossa] @multiple_init_allocstack : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  %0a = copy_value %0 : $Klass
+  %1 = alloc_stack $Klass
+  store %0 to [init] %1 : $*Klass
+  %2 = load [copy] %1 : $*Klass
+
+  %3 = function_ref @guaranteed_klass_user : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %3(%2) : $@convention(thin) (@guaranteed Klass) -> ()
+
+  destroy_value %2 : $Klass
+  destroy_addr %1 : $*Klass
+
+  store %0a to [init] %1 : $*Klass
+  destroy_addr %1 : $*Klass
+  dealloc_stack %1 : $*Klass
+
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// We could support this, but for now we are keeping things simple. If we do add
+// support, this test will need to be updated.
+//
+// CHECK-LABEL: sil [ossa] @single_init_wrongblock : $@convention(thin) (@owned Klass) -> () {
+// CHECK: load [copy]
+// CHECK: } // end sil function 'single_init_wrongblock'
+sil [ossa] @single_init_wrongblock : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  %1 = alloc_stack $Klass
+  br bb1
+
+bb1:
+  store %0 to [init] %1 : $*Klass
+  %2 = load [copy] %1 : $*Klass
+
+  %3 = function_ref @guaranteed_klass_user : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %3(%2) : $@convention(thin) (@guaranteed Klass) -> ()
+
+  destroy_value %2 : $Klass
+  destroy_addr %1 : $*Klass
+  dealloc_stack %1 : $*Klass
+
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// We could support this, but for now we are keeping things simple. If we do add
+// support, this test will need to be updated.
+//
+// CHECK-LABEL: sil [ossa] @single_init_loadtake : $@convention(thin) (@owned Klass) -> () {
+// CHECK: load [copy]
+// CHECK: } // end sil function 'single_init_loadtake'
+sil [ossa] @single_init_loadtake : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  %1 = alloc_stack $Klass
+  store %0 to [init] %1 : $*Klass
+  %2 = load [copy] %1 : $*Klass
+
+  %3 = function_ref @guaranteed_klass_user : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %3(%2) : $@convention(thin) (@guaranteed Klass) -> ()
+
+  destroy_value %2 : $Klass
+
+  %4 = load [take] %1 : $*Klass
+  destroy_value %4 : $Klass
+  dealloc_stack %1 : $*Klass
+
+  %9999 = tuple()
+  return %9999 : $()
+}