Re-instate "StackPromotion: Ignore unreachable blocks in post-dominator tree."

With a bug fix which should ensure that it doesn't violate the stack nesting.

Original commit: 3d050f7b43

lib/SILOptimizer/Transforms/StackPromotion.cpp

@@ -17,6 +17,9 @@
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
+#include "swift/SIL/CFG.h"
+#include "llvm/Support/GenericDomTree.h"
+#include "llvm/Support/GenericDomTreeConstruction.h"
 #include "llvm/ADT/Statistic.h"
 
 STATISTIC(NumStackPromoted, "Number of objects promoted to the stack");
@@ -48,9 +51,23 @@ class StackPromoter {
   SILFunction *F;
   EscapeAnalysis::ConnectionGraph *ConGraph;
   DominanceInfo *DT;
-  PostDominanceInfo *PDT;
   EscapeAnalysis *EA;
 
+  // We use our own post-dominator tree instead of PostDominatorAnalysis,
+  // because we ignore unreachable blocks (actually all unreachable sub-graphs).
+  // Example:
+  //              |
+  //             bb1
+  //            /   \
+  //   unreachable  bb2
+  //                  \
+  //
+  // We want to get bb2 as immediate post-domiator of bb1. This is not the case
+  // with the regualar post-dominator tree.
+  llvm::DominatorTreeBase<SILBasicBlock> PostDomTree;
+
+  bool PostDomTreeValid;
+
   // Pseudo-functions for (de-)allocating array buffers on the stack.
 
   SILFunction *BufferAllocFunc = nullptr;
@@ -120,16 +137,31 @@ class StackPromoter {
                        SILInstruction *&AllocInsertionPoint,
                        SILInstruction *&DeallocInsertionPoint);
 
+  /// Returns the place where to insert the deallocation.
+  /// Returns null if this doesn't succeed or, in case \p RestartPoint is set,
+  /// a new iteration should be triggered.
+  SILInstruction *findDeallocPoint(SILInstruction *StartInst,
+                               SILInstruction *&RestartPoint,
+                               EscapeAnalysis::CGNode *Node,
+                               int NumUsePointsToFind);
+
   bool strictlyDominates(SILBasicBlock *A, SILBasicBlock *B) {
     return A != B && DT->dominates(A, B);
   }
 
   bool strictlyPostDominates(SILBasicBlock *A, SILBasicBlock *B) {
-    return A != B && PDT->dominates(A, B);
+    calculatePostDomTree();
+    return A != B && PostDomTree.dominates(A, B);
+  }
+
+  bool postDominates(SILBasicBlock *A, SILBasicBlock *B) {
+    calculatePostDomTree();
+    return PostDomTree.dominates(A, B);
   }
 
   SILBasicBlock *getImmediatePostDom(SILBasicBlock *BB) {
-    auto *Node = PDT->getNode(BB);
+    calculatePostDomTree();
+    auto *Node = PostDomTree.getNode(BB);
     if (!Node)
       return nullptr;
     auto *IDomNode = Node->getIDom();
@@ -138,12 +170,21 @@ class StackPromoter {
     return IDomNode->getBlock();
   }
   
+  void calculatePostDomTree() {
+    if (!PostDomTreeValid) {
+      // The StackPromoter acts as a "graph" for which the post-dominator-tree
+      // is calculated.
+      PostDomTree.recalculate(*this);
+      PostDomTreeValid = true;
+    }
+  }
+
 public:
 
   StackPromoter(SILFunction *F, EscapeAnalysis::ConnectionGraph *ConGraph,
-                DominanceInfo *DT, PostDominanceInfo *PDT,
+                DominanceInfo *DT, EscapeAnalysis *EA) :
-                EscapeAnalysis *EA) :
+    F(F), ConGraph(ConGraph), DT(DT), EA(EA), PostDomTree(true),
-    F(F), ConGraph(ConGraph), DT(DT), PDT(PDT), EA(EA) { }
+    PostDomTreeValid(false) { }
 
   /// What did the optimization change?
   enum class ChangeState {
@@ -152,6 +193,8 @@ public:
     Calls
   };
 
+  SILFunction *getFunction() const { return F; }
+
   /// The main entry point for the optimization.
   ChangeState promote();
 };
