[ShrinkBorrowScope] Adopt VisitBarrierAccessScopes.

Avoids hoisting borrow scopes into unrelated access scopes which could
introduce exclusivity violations.

rdar://93060369

Author: nate-chandler

lib/SILOptimizer/Utils/ShrinkBorrowScope.cpp

@@ -21,6 +21,7 @@
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SILOptimizer/Analysis/Reachability.h"
+#include "swift/SILOptimizer/Analysis/VisitBarrierAccessScopes.h"
 #include "swift/SILOptimizer/Utils/CanonicalizeBorrowScope.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstructionDeleter.h"
@@ -134,10 +135,13 @@ struct DeinitBarriers final {
   DeinitBarriers &operator=(DeinitBarriers const &) = delete;
 };
 
+class BarrierAccessScopeFinder;
+
 /// Works backwards from the current location of end_borrows to the earliest
 /// place they can be hoisted to.
 ///
 /// Implements IterativeBackwardReachability::Effects.
+/// Implements BarrierAccessScopeFinder::Effects
 class DataFlow final {
 public:
   using Reachability = IterativeBackwardReachability<DataFlow>;
@@ -166,6 +170,8 @@ class DataFlow final {
 
 private:
   friend Reachability;
+  friend class BarrierAccessScopeFinder;
+  friend class VisitBarrierAccessScopes<DataFlow, BarrierAccessScopeFinder>;
 
   Classification classifyInstruction(SILInstruction *);
 
@@ -178,6 +184,33 @@ class DataFlow final {
   Optional<Effect> effectForPhi(SILBasicBlock *);
 };
 
+/// Finds end_access instructions which are barriers to hoisting because the
+/// access scopes they contain barriers to hoisting.  Hoisting end_borrows into
+/// such access scopes could introduce exclusivity violations.
+///
+/// Implements BarrierAccessScopeFinder::Visitor
+class BarrierAccessScopeFinder final {
+  using Impl = VisitBarrierAccessScopes<DataFlow, BarrierAccessScopeFinder>;
+  Context const &context;
+  Impl impl;
+  DataFlow &dataflow;
+  Optional<SmallVector<SILBasicBlock *, 16>> cachedRoots;
+
+public:
+  BarrierAccessScopeFinder(Context const &context, DataFlow &dataflow)
+      : context(context), impl(&context.function, dataflow, *this),
+        dataflow(dataflow) {}
+
+  void find();
+
+private:
+  friend Impl;
+
+  ArrayRef<SILBasicBlock *> roots();
+  bool isInRegion(SILBasicBlock *);
+  void visitBarrierAccessScope(BeginAccessInst *);
+};
+
 /// Whether the specified value is %lifetime or its iterated copy_value.
 ///
 /// In other words, it has to be a simple extended def of %lifetime.
@@ -198,6 +231,8 @@ bool isSimpleExtendedIntroducerDef(Context const &context, SILValue value) {
 
 void DataFlow::run() {
   reachability.initialize();
+  BarrierAccessScopeFinder finder(context, *this);
+  finder.find();
   reachability.solve();
   recordCopies = true;
   reachability.findBarriers(barriers.instructions, barriers.phis,
@@ -217,6 +252,11 @@ DataFlow::classifyInstruction(SILInstruction *instruction) {
   if (uses.users.contains(instruction)) {
     return Classification::Barrier;
   }
+  if (auto *eai = dyn_cast<EndAccessInst>(instruction)) {
+    return barrierAccessScopes.contains(eai->getBeginAccess())
+               ? Classification::Barrier
+               : Classification::Other;
+  }
   if (isDeinitBarrier(instruction)) {
     return Classification::Barrier;
   }
@@ -260,6 +300,33 @@ Optional<DataFlow::Effect> DataFlow::effectForPhi(SILBasicBlock *block) {
   return isBarrier ? Optional<Effect>{Effect::Kill} : llvm::None;
 }
 
+void BarrierAccessScopeFinder::find() { impl.visit(); }
+
+ArrayRef<SILBasicBlock *> BarrierAccessScopeFinder::roots() {
+  if (cachedRoots)
+    return *cachedRoots;
+
+  cachedRoots = SmallVector<SILBasicBlock *, 16>{};
+  BasicBlockSet seenRoots(&context.function);
+  for (auto *gen : dataflow.gens()) {
+    seenRoots.insert(gen->getParent());
+    cachedRoots->push_back(gen->getParent());
+  }
+
+  return *cachedRoots;
+}
+
+bool BarrierAccessScopeFinder::isInRegion(SILBasicBlock *block) {
+  return dataflow.result.discoveredBlocks.contains(block);
+}
+
+void BarrierAccessScopeFinder::visitBarrierAccessScope(BeginAccessInst *bai) {
+  dataflow.barrierAccessScopes.insert(bai);
+  for (auto *eai : bai->getEndAccesses()) {
+    dataflow.reachability.addKill(eai);
+  }
+}
+
 /// Hoist the scope ends of %lifetime, rewriting copies and borrows along the
 /// way.
 class Rewriter final {

test/SILOptimizer/shrink_borrow_scope.sil

@@ -1097,3 +1097,76 @@ bb0(%0 : $*ClassWrapper):
 // =============================================================================
 // instruction tests                                                          }}
 // =============================================================================
+
+// =============================================================================
+// access scope tests                                                         {{
+// =============================================================================
+
+// Don't hoist into an access scope that contains a barrier.
+//
+// CHECK-LABEL: sil [ossa] @nofold_scoped_load_barrier : {{.*}} {
+// CHECK: end_access
+// CHECK: end_access
+// CHECK: end_borrow
+// CHECK-LABEL: // end sil function 'nofold_scoped_load_barrier'
+sil [ossa] @nofold_scoped_load_barrier : $@convention(thin) (@owned C, @owned C) -> (@owned C) {
+entry(%instance : @owned $C, %other : @owned $C):
+  %lifetime = begin_borrow [lexical] %instance : $C
+  apply undef(%lifetime) : $@convention(thin) (@guaranteed C) -> ()
+  %addr = alloc_stack $C
+  %store_scope = begin_access [modify] [static] %addr : $*C
+  store %other to [init] %store_scope : $*C
+  end_access %store_scope : $*C
+  %load_scope = begin_access [read] [static] %addr : $*C
+  %value = load [copy] %load_scope : $*C
+  %barrier = function_ref @barrier : $@convention(thin) () -> ()
+  apply %barrier() : $@convention(thin) () -> ()
+  end_access %load_scope : $*C
+  destroy_addr %addr : $*C
+  dealloc_stack %addr : $*C
+  end_borrow %lifetime : $C
+  destroy_value %instance : $C
+  return %value : $C
+}
+
+// Access scopes that are open at barrier blocks are barriers.  Otherwise, we
+// would hoist end_borrows into the scopes when the end_borrows are hoisted up
+// to the begin of blocks whose predecessor is the barrier block.
+//
+// CHECK-LABEL: sil [ossa] @nohoist_into_access_scope_barred_by_barrier_block : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C, [[INOUT:%[^,]+]] : $*C):
+// CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [lexical] [[INSTANCE]]
+// CHECK:         [[SCOPE:%[^,]+]] = begin_access [modify] [static] [[INOUT]] : $*C
+// CHECK:         cond_br undef, [[LEFT:bb[0-9]+]],
+// CHECK:       [[LEFT]]:
+// CHECK:         end_access [[SCOPE]] : $*C
+// CHECK-NEXT:    end_borrow [[LIFETIME]] : $C
+// CHECK-LABEL: } // end sil function 'nohoist_into_access_scope_barred_by_barrier_block'
+sil [ossa] @nohoist_into_access_scope_barred_by_barrier_block : $@convention(thin) (@owned C, @inout C) -> () {
+entry(%instance : @owned $C, %second : $*C):
+    %lifetime = begin_borrow [lexical] %instance : $C
+    %scope = begin_access [modify] [static] %second : $*C
+    cond_br undef, left, right
+
+left:
+    end_access %scope : $*C
+    %ignore = tuple ()
+    end_borrow %lifetime : $C
+    destroy_value %instance : $C
+    br exit
+
+right:
+    end_access %scope : $*C
+    apply undef(%lifetime) : $@convention(thin) (@guaranteed C) -> ()
+    end_borrow %lifetime : $C
+    destroy_value %instance : $C
+    br exit
+
+exit:
+    %retval = tuple ()
+    return %retval : $()
+}
+
+// =============================================================================
+// access scope tests                                                         }}
+// =============================================================================