[ShrinkBorrowScope] Adopt iterative dataflow.

Adopt IterativeBackwardReachability and VisitBarrierAccessScopes.
Enables hoisting borrow scopes over loops and avoiding hoisting borrow
scopes into unrelated access scopes, respectively.

rdar://93186505
rdar://93060369

Author: nate-chandler

include/swift/SILOptimizer/Analysis/Reachability.h

@@ -13,12 +13,6 @@
 /// Reachability data flow analysis using a path-discovery worklist. For
 /// efficient data flow propagation based on a single SSA value and its uses.
 ///
-/// TODO: Add an optimistic data flow for more aggresive optimization:
-/// - Add another set for blocks reachable by barriers
-/// - Change the meet operation to a union
-/// - Propagate past barriers up to the SSA def
-/// - Iterate to a fix-point.
-///
 //===----------------------------------------------------------------------===//
 
 #ifndef SWIFT_SILOPTIMIZER_ANALYSIS_REACHABILITY_H

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,94 @@ 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;
+  SmallSetVector<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;
+  SmallSetVector<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.
+  SmallSetVector<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.
+/// Implements IterativeBackwardReachability::Effects.
 class DataFlow final {
+public:
+  using Reachability = IterativeBackwardReachability<DataFlow>;
+  using Effect = Reachability::Effect;
+  using VBAS = VisitBarrierAccessScopes<DataFlow, BarrierAccessScopeFinder>;
+
+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(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 class IterativeBackwardReachability<DataFlow>;
+  friend class BarrierAccessScopeFinder;
+  friend VBAS;
 
-  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);
-  }
+  ArrayRef<SILInstruction *> gens();
 
-  Classification classifyInstruction(SILInstruction *);
+  Optional<Effect> effectForInstruction(SILInstruction *);
 
-  bool classificationIsBarrier(Classification);
+  Optional<Effect> effectForPhi(SILBasicBlock *);
+};
+
+class BarrierAccessScopeFinder final {
+  using Impl = VisitBarrierAccessScopes<DataFlow, BarrierAccessScopeFinder>;
+  Context const &context;
+  Impl impl;
+  DataFlow &dataflow;
+  Optional<SmallVector<SILBasicBlock *, 16>> cachedRoots;
 
-  void visitedInstruction(SILInstruction *, Classification);
+public:
+  BarrierAccessScopeFinder(Context const &context, DataFlow &dataflow)
+      : context(context), impl(&context.function, dataflow, *this),
+        dataflow(dataflow) {}
 
-  bool checkReachableBarrier(SILInstruction *);
+  void find();
 
-  bool checkReachablePhiBarrier(SILBasicBlock *);
+private:
+  friend Impl;
+
+  ArrayRef<SILBasicBlock *> roots();
+  bool isInRegion(SILBasicBlock *);
+  void visitBarrierAccessScope(BeginAccessInst *);
 };
 
 /// Whether the specified value is %lifetime or its iterated copy_value.
@@ -207,6 +224,16 @@ 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,
+                            barriers.blocks);
+}
+
 DataFlow::Classification
 DataFlow::classifyInstruction(SILInstruction *instruction) {
   if (instruction == &context.introducer) {
@@ -220,6 +247,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;
   }
@@ -237,29 +269,21 @@ 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?!");
-}
+ArrayRef<SILInstruction *> DataFlow::gens() { return uses.ends; }
 
-bool DataFlow::checkReachableBarrier(SILInstruction *instruction) {
+Optional<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)
+             ? Optional<Effect>{Effect::Kill}
+             : llvm::None;
 }
 
-bool DataFlow::checkReachablePhiBarrier(SILBasicBlock *block) {
+Optional<DataFlow::Effect> DataFlow::effectForPhi(SILBasicBlock *block) {
   assert(llvm::all_of(block->getArguments(),
                       [&](auto argument) { return PhiValue(argument); }));
 
@@ -268,10 +292,34 @@ bool DataFlow::checkReachablePhiBarrier(SILBasicBlock *block) {
         return classificationIsBarrier(
             classifyInstruction(predecessor->getTerminator()));
       });
-  if (isBarrier) {
-    result.phiBarriers.push_back(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);
   }
-  return isBarrier;
 }
 
 /// Hoist the scope ends of %lifetime, rewriting copies and borrows along the
@@ -311,7 +359,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 +372,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 +400,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 +419,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;
     }

test/SILOptimizer/shrink_borrow_scope.sil

@@ -458,12 +458,13 @@ exit:
 // =============================================================================
 
 // Don't hoist over loop without uses.
-// TODO: Eventually, we should hoist over such loops.
+//
 // CHECK-LABEL: sil [ossa] @hoist_over_loop_1 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C):
 // CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [[INSTANCE]]
 // CHECK:         [[CALLEE_GUARANTEED:%[^,]+]] = function_ref @callee_guaranteed
 // CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[LIFETIME]])
+// CHECK:         end_borrow [[LIFETIME]]
 // CHECK:         br [[LOOP_HEADER:bb[0-9]+]]
 // CHECK:       [[LOOP_HEADER]]:
 // CHECK:         br [[LOOP_BODY:bb[0-9]+]]
@@ -474,7 +475,6 @@ exit:
 // CHECK:       [[LOOP_BACKEDGE]]:
 // CHECK:         br [[LOOP_HEADER]]
 // CHECK:       [[EXIT]]:
-// CHECK:         end_borrow [[LIFETIME]]
 // CHECK:         return [[INSTANCE]]
 // CHECK-LABEL: } // end sil function 'hoist_over_loop_1'
 sil [ossa] @hoist_over_loop_1 : $@convention(thin) (@owned C) -> @owned C {