@@ -284,6 +327,87 @@ SILFunction *StackPromoter::getBufferDeallocFunc(SILFunction *OrigFunc,
   return BufferDeallocFunc;
 }
 
+namespace {
+
+/// Iterator which iterates over all basic blocks of a function which are not
+/// terminated by an unreachable inst.
+class NonUnreachableBlockIter :
+public std::iterator<std::forward_iterator_tag, SILBasicBlock, ptrdiff_t> {
+
+  SILFunction::iterator BaseIterator;
+  SILFunction::iterator End;
+
+  void skipUnreachables() {
+    while (true) {
+      if (BaseIterator == End)
+        return;
+      if (!isa<UnreachableInst>(BaseIterator->getTerminator()))
+        return;
+      BaseIterator++;
+    }
+  }
+
+public:
+  NonUnreachableBlockIter(SILFunction::iterator BaseIterator,
+                          SILFunction::iterator End) :
+      BaseIterator(BaseIterator), End(End) {
+    skipUnreachables();
+  }
+
+  NonUnreachableBlockIter() = default;
+  
+  SILBasicBlock &operator*() const { return *BaseIterator; }
+  SILBasicBlock &operator->() const { return *BaseIterator; }
+  
+  NonUnreachableBlockIter &operator++() {
+    BaseIterator++;
+    skipUnreachables();
+    return *this;
+  }
+  
+  NonUnreachableBlockIter operator++(int unused) {
+    NonUnreachableBlockIter Copy = *this;
+    ++*this;
+    return Copy;
+  }
+
+  friend bool operator==(NonUnreachableBlockIter lhs,
+                         NonUnreachableBlockIter rhs) {
+    return lhs.BaseIterator == rhs.BaseIterator;
+  }
+  friend bool operator!=(NonUnreachableBlockIter lhs,
+                         NonUnreachableBlockIter rhs) {
+    return !(lhs == rhs);
+  }
+};
+}
+
+namespace llvm {
+
+/// Use the StackPromoter as a wrapper for the function. It holds the list of
+/// basic blocks excluding all unreachable blocks.
+template <> struct GraphTraits<StackPromoter *>
+    : public GraphTraits<swift::SILBasicBlock*> {
+  typedef StackPromoter *GraphType;
+
+  static NodeType *getEntryNode(GraphType SP) {
+    return &SP->getFunction()->front();
+  }
+
+  typedef NonUnreachableBlockIter nodes_iterator;
+  static nodes_iterator nodes_begin(GraphType SP) {
+    return nodes_iterator(SP->getFunction()->begin(), SP->getFunction()->end());
+  }
+  static nodes_iterator nodes_end(GraphType SP) {
+    return nodes_iterator(SP->getFunction()->end(), SP->getFunction()->end());
+  }
+  static unsigned size(GraphType SP) {
+    return std::distance(nodes_begin(SP), nodes_end(SP));
+  }
+};
+
+}
+
 bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
                                     SILInstruction *&AllocInsertionPoint,
                                     SILInstruction *&DeallocInsertionPoint) {
@@ -311,13 +435,42 @@ bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
     return false;
   }
 
