Merge pull request #39761 from atrick/fix-accessopt

Fix AccessEnforcementOpts for OSSA

Author: atrick

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -22,6 +22,7 @@
 #include "swift/Basic/Defer.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/OwnershipUtils.h"
+#include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 
@@ -34,6 +35,31 @@ inline bool requiresOSSACleanup(SILValue v) {
     && v.getOwnershipKind() != OwnershipKind::Unowned;
 }
 
+/// Rewrite the lifetime of \p ownedValue to match \p lifetimeBoundary. This may
+/// insert copies at forwarding consumes, including phis.
+///
+/// Precondition: lifetimeBoundary is dominated by ownedValue.
+///
+/// Precondition: lifetimeBoundary is a superset of ownedValue's current
+/// lifetime (therefore, none of the safety checks done during
+/// CanonicalizeOSSALifetime are needed here).
+void extendOwnedLifetime(SILValue ownedValue,
+                         PrunedLivenessBoundary &lifetimeBoundary,
+                         InstructionDeleter &deleter);
+
+/// Rewrite the local borrow scope introduced by \p beginBorrow to match \p
+/// guaranteedBoundary.
+///
+/// Precondition: guaranteedBoundary is dominated by beginBorrow which has no
+/// reborrows.
+///
+/// Precondition: guaranteedBoundary is a superset of beginBorrow's current
+/// scope (therefore, none of the safety checks done during
+/// CanonicalizeBorrowScope are needed here).
+void extendLocalBorrow(BeginBorrowInst *beginBorrow,
+                       PrunedLivenessBoundary &guaranteedBoundary,
+                       InstructionDeleter &deleter);
+
 /// Given a new phi that may use a guaranteed value, create nested borrow scopes
 /// for its incoming operands and end_borrows that cover the phi's extended
 /// borrow scope, which transitively includes any phis that use this phi.
@@ -47,6 +73,73 @@ inline bool requiresOSSACleanup(SILValue v) {
 /// newly created phis do not yet have a borrow scope.
 bool createBorrowScopeForPhiOperands(SILPhiArgument *newPhi);
 
+//===----------------------------------------------------------------------===//
+//                        GuaranteedOwnershipExtension
+//===----------------------------------------------------------------------===//
+
+/// Extend existing guaranteed ownership to cover new guaranteeed uses that are
+/// dominated by the borrow introducer.
+class GuaranteedOwnershipExtension {
+  // --- context
+  InstructionDeleter &deleter;
+  DeadEndBlocks &deBlocks;
+
+  // --- analysis state
+  PrunedLiveness guaranteedLiveness;
+  PrunedLiveness ownedLifetime;
+  SmallVector<SILBasicBlock *, 4> ownedConsumeBlocks;
+  BeginBorrowInst *beginBorrow = nullptr;
+
+public:
+  GuaranteedOwnershipExtension(InstructionDeleter &deleter,
+                               DeadEndBlocks &deBlocks)
+      : deleter(deleter), deBlocks(deBlocks) {}
+
+  void clear() {
+    guaranteedLiveness.clear();
+    ownedLifetime.clear();
+    ownedConsumeBlocks.clear();
+    beginBorrow = nullptr;
+  }
+
+  /// Invalid indicates that the current guaranteed scope is insufficient, and
+  /// it does not meet the precondition for scope extension.
+  ///
+  /// Valid indicates that the current guaranteed scope is sufficient with no
+  /// transformation required.
+  ///
+  /// ExtendBorrow indicates that the local borrow scope can be extended without
+  /// affecting the owned lifetime or introducing copies.
+  ///
+  /// ExtendLifetime indicates that the owned lifetime can be extended possibly
+  /// requiring additional copies.
+  enum Status { Invalid, Valid, ExtendBorrow, ExtendLifetime };
+
+  /// Can the OSSA ownership of the \p parentAddress cover all uses of the \p
+  /// childAddress?
+  ///
+  /// Precondition: \p parentAddress dominates \p childAddress
+  Status checkAddressOwnership(SILValue parentAddress, SILValue childAddress);
+
+  /// Can the OSSA scope of \p borrow cover all \p newUses?
+  ///
+  /// Precondition: \p borrow dominates \p newUses
+  Status checkBorrowExtension(BorrowedValue borrow,
+                              ArrayRef<Operand *> newUses);
+
+  /// Can the OSSA scope of \p ownedValue cover all the guaranteed \p newUses?
+  ///
+  /// Precondition: \p ownedValue dominates \p newUses
+  Status checkLifetimeExtension(SILValue ownedValue,
+                                ArrayRef<Operand *> newUses);
+
+  void transform(Status status);
+};
+
+//===----------------------------------------------------------------------===//
+//                      RAUW - Replace All Uses With...
+//===----------------------------------------------------------------------===//
+
 /// A struct that contains context shared in between different operation +
 /// "ownership fixup" utilities. Please do not put actual methods on this, it is
 /// meant to be composed with.

lib/SIL/Utils/PrunedLiveness.cpp

@@ -209,6 +209,7 @@ static void findLastUserInBlock(SILBasicBlock *bb,
       continue;
 
     boundary.lastUsers.push_back(inst);
+    return;
   }
   llvm_unreachable("No user in LiveWithin block");
 }

