Fix AccessEnforcementOpts for OSSA.

Add extendOwnership to fixup OSSA before merging access scopes.

Uses the new GuaranteedOwnershipExtension utility.

Author: atrick

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);
   }
 };

test/SILOptimizer/access_enforcement_opts_ossa.sil

@@ -18,6 +18,10 @@ class Klass {
 
 }
 
+class NonTrivial {
+  @_hasStorage var object: Klass { get set }
+}
+
 var globalX: X
 
 var globalOtherX: X
@@ -26,6 +30,8 @@ sil_global hidden @globalX : $X
 
 sil_global hidden @globalOtherX : $X
 
+sil [ossa] @hasIndirectArg : $@convention(thin) (@inout Klass) -> ()
+
 sil hidden [ossa] @Xinit : $@convention(method) (@thin X.Type) -> X {
 bb0(%0 : $@thin X.Type):
   %1 = alloc_stack $X, var, name "self"
@@ -1821,6 +1827,10 @@ class Bar {
   func foo()
 }
 
+// Nothing happens here. I guess we check that the analysis does not crash.
+//
+// CHECK-LABEL: sil [ossa] @$testEndLifetime : $@convention(method) (@owned Bar, @in Bar) -> () {
+// CHECK-LABEL: } // end sil function '$testEndLifetime'
 sil [ossa] @$testEndLifetime : $@convention(method) (@owned Bar, @in Bar) -> () {
 bb0(%0 : @owned $Bar, %1 : $*Bar):
   %b = begin_access [modify] [dynamic] %1 : $*Bar
@@ -1832,3 +1842,35 @@ bb0(%0 : @owned $Bar, %1 : $*Bar):
   %r = tuple ()
   return %r : $()
 }
+
+// CHECK-LABEL: sil [ossa] @testBorrowedAccess : $@convention(thin) (@owned NonTrivial) -> () {
+// CHECK: bb0(%0 : @owned $NonTrivial):
+// CHECK:   [[BORROW:%.*]] = begin_borrow %0 : $NonTrivial
+// CHECK:   [[ADR:%.*]] = ref_element_addr [[BORROW]] : $NonTrivial, #NonTrivial.object
+// CHECK:   [[ACCESS:%.*]] = begin_access [modify] [dynamic] [[ADR]] : $*Klass
+// CHECK-NOT: access
+// CHECK-NOT: borrow
+// CHECK:   apply %{{.*}}([[ACCESS]]) : $@convention(thin) (@inout Klass) -> ()
+// CHECK:   apply %{{.*}}([[ACCESS]]) : $@convention(thin) (@inout Klass) -> ()
+// CHECK:   end_access [[ACCESS]] : $*Klass
+// CHECK:   end_borrow [[BORROW]] : $NonTrivial
+// CHECK:   destroy_value %0 : $NonTrivial
+// CHECK-LABEL: } // end sil function 'testBorrowedAccess'
+sil [ossa] @testBorrowedAccess : $@convention(thin) (@owned NonTrivial) -> () {
+bb0(%0 : @owned $NonTrivial):
+  %f = function_ref @hasIndirectArg : $@convention(thin) (@inout Klass) -> ()
+  %b1 = begin_borrow %0 : $NonTrivial
+  %e1 = ref_element_addr %b1 : $NonTrivial, #NonTrivial.object
+  %a1 = begin_access [modify] [dynamic] %e1 : $*Klass
+  %c1 = apply %f(%a1) : $@convention(thin) (@inout Klass) -> ()
+  end_access %a1 : $*Klass
+  end_borrow %b1 : $NonTrivial
+  %b2 = begin_borrow %0 : $NonTrivial
+  %e2 = ref_element_addr %b2 : $NonTrivial, #NonTrivial.object
+  %a2 = begin_access [modify] [dynamic] %e2 : $*Klass
+  %c2 = apply %f(%a2) : $@convention(thin) (@inout Klass) -> ()
+  end_access %a2 : $*Klass
+  end_borrow %b2 : $NonTrivial
+  destroy_value %0 : $NonTrivial
+  return undef : $()
+}