+  // Try to find the point where to insert the deallocation.
+  // This might need more than one try in case we need to move the allocation
+  // out of a stack-alloc-dealloc pair. See findDeallocPoint().
+  SILInstruction *StartInst = AI;
+  for (;;) {
+    SILInstruction *RestartPoint = nullptr;
+    DeallocInsertionPoint = findDeallocPoint(StartInst, RestartPoint, Node,
+                                             NumUsePointsToFind);
+    if (DeallocInsertionPoint)
+      return true;
+
+    if (!RestartPoint)
+      return false;
+
+    // Moving a buffer allocation call is not trivial because we would need to
+    // move all the parameter calculations as well. So we just don't do it.
+    if (!isa<AllocRefInst>(AI))
+      return false;
+
+    // Retry with moving the allocation up.
+    AllocInsertionPoint = RestartPoint;
+    StartInst = RestartPoint;
+  }
+}
+
+SILInstruction *StackPromoter::findDeallocPoint(SILInstruction *StartInst,
+                                                SILInstruction *&RestartPoint,
+                                                EscapeAnalysis::CGNode *Node,
+                                                int NumUsePointsToFind) {
   // In the following we check two requirements for stack promotion:
   // 1) Are all uses in the same control region as the alloc? E.g. if the
   //    allocation is in a loop then there may not be any uses of the object
   //    outside the loop.
   // 2) We need to find an insertion place for the deallocation so that it
   //    preserves a properly nested stack allocation-deallocation structure.
-  SILBasicBlock *StartBlock = AI->getParent();
+  SILBasicBlock *StartBlock = StartInst->getParent();
 
   // The block where we assume we can insert the deallocation.
   SILBasicBlock *EndBlock = StartBlock;
@@ -336,7 +489,7 @@ bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
     if (BB == StartBlock) {
       // In the first block we start at the allocation instruction and not at
       // the begin of the block.
-      Iter = AI->getIterator();
+      Iter = StartInst->getIterator();
     } else {
       // Track all uses in the block arguments.
       for (SILArgument *BBArg : BB->getBBArgs()) {
@@ -359,7 +512,7 @@ bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
         while (!strictlyPostDominates(EndBlock, Pred)) {
           EndBlock = getImmediatePostDom(EndBlock);
           if (!EndBlock)
-            return false;
+            return nullptr;
         }
       }
       Iter = BB->begin();
@@ -370,8 +523,7 @@ bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
       SILInstruction &I = *Iter++;
       if (BB == EndBlock && StackDepth == 0 && NumUsePointsToFind == 0) {
         // We found a place to insert the stack deallocation.
-        DeallocInsertionPoint = &I;
+        return &I;
-        return true;
       }
       if (I.isAllocatingStack()) {
         StackDepth++;
@@ -388,16 +540,17 @@ bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
           //
           // In this case we can move the alloc_ref before the alloc_stack
           // to fix the nesting.
-          if (!isa<AllocRefInst>(AI))
-            return false;
           auto *Alloc = dyn_cast<SILInstruction>(I.getOperand(0));
           if (!Alloc)
-            return false;
+            return nullptr;
+
           // This should always be the case, but let's be on the safe side.
-          if (!PDT->dominates(StartBlock, Alloc->getParent()))
+          if (!postDominates(StartBlock, Alloc->getParent()))
-            return false;
+            return nullptr;
-          AllocInsertionPoint = Alloc;
+
-          StackDepth++;
+          // Trigger another iteration with a new start point;
+          RestartPoint = Alloc;
+          return nullptr;
         }
         StackDepth--;
       }
@@ -422,7 +575,7 @@ bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
         //     dealloc_stack %1 // this is the new EndBlock
         EndBlock = getImmediatePostDom(EndBlock);
         if (!EndBlock)
-          return false;
+          return nullptr;
       }
       // Again, it's important that the EndBlock is the first in the WorkList.
       WorkList.insert(EndBlock, -1);
@@ -441,7 +594,7 @@ bool StackPromoter::canPromoteAlloc(SILInstruction *AI,
         //     cond_br ..., loop, exit
         //   exit:
         //     use(%container)
-        return false;
+        return nullptr;
       }
       WorkList.insert(Succ, StackDepth);
     }
@@ -466,20 +619,18 @@ private:
 
     auto *EA = PM->getAnalysis<EscapeAnalysis>();
     auto *DA = PM->getAnalysis<DominanceAnalysis>();
-    auto *PDA = PM->getAnalysis<PostDominanceAnalysis>();
 
     SILFunction *F = getFunction();
     if (auto *ConGraph = EA->getConnectionGraph(F)) {
-      StackPromoter promoter(F, ConGraph, DA->get(F), PDA->get(F), EA);
+      StackPromoter promoter(F, ConGraph, DA->get(F), EA);
       switch (promoter.promote()) {
         case StackPromoter::ChangeState::None:
           break;
         case StackPromoter::ChangeState::Insts:
           invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
           break;
-        case StackPromoter::ChangeState::Calls: {
+        case StackPromoter::ChangeState::Calls:
           invalidateAnalysis(SILAnalysis::InvalidationKind::CallsAndInstructions);
-        }
           break;
       }
     }

test/SILOptimizer/stack_promotion.sil

@@ -88,6 +88,49 @@ bb0:
   return %t : $()
 }
 
