SILOptimizer: fix a stack nesting problem when inlining coroutines

include/swift/SILOptimizer/Utils/SILInliner.h

@@ -67,6 +67,17 @@ class SILInliner {
   /// attempts to inline \arg AI and this returns false, an assert will fire.
   static bool canInlineApplySite(FullApplySite apply);
 
+  /// Returns true if inlining \arg apply can result in improperly nested stack
+  /// allocations.
+  ///
+  /// In this case stack nesting must be corrected after inlining with the
+  /// StackNesting utility.
+  static bool needsUpdateStackNesting(FullApplySite apply) {
+    // Inlining of coroutines can result in improperly nested stack
+    // allocations.
+    return isa<BeginApplyInst>(apply);
+  }
+
   /// Allow the client to track instructions before they are deleted. The
   /// registered callback is called from
   /// recursivelyDeleteTriviallyDeadInstructions.
@@ -99,6 +110,9 @@ class SILInliner {
   /// to be inlined using SILInliner::canInlineApplySite before calling this
   /// function.
   ///
+  /// *NOTE*: Inlining can result in improperly nested stack allocations, which
+  /// must be corrected after inlining. See needsUpdateStackNesting().
+  ///
   /// Returns an iterator to the first inlined instruction (or the end of the
   /// caller block for empty functions) and the last block in function order
   /// containing inlined instructions (the original caller block for

include/swift/SILOptimizer/Utils/StackNesting.h

@@ -20,25 +20,6 @@
 
 namespace swift {
 
-struct InstructionMatchResult {
-  /// The instruction matches.
-  bool matches;
-
-  /// The search should be halted.
-  bool halt;
-};
-
-using InstructionMatcher =
-  llvm::function_ref<InstructionMatchResult(SILInstruction *i)>;
-
-/// Given a predicate defining interesting instructions, check whether
-/// the stack is different at any of them from how it stood at the start
-/// of the search.
-///
-/// The matcher function must halt the search before any edges which would
-/// lead back to before the start point.
-bool hasStackDifferencesAt(SILInstruction *start, InstructionMatcher matcher);
-
 /// A utility to correct the nesting of stack allocating/deallocating
 /// instructions.
 ///

lib/SILOptimizer/Mandatory/MandatoryInlining.cpp

@@ -25,6 +25,7 @@
 #include "swift/SILOptimizer/Utils/Local.h"
 #include "swift/SILOptimizer/Utils/SILInliner.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/ImmutableSet.h"
 #include "llvm/ADT/Statistic.h"
@@ -768,6 +769,7 @@ runOnFunctionRecursively(SILOptFunctionBuilder &FuncBuilder,
 
   SmallVector<ParameterConvention, 16> CapturedArgConventions;
   SmallVector<SILValue, 32> FullArgs;
+  bool needUpdateStackNesting = false;
 
   // Visiting blocks in reverse order avoids revisiting instructions after block
   // splitting, which would be quadratic.
@@ -886,6 +888,8 @@ runOnFunctionRecursively(SILOptFunctionBuilder &FuncBuilder,
         closureCleanup.recordDeadFunction(I);
       });
 
+      needUpdateStackNesting |= Inliner.needsUpdateStackNesting(InnerAI);
+
       // Inlining deletes the apply, and can introduce multiple new basic
       // blocks. After this, CalleeValue and other instructions may be invalid.
       // nextBB will point to the last inlined block
@@ -905,6 +909,11 @@ runOnFunctionRecursively(SILOptFunctionBuilder &FuncBuilder,
       break;
     }
   }
+
+  if (needUpdateStackNesting) {
+    StackNesting().correctStackNesting(F);
+  }
+
   // Keep track of full inlined functions so we don't waste time recursively
   // reprocessing them.
   FullyInlinedSet.insert(F);

lib/SILOptimizer/Transforms/PerformanceInliner.cpp

@@ -22,6 +22,7 @@
 #include "swift/SILOptimizer/Utils/Generics.h"
 #include "swift/SILOptimizer/Utils/PerformanceInlinerUtils.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -877,6 +878,7 @@ bool SILPerformanceInliner::inlineCallsIntoFunction(SILFunction *Caller) {
   // remains valid.
   SmallVector<FullApplySite, 8> AppliesToInline;
   collectAppliesToInline(Caller, AppliesToInline);
+  bool needUpdateStackNesting = false;
 
   if (AppliesToInline.empty())
     return false;
@@ -909,6 +911,8 @@ bool SILPerformanceInliner::inlineCallsIntoFunction(SILFunction *Caller) {
     SILInliner Inliner(FuncBuilder, SILInliner::InlineKind::PerformanceInline,
                        AI.getSubstitutionMap(), OpenedArchetypesTracker);
 
+    needUpdateStackNesting |= Inliner.needsUpdateStackNesting(AI);
+
     // We've already determined we should be able to inline this, so
     // unconditionally inline the function.
     //
@@ -920,6 +924,11 @@ bool SILPerformanceInliner::inlineCallsIntoFunction(SILFunction *Caller) {
   // The inliner splits blocks at call sites. Re-merge trivial branches to
   // reestablish a canonical CFG.
   mergeBasicBlocks(Caller);
+
+  if (needUpdateStackNesting) {
+    StackNesting().correctStackNesting(Caller);
+  }
+
   return true;
 }
 

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -17,54 +17,10 @@
 #include "swift/SIL/TypeSubstCloner.h"
 #include "swift/SILOptimizer/Utils/CFG.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
-#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 using namespace swift;
 
-/// Does the given coroutine make any stack allocations that are live across
-/// its yields?
-static bool allocatesStackAcrossYields(SILFunction *F) {
-  assert(F->getLoweredFunctionType()->isCoroutine());
-
-  return hasStackDifferencesAt(&F->getEntryBlock()->front(),
-                               [](SILInstruction *i) -> InstructionMatchResult {
-    if (isa<YieldInst>(i)) {
-      return {
-        /*matches*/ true,
-        /*halt*/ i->getFunction()->getLoweredFunctionType()->getCoroutineKind()
-                       == SILCoroutineKind::YieldOnce
-      };
-    }
-
-    // Otherwise, search until the end of the function.
-    return { false, false };
-  });
-}
-
-static bool isEndOfApply(SILInstruction *i, BeginApplyInst *beginApply) {
-  if (auto endApply = dyn_cast<EndApplyInst>(i)) {
-    return endApply->getBeginApply() == beginApply;
-  } else if (auto abortApply = dyn_cast<AbortApplyInst>(i)) {
-    return abortApply->getBeginApply() == beginApply;
-  } else {
-    return false;
-  }
-}
-
-/// Are there any stack differences from the given begin_apply to any
-/// corresponding end_apply/abort_apply?
-static bool hasStackDifferencesAtEnds(BeginApplyInst *apply) {
-  return hasStackDifferencesAt(apply, [apply](SILInstruction *i)
-                                                     -> InstructionMatchResult {
-    // Search for ends of the original apply.  We can stop searching
-    // at these points.
-    if (isEndOfApply(i, apply))
-      return { true, true };
-    return { false, false };
-  });
-}
-
 static bool canInlineBeginApply(BeginApplyInst *BA) {
   // Don't inline if we have multiple resumption sites (i.e. end_apply or
   // abort_apply instructions).  The current implementation clones a single
@@ -122,7 +78,6 @@ class BeginApplySite {
   SILBuilder *Builder;
   BeginApplyInst *BeginApply;
   bool HasYield = false;
-  bool NeedsStackCorrection;
 
   EndApplyInst *EndApply = nullptr;
   SILBasicBlock *EndApplyBB = nullptr;
@@ -134,31 +89,15 @@ class BeginApplySite {
 
 public:
   BeginApplySite(BeginApplyInst *BeginApply, SILLocation Loc,
-                 SILBuilder *Builder, bool NeedsStackCorrection)
-      : Loc(Loc), Builder(Builder), BeginApply(BeginApply),
-        NeedsStackCorrection(NeedsStackCorrection) {}
+                 SILBuilder *Builder)
+      : Loc(Loc), Builder(Builder), BeginApply(BeginApply) {}
 
   static Optional<BeginApplySite> get(FullApplySite AI, SILLocation Loc,
                                       SILBuilder *Builder) {
     auto *BeginApply = dyn_cast<BeginApplyInst>(AI);
     if (!BeginApply)
       return None;
-
-    // We need stack correction if there are both:
-    //   - stack allocations in the callee that are live across the yield and
-    //   - stack differences in the caller from the begin_apply to any
-    //     end_apply or abort_apply.
-    // In these cases, naive cloning will cause the allocations to become
-    // improperly nested.
-    //
-    // We need to compute this here before we do any splitting in the parent
-    // function.
-    bool NeedsStackCorrection = false;
-    if (allocatesStackAcrossYields(BeginApply->getReferencedFunction()) &&
-        hasStackDifferencesAtEnds(BeginApply))
-      NeedsStackCorrection = true;
-
-    return BeginApplySite(BeginApply, Loc, Builder, NeedsStackCorrection);
+    return BeginApplySite(BeginApply, Loc, Builder);
   }
 
   void preprocess(SILBasicBlock *returnToBB) {
@@ -286,11 +225,6 @@ class BeginApplySite {
       AbortApply->eraseFromParent();
 
     assert(!BeginApply->hasUsesOfAnyResult());
-
-    // Correct the stack if necessary.
-    if (NeedsStackCorrection) {
-      StackNesting().correctStackNesting(BeginApply->getFunction());
-    }
   }
 };
 } // namespace swift

lib/SILOptimizer/Utils/StackNesting.cpp

@@ -18,75 +18,6 @@
 
 using namespace swift;
 
-bool swift::hasStackDifferencesAt(SILInstruction *start,
-                                  InstructionMatcher matcher) {
-  struct StackStatus {
-    /// The number of currently-active stack allocations on this path.
-    /// Because of the stack discipline, we only have to maintain a count.
-    unsigned depth;
-
-    /// Whether we've popped anything off of the stack that was active
-    /// at the start of the search.
-    bool hasPops;
-
-    bool hasDifferences() const { return hasPops || depth != 0; }
-  };
-
-  SmallPtrSet<SILBasicBlock*, 8> visited;
-  SmallVector<std::pair<SILInstruction *, StackStatus>, 8> worklist;
-  worklist.push_back({start, {0, false}});
-
-  while (!worklist.empty()) {
-    // Pull a block and depth off the worklist.
-    // Visitation order doesn't matter.
-    auto pair = worklist.pop_back_val();
-    auto status = pair.second;
-    auto firstInst = pair.first;
-    auto block = firstInst->getParent();
-
-    for (SILBasicBlock::iterator i(firstInst), e = block->end(); i != e; ++i) {
-      auto match = matcher(&*i);
-
-      // If this is an interesting instruction, check for stack
-      // differences here.
-      if (match.matches && status.hasDifferences())
-        return true;
-
-      // If we should halt, just go to the next work-list item, bypassing
-      // visiting the successors.
-      if (match.halt)
-        goto nextWorkListItem; // a labelled continue would be nice
-
-      // Otherwise, do stack-depth bookkeeping.
-      if (i->isAllocatingStack()) {
-        status.depth++;
-      } else if (i->isDeallocatingStack()) {
-        // If the stack depth is already zero, we're deallocating a stack
-        // allocation that was live into the search.
-        if (status.depth > 0) {
-          status.depth--;
-        } else {
-          status.hasPops = true;
-        }
-      }
-    }
-
-    // Add the successors to the worklist.
-    for (auto &succ : block->getSuccessors()) {
-      auto succBlock = succ.getBB();
-      auto insertResult = visited.insert(succBlock);
-      if (insertResult.second) {
-        worklist.push_back({&succBlock->front(), status});
-      }
-    }
-
-  nextWorkListItem: ;
-  }
-
-  return false;
-}
-
-
 void StackNesting::setup(SILFunction *F) {
   SmallVector<BlockInfo *, 8> WorkList;
   llvm::DenseMap<SILBasicBlock *, BlockInfo *> BlockMapping;

test/SILOptimizer/inline_begin_apply.sil

@@ -485,3 +485,38 @@ bb0:
   %ret = tuple ()
   return %ret : $()
 }
+
+// CHECK-LABEL: sil [ossa] @test_stack_overlap_unreachable
+// CHECK: bb0:
+// CHECK:   [[A16:%.*]] = alloc_stack $Builtin.Int16
+// CHECK:   [[A32:%.*]] = alloc_stack $Builtin.Int32
+// CHECK: bb1:
+// CHECK:   dealloc_stack [[A32]] : $*Builtin.Int32
+// CHECK:   unreachable
+// CHECK: bb2:
+// CHECK:   dealloc_stack [[A32]] : $*Builtin.Int32
+// CHECK:   dealloc_stack [[A16]] : $*Builtin.Int16
+// CHECK:   unreachable
+// CHECK: bb3:
+// CHECK:   dealloc_stack [[A32]] : $*Builtin.Int32
+// CHECK:   dealloc_stack [[A16]] : $*Builtin.Int16
+// CHECK:   return
+// CHECK-LABEL: // end sil function 'test_stack_overlap_unreachable'
+sil [ossa] @test_stack_overlap_unreachable : $() -> () {
+bb0:
+  %stack = alloc_stack $Builtin.Int16
+  %0 = function_ref @stack_overlap : $@convention(thin) @yield_once () -> (@yields @inout Builtin.Int32)
+  (%value, %token) = begin_apply %0() : $@convention(thin) @yield_once () -> (@yields @inout Builtin.Int32)
+  cond_br undef, bb1, bb2
+
+bb1:
+  dealloc_stack %stack : $*Builtin.Int16
+  unreachable
+
+bb2:
+  end_apply %token
+  dealloc_stack %stack : $*Builtin.Int16
+  %ret = tuple ()
+  return %ret : $()
+}
+