lib/SILOptimizer/Transforms/AccessEnforcementOpts.cpp

@@ -83,10 +83,12 @@
 #include "swift/SIL/MemAccessUtils.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SILOptimizer/Analysis/AccessStorageAnalysis.h"
+#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SILOptimizer/Analysis/LoopRegionAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
+#include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SCCIterator.h"
 
@@ -1003,10 +1005,30 @@ canMerge(PostDominanceInfo *postDomTree,
   return canMergeEnd(parentIns, childIns);
 }
 
+static bool extendOwnership(BeginAccessInst *parentInst,
+                            BeginAccessInst *childInst,
+                            InstructionDeleter &deleter,
+                            DeadEndBlocks &deBlocks) {
+  GuaranteedOwnershipExtension extension(deleter, deBlocks);
+  auto status = extension.checkAddressOwnership(parentInst, childInst);
+  switch (status) {
+  case GuaranteedOwnershipExtension::Invalid:
+    return false;
+  case GuaranteedOwnershipExtension::Valid:
+    return true;
+  case GuaranteedOwnershipExtension::ExtendLifetime:
+  case GuaranteedOwnershipExtension::ExtendBorrow:
+    break;
+  }
+  extension.transform(status);
+  return true;
+}
+
 /// Perform access merging.
-static bool mergeAccesses(
-    SILFunction *F, PostDominanceInfo *postDomTree,
-    const AccessConflictAndMergeAnalysis::MergeablePairs &mergePairs) {
+static bool
+mergeAccesses(SILFunction *F, PostDominanceInfo *postDomTree,
+              const AccessConflictAndMergeAnalysis::MergeablePairs &mergePairs,
+              InstModCallbacks &callbacks, DeadEndBlocks &deBlocks) {
 
   if (mergePairs.empty()) {
     LLVM_DEBUG(llvm::dbgs() << "Skipping SCC Analysis...\n");
@@ -1042,6 +1064,7 @@ static bool mergeAccesses(
   // begin_access instruction. We store (begin_access %2 -> begin_access %1)
   // to re-map a merged begin_access to it's replaced instruction.
   llvm::DenseMap<BeginAccessInst *, BeginAccessInst *> oldToNewMap;
+  InstructionDeleter deleter(callbacks);
 
   while (!workPairs.empty()) {
     auto curr = workPairs.pop_back_val();
@@ -1058,6 +1081,9 @@ static bool mergeAccesses(
     if (!canMerge(postDomTree, blockToSCCMap, parentIns, childIns))
       continue;
 
+    if (!extendOwnership(parentIns, childIns, deleter, deBlocks))
+      return false;
+
     LLVM_DEBUG(llvm::dbgs()
                << "Merging " << *childIns << " into " << *parentIns << "\n");
 
@@ -1084,8 +1110,9 @@ static bool mergeAccesses(
     auto curr = oldToNewMap.begin();
     auto *oldIns = curr->getFirst();
     oldToNewMap.erase(oldIns);
-    oldIns->eraseFromParent();
+    deleter.forceDelete(oldIns);
   }
+  deleter.cleanupDeadInstructions();
   return changed;
 }
 
@@ -1101,6 +1128,8 @@ struct AccessEnforcementOpts : public SILFunctionTransform {
 
     LoopRegionFunctionInfo *LRFI = getAnalysis<LoopRegionAnalysis>()->get(F);
     PostOrderFunctionInfo *PO = getAnalysis<PostOrderAnalysis>()->get(F);
+    DeadEndBlocksAnalysis *deBlocksAnalysis =
+        PM->getAnalysis<DeadEndBlocksAnalysis>();
     AccessStorageAnalysis *ASA = getAnalysis<AccessStorageAnalysis>();
     AccessConflictAndMergeAnalysis a(LRFI, PO, ASA);
     if (!a.analyze())
@@ -1132,7 +1161,9 @@ struct AccessEnforcementOpts : public SILFunctionTransform {
     PostDominanceAnalysis *postDomAnalysis =
         getAnalysis<PostDominanceAnalysis>();
     PostDominanceInfo *postDomTree = postDomAnalysis->get(F);
-    if (mergeAccesses(F, postDomTree, result.mergePairs))
+    DeadEndBlocks *deBlocks = deBlocksAnalysis->get(F);
+    InstModCallbacks callbacks;
+    if (mergeAccesses(F, postDomTree, result.mergePairs, callbacks, *deBlocks))
       invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
   }
 };