+// CHECK-LABEL: sil @promote_with_unreachable_block
+// CHECK: [[O:%[0-9]+]] = alloc_ref [stack] $XX
+// CHECK: bb1:
+// CHECK-NEXT: unreachable
+// CHECK: bb2:
+// CHECK: strong_release
+// CHECK: dealloc_ref [stack] [[O]] : $XX
+// CHECK: return
+sil @promote_with_unreachable_block : $@convention(thin) () -> Int32 {
+bb0:
+  %o1 = alloc_ref $XX
+  %f1 = function_ref @xx_init : $@convention(thin) (@guaranteed XX) -> XX
+  %n1 = apply %f1(%o1) : $@convention(thin) (@guaranteed XX) -> XX
+  cond_br undef, bb1, bb2
+
+bb1:
+  unreachable
+
+bb2:
+  %l1 = ref_element_addr %n1 : $XX, #XX.x
+  %l2 = load %l1 : $*Int32
+  strong_release %n1 : $XX
+  return %l2 : $Int32
+}
+
+// CHECK-LABEL: sil @no_return_function
+// Just check that we don't crash on this.
+// It's a corner case, so we don't care if stack promotion is done or not.
+// CHECK: unreachable
+sil @no_return_function : $@convention(thin) () -> Int32 {
+bb0:
+  %o1 = alloc_ref $XX
+  %f1 = function_ref @xx_init : $@convention(thin) (@guaranteed XX) -> XX
+  %n1 = apply %f1(%o1) : $@convention(thin) (@guaranteed XX) -> XX
+  br bb1
+
+bb1:
+  %l1 = ref_element_addr %n1 : $XX, #XX.x
+  %l2 = load %l1 : $*Int32
+  strong_release %n1 : $XX
+  unreachable
+}
+
 // CHECK-LABEL: sil @promote_in_loop_with_if
 // CHECK: [[O:%[0-9]+]] = alloc_ref [stack] $XX
 // CHECK: {{^}}bb4({{.*}}):
@@ -247,7 +290,7 @@ bb5:
   return %a1 : $Int32
 }
 
-// CHECK-LABEL: sil @promote_and_move_alloc_before_alloc_stack
+// CHECK-LABEL: sil @promote_and_move_alloc_before_alloc_stack1
 // CHECK: [[O:%[0-9]+]] = alloc_ref [stack] $XX
 // CHECK: alloc_stack
 // CHECK: {{^}}bb2:
@@ -255,7 +298,7 @@ bb5:
 // CHECK: strong_release
 // CHECK: dealloc_ref [stack] [[O]] : $XX
 // CHECK: return
-sil @promote_and_move_alloc_before_alloc_stack : $@convention(thin) () -> Int32 {
+sil @promote_and_move_alloc_before_alloc_stack1 : $@convention(thin) () -> Int32 {
 bb0:
   %s1 = alloc_stack $Int32
   cond_br undef, bb1, bb2
@@ -274,6 +317,36 @@ bb2:
   return %l2 : $Int32
 }
 
+// CHECK-LABEL: sil @promote_and_move_alloc_before_alloc_stack2
+// CHECK: alloc_ref [stack] $XX
+// CHECK-NEXT: alloc_stack
+// CHECK: {{^}}bb3:
+// CHECK: strong_release
+// CHECK-NEXT: dealloc_ref [stack]
+// CHECK-NEXT: return
+sil @promote_and_move_alloc_before_alloc_stack2 : $@convention(thin) () -> Int32 {
+bb0:
+  %s1 = alloc_stack $Int32
+  %o1 = alloc_ref $XX
+  %f1 = function_ref @xx_init : $@convention(thin) (@guaranteed XX) -> XX
+  %n1 = apply %f1(%o1) : $@convention(thin) (@guaranteed XX) -> XX
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb3
+ 
+bb2:
+  dealloc_stack %s1 : $*Int32
+  unreachable
+
+bb3:
+  dealloc_stack %s1 : $*Int32
+  %l1 = ref_element_addr %n1 : $XX, #XX.x
+  %l2 = load %l1 : $*Int32
+  strong_release %n1 : $XX
+  return %l2 : $Int32
+}
+
 // CHECK-LABEL: sil @promote_array
 // CHECK: [[AF:%[0-9]+]] = function_ref @swift_bufferAllocateOnStack : $@convention(thin) (@thick AnyObject.Type, Int, Int) -> @owned AnyObject
 // CHECK: [[B:%[0-9]+]] = apply [[AF]](