Merge pull request #22120 from eeckstein/fix-stacknesting

SILOptimizer: correctly handle unreachable blocks in StackNesting.

Author: eeckstein

include/swift/SILOptimizer/Utils/StackNesting.h

@@ -84,12 +84,8 @@ class StackNesting {
     /// The bit-set of alive stack locations at the block entry.
     BitVector AliveStackLocsAtEntry;
 
-    /// True if there is a path from this block to a function-exit.
-    ///
-    /// In other words: this block does not end in an unreachable-instruction.
-    /// This flag is only used for verifying that the lifetime of a stack
-    /// location does not end at the end of a block.
-    bool ExitReachable = false;
+    /// The bit-set of alive stack locations at the block exit.
+    BitVector AliveStackLocsAtExit;
 
     BlockInfo(SILBasicBlock *Block) : Block(Block) { }
   };
@@ -165,6 +161,19 @@ class StackNesting {
                       SILInstruction *InsertionPoint,
                       Optional<SILLocation> Location);
 
+  /// Reeturns the location bit number for a stack allocation instruction.
+  int bitNumberForAlloc(SILInstruction *AllocInst) {
+    assert(AllocInst->isAllocatingStack());
+    return StackLoc2BitNumbers[cast<SingleValueInstruction>(AllocInst)];
+  }
+
+  /// Reeturns the location bit number for a stack deallocation instruction.
+  int bitNumberForDealloc(SILInstruction *DeallocInst) {
+    assert(DeallocInst->isDeallocatingStack());
+    auto *AllocInst = cast<SingleValueInstruction>(DeallocInst->getOperand(0));
+    return bitNumberForAlloc(AllocInst);
+  }
+
   /// Modifies the SIL to end up with a correct stack nesting.
   ///
   /// Returns the status of what changes were made.

lib/SILOptimizer/Utils/StackNesting.cpp

@@ -134,9 +134,6 @@ void StackNesting::setup(SILFunction *F) {
         }
       }
     }
