[ShrinkBorrowScope] Adopt iterative dataflow.

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"
@@ -49,6 +50,8 @@ struct Context final {
   /// introducer->getOperand()
   SILValue const borrowee;
 
+  SILBasicBlock *defBlock;
+
   SILFunction &function;
 
   /// The copy_value instructions that the utility creates or changes.
@@ -62,7 +65,8 @@ struct Context final {
           SmallVectorImpl<CopyValueInst *> &modifiedCopyValueInsts,
           InstructionDeleter &deleter)
       : introducer(introducer), borrowedValue(BorrowedValue(&introducer)),
-        borrowee(introducer.getOperand()), function(*introducer.getFunction()),
+        borrowee(introducer.getOperand()), defBlock(introducer.getParent()),
+        function(*introducer.getFunction()),
         modifiedCopyValueInsts(modifiedCopyValueInsts), deleter(deleter) {}
   Context(Context const &) = delete;
   Context &operator=(Context const &) = delete;
@@ -75,7 +79,7 @@ struct Usage final {
   SmallPtrSet<SILInstruction *, 16> users;
   // The instructions from which the shrinking starts, the scope ending
   // instructions.
-  llvm::SmallSetVector<SILInstruction *, 4> ends;
+  llvm::SmallVector<SILInstruction *, 4> ends;
 
   Usage(){};
   Usage(Usage const &) = delete;
@@ -95,7 +99,7 @@ bool findUsage(Context const &context, Usage &usage) {
     // If a scope ending instruction is not an end_borrow, bail out.
     if (!isa<EndBorrowInst>(instruction))
       return false;
-    usage.ends.insert(instruction);
+    usage.ends.push_back(instruction);
   }
 
   SmallVector<Operand *, 16> uses;
@@ -112,81 +116,96 @@ bool findUsage(Context const &context, Usage &usage) {
 
 /// How end_borrow hoisting is obstructed.
 struct DeinitBarriers final {
-  /// Blocks up to "before the beginning" of which hoisting was able to proceed.
-  BasicBlockSetVector hoistingReachesBeginBlocks;
-
-  /// Blocks to "after the end" of which hoisting was able to proceed.
-  BasicBlockSet hoistingReachesEndBlocks;
-
   /// Copies to be rewritten as copies of %borrowee.
   SmallVector<CopyValueInst *, 4> copies;
 
   /// Instructions above which end_borrows cannot be hoisted.
-  SmallVector<SILInstruction *, 4> barriers;
+  SmallVector<SILInstruction *, 4> instructions;
 
   /// Blocks one of whose phis is a barrier and consequently out of which
   /// end_borrows cannot be hoisted.
-  SmallVector<SILBasicBlock *, 4> phiBarriers;
+  SmallVector<SILBasicBlock *, 4> phis;
 
-  DeinitBarriers(Context &context)
-      : hoistingReachesBeginBlocks(&context.function),
-        hoistingReachesEndBlocks(&context.function) {}
+  /// Blocks whose single predecessors has another successor to the top of which
+  /// end_borrows cannot be hoisted.
+  SmallVector<SILBasicBlock *, 4> blocks;
+
+  DeinitBarriers(Context &context) {}
   DeinitBarriers(DeinitBarriers const &) = delete;
   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 BackwardReachability::BlockReachability.
-class DataFlow final {
+/// Implements IterativeBackwardReachability::Effects.
+/// Implements IterativeBackwardReachability::findBarrier::Visitor.
+/// Implements VisitBarrierAccessScopes::Effects
+class Dataflow final {
+public:
+  using Reachability = IterativeBackwardReachability<Dataflow>;
+  using Effect = Reachability::Effect;
+
+private:
   Context const &context;
   Usage const &uses;
-  DeinitBarriers &result;
+  DeinitBarriers &barriers;
+  Reachability::Result result;
+  Reachability reachability;
+  SmallPtrSet<BeginAccessInst *, 8> barrierAccessScopes;
+  bool recordCopies = false;
 
   enum class Classification { Barrier, Copy, Other };
 
-  BackwardReachability<DataFlow> reachability;
-
 public:
-  DataFlow(Context const &context, Usage const &uses, DeinitBarriers &result)
-      : context(context), uses(uses), result(result),
-        reachability(&context.function, *this) {
-    // Seed reachability with the scope ending uses from which the backwards
-    // data flow will begin.
-    for (auto *end : uses.ends) {
-      reachability.initLastUse(end);
-    }
-  }
-  DataFlow(DataFlow const &) = delete;
-  DataFlow &operator=(DataFlow const &) = delete;
+  Dataflow(Context const &context, Usage const &uses, DeinitBarriers &barriers)
+      : context(context), uses(uses), barriers(barriers),
+        result(&context.function),
+        reachability(&context.function, context.defBlock, *this, result) {}
+  Dataflow(Dataflow const &) = delete;
+  Dataflow &operator=(Dataflow const &) = delete;
 
-  void run() { reachability.solveBackward(); }
+  void run();
 
 private:
-  friend class BackwardReachability<DataFlow>;
+  friend Reachability;
+  friend class BarrierAccessScopeFinder;
+  friend class VisitBarrierAccessScopes<Dataflow, BarrierAccessScopeFinder>;
 
-  bool hasReachableBegin(SILBasicBlock *block) {
-    return result.hoistingReachesBeginBlocks.contains(block);
-  }
+  Classification classifyInstruction(SILInstruction *);
 
-  void markReachableBegin(SILBasicBlock *block) {
-    result.hoistingReachesBeginBlocks.insert(block);
-  }
+  bool classificationIsBarrier(Classification);
 
-  void markReachableEnd(SILBasicBlock *block) {
-    result.hoistingReachesEndBlocks.insert(block);
-  }
+  /// IterativeBackwardReachability::Effects
+  /// VisitBarrierAccessScopes::Effects
 
-  Classification classifyInstruction(SILInstruction *);
+  ArrayRef<SILInstruction *> gens() { return uses.ends; }
 
-  bool classificationIsBarrier(Classification);
+  Effect effectForInstruction(SILInstruction *);
+
+  Effect effectForPhi(SILBasicBlock *);
 
-  void visitedInstruction(SILInstruction *, Classification);
+  /// VisitBarrierAccessScopes::Effects
 
-  bool checkReachableBarrier(SILInstruction *);
+  auto localGens() { return result.localGens; }
 
-  bool checkReachablePhiBarrier(SILBasicBlock *);
+  bool isLocalGen(SILInstruction *instruction) {
+    return result.localGens.contains(instruction);
+  }
+
+  /// IterativeBackwardReachability::findBarrier::Visitor.
+
+  void visitBarrierInstruction(SILInstruction *instruction) {
+    barriers.instructions.push_back(instruction);
+  }
+
+  void visitBarrierPhi(SILBasicBlock *block) { barriers.phis.push_back(block); }
+
+  void visitBarrierBlock(SILBasicBlock *block) {
+    barriers.blocks.push_back(block);
+  }
 };
 
 /// Whether the specified value is %lifetime or its iterated copy_value.
@@ -207,8 +226,8 @@ bool isSimpleExtendedIntroducerDef(Context const &context, SILValue value) {
   }
 }
 
-DataFlow::Classification
-DataFlow::classifyInstruction(SILInstruction *instruction) {
+Dataflow::Classification
+Dataflow::classifyInstruction(SILInstruction *instruction) {
   if (instruction == &context.introducer) {
     return Classification::Barrier;
   }
@@ -220,13 +239,18 @@ 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;
   }
   return Classification::Other;
 }
 
-bool DataFlow::classificationIsBarrier(Classification classification) {
+bool Dataflow::classificationIsBarrier(Classification classification) {
   switch (classification) {
   case Classification::Barrier:
     return true;
@@ -237,29 +261,17 @@ bool DataFlow::classificationIsBarrier(Classification classification) {
   llvm_unreachable("exhaustive switch not exhaustive?!");
 }
 
-void DataFlow::visitedInstruction(SILInstruction *instruction,
-                                  Classification classification) {
-  assert(classifyInstruction(instruction) == classification);
-  switch (classification) {
-  case Classification::Barrier:
-    result.barriers.push_back(instruction);
-    return;
-  case Classification::Copy:
-    result.copies.push_back(cast<CopyValueInst>(instruction));
-    return;
-  case Classification::Other:
-    return;
-  }
-  llvm_unreachable("exhaustive switch not exhaustive?!");
-}
-
-bool DataFlow::checkReachableBarrier(SILInstruction *instruction) {
+Dataflow::Effect Dataflow::effectForInstruction(SILInstruction *instruction) {
+  if (llvm::find(uses.ends, instruction) != uses.ends.end())
+    return Effect::Gen();
   auto classification = classifyInstruction(instruction);
-  visitedInstruction(instruction, classification);
-  return classificationIsBarrier(classification);
+  if (recordCopies && classification == Classification::Copy)
+    barriers.copies.push_back(cast<CopyValueInst>(instruction));
+  return classificationIsBarrier(classification) ? Effect::Kill()
+                                                 : Effect::NoEffect();
 }
 
-bool DataFlow::checkReachablePhiBarrier(SILBasicBlock *block) {
+Dataflow::Effect Dataflow::effectForPhi(SILBasicBlock *block) {
   assert(llvm::all_of(block->getArguments(),
                       [&](auto argument) { return PhiValue(argument); }));
 
@@ -268,10 +280,49 @@ bool DataFlow::checkReachablePhiBarrier(SILBasicBlock *block) {
         return classificationIsBarrier(
             classifyInstruction(predecessor->getTerminator()));
       });
-  if (isBarrier) {
-    result.phiBarriers.push_back(block);
+  return isBarrier ? Effect::Kill() : Effect::NoEffect();
+}
+
+/// 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;
+
+public:
+  BarrierAccessScopeFinder(Context const &context, Dataflow &dataflow)
+      : context(context), impl(&context.function, dataflow, *this),
+        dataflow(dataflow) {}
+
+  void find() { impl.visit(); }
+
+private:
+  friend Impl;
+
+  bool isInRegion(SILBasicBlock *block) {
+    return dataflow.result.discoveredBlocks.contains(block);
   }
-  return isBarrier;
+
+  void visitBarrierAccessScope(BeginAccessInst *bai) {
+    dataflow.barrierAccessScopes.insert(bai);
+    for (auto *eai : bai->getEndAccesses()) {
+      dataflow.reachability.addKill(eai);
+    }
+  }
+};
+
+void Dataflow::run() {
+  reachability.initialize();
+  BarrierAccessScopeFinder finder(context, *this);
+  finder.find();
+  reachability.solve();
+  recordCopies = true;
+  reachability.findBarriers(*this);
 }
 
 /// Hoist the scope ends of %lifetime, rewriting copies and borrows along the
@@ -311,7 +362,7 @@ bool Rewriter::run() {
   // A block is a phi barrier iff any of its predecessors' terminators get
   // classified as barriers.  That happens when a copy of %lifetime is passed
   // to a phi.
-  for (auto *block : barriers.phiBarriers) {
+  for (auto *block : barriers.phis) {
     madeChange |= createEndBorrow(&block->front());
   }
 
@@ -324,15 +375,11 @@ bool Rewriter::run() {
   // of a block P's successors B had reachable beginnings.  If any of them
   // didn't, then BackwardReachability::meetOverSuccessors would never have
   // returned true for P, so none of its instructions would ever have been
-  // classified (except for via checkReachablePhiBarrier, which doesn't record
-  // terminator barriers).
-  for (auto instruction : barriers.barriers) {
+  // classified (except for via effectForPhi, which doesn't record terminator
+  // barriers).
+  for (auto instruction : barriers.instructions) {
     if (auto *terminator = dyn_cast<TermInst>(instruction)) {
       auto successors = terminator->getParentBlock()->getSuccessorBlocks();
-      // In order for the instruction to have been classified as a barrier,
-      // reachability would have had to reach the block containing it.
-      assert(barriers.hoistingReachesEndBlocks.contains(
-                terminator->getParentBlock()));
       for (auto *successor : successors) {
         madeChange |= createEndBorrow(&successor->front());
       }
@@ -356,12 +403,8 @@ bool Rewriter::run() {
   // P not having a reachable end--see BackwardReachability::meetOverSuccessors.
   //
   // control-flow-boundary(B) := beginning-reachable(B) && !end-reachable(P)
-  for (auto *block : barriers.hoistingReachesBeginBlocks) {
-    if (auto *predecessor = block->getSinglePredecessorBlock()) {
-      if (!barriers.hoistingReachesEndBlocks.contains(predecessor)) {
-        madeChange |= createEndBorrow(&block->front());
-      }
-    }
+  for (auto *block : barriers.blocks) {
+    madeChange |= createEndBorrow(&block->front());
   }
 
   if (madeChange) {
@@ -379,7 +422,7 @@ bool Rewriter::run() {
 
 bool Rewriter::createEndBorrow(SILInstruction *insertionPoint) {
   if (auto *ebi = dyn_cast<EndBorrowInst>(insertionPoint)) {
-    if (uses.ends.contains(insertionPoint)) {
+    if (llvm::find(uses.ends, insertionPoint) != uses.ends.end()) {
       reusedEndBorrowInsts.insert(insertionPoint);
       return false;
     }
@@ -397,7 +440,7 @@ bool run(Context &context) {
     return false;
 
   DeinitBarriers barriers(context);
-  DataFlow flow(context, usage, barriers);
+  Dataflow flow(context, usage, barriers);
   flow.run();
 
   Rewriter rewriter(context, usage, barriers);