-    if (BI->Block->getTerminator()->isFunctionExiting())
-      BI->ExitReachable = true;
-
     for (auto *SuccBB : BI->Block->getSuccessorBlocks()) {
       BlockInfo *&SuccBI = BlockMapping[SuccBB];
       if (!SuccBI) {
@@ -165,33 +162,78 @@ bool StackNesting::solve() {
   bool isNested = false;
   BitVector Bits(StackLocs.size());
 
-  // Iterate until we reach a fixed-point.
+  // Initialize all bit fields to 1s, expect 0s for the entry block.
+  bool initVal = false;
+  for (BlockInfo &BI : BlockInfos) {
+    BI.AliveStackLocsAtEntry.resize(StackLocs.size(), initVal);
+    initVal = false;
+  }
+
+  // First step: do a forward dataflow analysis to get the live stack locations
+  // at the block exits.
+  // This is necessary to get the live locations at blocks which end in
+  // unreachable instructions (otherwise the backward data flow would be
+  // sufficient). The special thing about unreachable-blocks is that it's
+  // okay to have alive locations at that point, i.e. locations which are never
+  // dealloced. We cannot get such locations with a purly backward dataflow.
+  do {
+    changed = false;
+
+    for (BlockInfo &BI : BlockInfos) {
+      Bits = BI.AliveStackLocsAtEntry;
+      for (SILInstruction *StackInst : BI.StackInsts) {
+        if (StackInst->isAllocatingStack()) {
+          Bits.set(bitNumberForAlloc(StackInst));
+        } else if (StackInst->isDeallocatingStack()) {
+          Bits.reset(bitNumberForDealloc(StackInst));
+        }
+      }
+      if (Bits != BI.AliveStackLocsAtExit) {
+        BI.AliveStackLocsAtExit = Bits;
+        assert(!(BI.Block->getTerminator()->isFunctionExiting() && Bits.any())
+               && "stack location is missing dealloc");
+        changed = true;
+      }
+      // Merge the bits into the successors.
+      for (BlockInfo *SuccBI : BI.Successors) {
+        SuccBI->AliveStackLocsAtEntry &= Bits;
+      }
+    }
+  } while (changed);
+
+  // Second step: do a backward dataflow analysis to extend the lifetimes of
+  // no properly nested allocations.
   do {
     changed = false;
 
     // It's a backward dataflow problem.
     for (BlockInfo &BI : reversed(BlockInfos)) {
       // Collect the alive-bits (at the block exit) from the successor blocks.
-      Bits.reset();
       for (BlockInfo *SuccBI : BI.Successors) {
-        Bits |= SuccBI->AliveStackLocsAtEntry;
-
-        // Also get the ExitReachable flag from the successor blocks.
-        if (!BI.ExitReachable && SuccBI->ExitReachable) {
-          BI.ExitReachable = true;
-          changed = true;
+        BI.AliveStackLocsAtExit |= SuccBI->AliveStackLocsAtEntry;
+      }
+      Bits = BI.AliveStackLocsAtExit;
+
+      if (isa<UnreachableInst>(BI.Block->getTerminator())) {
+        // We treat unreachable as an implicit deallocation for all locations
+        // which are still alive at this point.
+        for (int BitNr = Bits.find_first(); BitNr >= 0;
+             BitNr = Bits.find_next(BitNr)) {
+          // For each alive location extend the lifetime of all locations which
+          // are alive at the allocation point. This is the same as we do for
+          // a "real" deallocation instruction (see below).
+          Bits |= StackLocs[BitNr].AliveLocs;
         }
+        BI.AliveStackLocsAtExit = Bits;
       }
       for (SILInstruction *StackInst : reversed(BI.StackInsts)) {
         if (StackInst->isAllocatingStack()) {
-          auto AllocInst = cast<SingleValueInstruction>(StackInst);
-          int BitNr = StackLoc2BitNumbers[AllocInst];
+          int BitNr = bitNumberForAlloc(StackInst);
           if (Bits != StackLocs[BitNr].AliveLocs) {
             // More locations are alive around the StackInst's location.
             // Update the AlivaLocs bitset, which contains all those alive
             // locations.
-            assert((Bits.test(BitNr) || (!BI.ExitReachable && !Bits.any()))
-                   && "no dealloc found for alloc stack");
+            assert(Bits.test(BitNr) && "no dealloc found for alloc stack");
             StackLocs[BitNr].AliveLocs = Bits;
             changed = true;
             isNested = true;
@@ -203,10 +245,7 @@ bool StackNesting::solve() {
           // A stack deallocation begins the lifetime of its location (in
           // reverse order). And it also begins the lifetime of all other
           // locations which are alive at the allocation point.
-          auto *AllocInst =
-            cast<SingleValueInstruction>(StackInst->getOperand(0));
-          int BitNr = StackLoc2BitNumbers[AllocInst];
-          Bits |= StackLocs[BitNr].AliveLocs;
+          Bits |= StackLocs[bitNumberForDealloc(StackInst)].AliveLocs;
         }
       }
       if (Bits != BI.AliveStackLocsAtEntry) {
@@ -271,10 +310,7 @@ StackNesting::Changes StackNesting::adaptDeallocs() {
   // the same order as in solve().
   for (const BlockInfo &BI : reversed(BlockInfos)) {
     // Collect the alive-bits (at the block exit) from the successor blocks.
-    Bits.reset();
-    for (BlockInfo *SuccBI : BI.Successors) {
-      Bits |= SuccBI->AliveStackLocsAtEntry;
-    }
+    Bits = BI.AliveStackLocsAtExit;
 
     // Insert deallocations at block boundaries.
     // This can be necessary for unreachable blocks. Example:
@@ -293,12 +329,6 @@ StackNesting::Changes StackNesting::adaptDeallocs() {
          SuccIdx < NumSuccs; ++ SuccIdx) {
       BlockInfo *SuccBI = BI.Successors[SuccIdx];
 
-      // It's acceptable to not deallocate alive locations in unreachable
-      // blocks - as long as the nesting is not violated. So if there are no
-      // alive locations at the unreachable successor block, we can ignore it.
-      if (!SuccBI->ExitReachable && !SuccBI->AliveStackLocsAtEntry.any())
-        continue;
-
       if (SuccBI->AliveStackLocsAtEntry == Bits)
         continue;
 
@@ -320,16 +350,14 @@ StackNesting::Changes StackNesting::adaptDeallocs() {
     for (SILInstruction *StackInst : reversed(BI.StackInsts)) {
       if (StackInst->isAllocatingStack()) {
         // For allocations we just update the bit-set.
-        auto AllocInst = cast<SingleValueInstruction>(StackInst);
-        int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
+        int BitNr = bitNumberForAlloc(StackInst);
         assert(Bits == StackLocs[BitNr].AliveLocs &&
                "dataflow didn't converge");
         Bits.reset(BitNr);
       } else if (StackInst->isDeallocatingStack()) {
         // Handle deallocations.
-        auto *AllocInst = cast<SingleValueInstruction>(StackInst->getOperand(0));
         SILLocation Loc = StackInst->getLoc();
-        int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
+        int BitNr = bitNumberForDealloc(StackInst);
         SILInstruction *InsertionPoint = &*std::next(StackInst->getIterator());
         if (Bits.test(BitNr)) {
           // The location of StackInst is alive after StackInst. So we have to
@@ -372,17 +400,18 @@ void StackNesting::dump() const {
       continue;
 
     llvm::dbgs() << "Block " << BI.Block->getDebugID();
-    if (!BI.ExitReachable)
-      llvm::dbgs() << " (unreachable exit)";
-    llvm::dbgs() << ": bits=";
+    llvm::dbgs() << ": entry-bits=";
     dumpBits(BI.AliveStackLocsAtEntry);
+    llvm::dbgs() << ": exit-bits=";
+    dumpBits(BI.AliveStackLocsAtExit);
+    llvm::dbgs() << '\n';
     for (SILInstruction *StackInst : BI.StackInsts) {
       if (StackInst->isAllocatingStack()) {
         auto AllocInst = cast<SingleValueInstruction>(StackInst);
         int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
         llvm::dbgs() << "  alloc #" << BitNr << ": alive=";
         dumpBits(StackLocs[BitNr].AliveLocs);
-        llvm::dbgs() << "    " << *StackInst;
+        llvm::dbgs() << ",     " << *StackInst;
       } else if (StackInst->isDeallocatingStack()) {
         auto *AllocInst = cast<SingleValueInstruction>(StackInst->getOperand(0));
         int BitNr = StackLoc2BitNumbers.lookup(AllocInst);
@@ -405,6 +434,5 @@ void StackNesting::dumpBits(const BitVector &Bits) {
     llvm::dbgs() << separator << Bit;
     separator = ",";
   }
-  llvm::dbgs() << ">\n";
+  llvm::dbgs() << '>';
 }
-

test/SILOptimizer/allocbox_to_stack.sil

@@ -819,6 +819,8 @@ bb0(%0 : $Int):
 // CHECK-NEXT:   [[STACK1:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
+// CHECK-NEXT:   dealloc_stack [[STACK1]]
 // CHECK-NEXT:   [[STACK2:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   dealloc_stack [[STACK2]]
 // CHECK-NEXT:   unreachable
@@ -952,6 +954,7 @@ bb2:
 // CHECK:      bb0(%0 : $Int):
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   unreachable
@@ -979,6 +982,38 @@ bb2:
   return %r : $()
 }
 
+// CHECK-LABEL: sil @nesting_and_unreachable6
+// CHECK:      bb0(%0 : $Int):
+// CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
+// CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
+// CHECK-NEXT:   br bb3
+// CHECK:      bb2:
+// CHECK-NEXT:   alloc_stack
+// CHECK-NEXT:   unreachable
+// CHECK:      bb3:
+// CHECK-NEXT:   return
+sil @nesting_and_unreachable6 : $@convention(thin) (Int) -> Int {
+bb0(%0 : $Int):
+  %1 = alloc_box ${ var Int }
+  %1a = project_box %1 : ${ var Int }, 0
+  store %0 to %1a : $*Int
+  %3 = load %1a : $*Int
+  cond_br undef, bb1, bb2
+
+bb1:
+  strong_release %1 : ${ var Int }
+  br bb3
+
+bb2:
+  %241 = alloc_stack $Int
+  strong_release %1 : ${ var Int }
+  unreachable
+
+bb3:
+  return %3 : $Int
+}
+
 // CHECK-LABEL: sil @nesting_and_unreachable_critical_edge
 // CHECK:      bb0(%0 : $Int):
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -814,6 +814,8 @@ bb0(%0 : $Int):
 // CHECK-NEXT:   [[STACK1:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
+// CHECK-NEXT:   dealloc_stack [[STACK1]]
 // CHECK-NEXT:   [[STACK2:%[0-9]+]] = alloc_stack $Bool
 // CHECK-NEXT:   dealloc_stack [[STACK2]]
 // CHECK-NEXT:   unreachable
@@ -947,6 +949,7 @@ bb2:
 // CHECK:      bb0(%0 : $Int):
 // CHECK-NEXT:   [[BOX:%[0-9]+]] = alloc_stack $Int
 // CHECK:      bb1:
+// CHECK-NEXT:   dealloc_stack [[BOX]]
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   {{.*}} = alloc_stack $Bool
 // CHECK-NEXT:   unreachable