[Inliner] Fix stack nesting when lowering OSSA.

lib/SILOptimizer/Mandatory/MandatoryInlining.cpp

@@ -50,6 +50,8 @@ extern llvm::cl::opt<bool> SILPrintInliningCallerBefore;
 
 extern llvm::cl::opt<bool> SILPrintInliningCallerAfter;
 
+extern llvm::cl::opt<bool> EnableVerifyAfterEachInlining;
+
 extern void printInliningDetailsCallee(StringRef passName, SILFunction *caller,
                                        SILFunction *callee);
 
@@ -954,6 +956,14 @@ runOnFunctionRecursively(SILOptFunctionBuilder &FuncBuilder, SILFunction *F,
       // nextBB will point to the last inlined block
       SILBasicBlock *lastBB =
           Inliner.inlineFunction(CalleeFunction, InnerAI, FullArgs);
+
+      // When inlining an OSSA function into a non-OSSA function, ownership of
+      // nonescaping closures is lowered.  At that point, they are recognized
+      // as stack users.  Since they weren't recognized as such before, they
+      // may not satisfy stack discipline.  Fix that up now.
+      invalidatedStackNesting |=
+          (CalleeFunction->hasOwnership() && !F->hasOwnership());
+
       if (SILPrintInliningCallerAfter) {
         printInliningDetailsCallerAfter("MandatoryInlining", F, CalleeFunction);
       }
@@ -972,6 +982,16 @@ runOnFunctionRecursively(SILOptFunctionBuilder &FuncBuilder, SILFunction *F,
       // later.
       changedFunctions.insert(F);
 
+      if (EnableVerifyAfterEachInlining) {
+        if (invalidatedStackNesting) {
+          StackNesting::fixNesting(F);
+          changedFunctions.insert(F);
+          invalidatedStackNesting = false;
+        }
+
+        F->verify();
+      }
+
       // Resume inlining within nextBB, which contains only the inlined
       // instructions and possibly instructions in the original call block that
       // have not yet been visited.

lib/SILOptimizer/Transforms/PerformanceInliner.cpp

@@ -63,6 +63,12 @@ llvm::cl::opt<bool> SILPrintInliningCallerAfter(
     llvm::cl::desc(
         "Print functions into which another function has been inlined."));
 
+llvm::cl::opt<bool> EnableVerifyAfterEachInlining(
+    "sil-inline-verify-after-each-inline", llvm::cl::init(false),
+    llvm::cl::desc(
+        "Run sil verification after inlining each found callee apply "
+        "site into a caller."));
+
 //===----------------------------------------------------------------------===//
 //                           Printing Helpers
 //===----------------------------------------------------------------------===//
@@ -1034,10 +1040,30 @@ bool SILPerformanceInliner::inlineCallsIntoFunction(SILFunction *Caller) {
     // will assert, so we are safe making this assumption.
     SILInliner::inlineFullApply(AI, SILInliner::InlineKind::PerformanceInline,
                                 FuncBuilder, deleter);
+    // When inlining an OSSA function into a non-OSSA function, ownership of
+    // nonescaping closures is lowered.  At that point, they are recognized as
+    // stack users.  Since they weren't recognized as such before, they may not
+    // satisfy stack discipline.  Fix that up now.
+    invalidatedStackNesting |=
+        Callee->hasOwnership() && !Caller->hasOwnership();
     ++NumFunctionsInlined;
     if (SILPrintInliningCallerAfter) {
       printInliningDetailsCallerAfter(PassName, Caller, Callee);
     }
+    if (EnableVerifyAfterEachInlining) {
+      deleter.cleanupDeadInstructions();
+
+      // The inliner splits blocks at call sites. Re-merge trivial branches to
+      // reestablish a canonical CFG.
+      mergeBasicBlocks(Caller);
+
+      if (invalidatedStackNesting) {
+        StackNesting::fixNesting(Caller);
+        invalidatedStackNesting = false;
+      }
+
+      Caller->verify();
+    }
   }
   deleter.cleanupDeadInstructions();
 

test/SILOptimizer/inline_ossa_to_non_ossa.sil

@@ -0,0 +1,36 @@
+// RUN: %target-sil-opt -enable-sil-verify-all %s -early-inline | %FileCheck %s
+
+import Builtin
+
+sil @paable : $@convention(thin) (Builtin.Int64) -> ()
+
+sil [ossa] @partial_apply_on_stack_nesting_violator : $@convention(thin) () -> () {
+    %paable = function_ref @paable : $@convention(thin) (Builtin.Int64) -> ()
+    %one = integer_literal $Builtin.Int64, 1
+    %first = partial_apply [callee_guaranteed] [on_stack] %paable(%one) : $@convention(thin) (Builtin.Int64) -> ()
+    %two = integer_literal $Builtin.Int64, 2
+    %second = partial_apply [callee_guaranteed] [on_stack] %paable(%two) : $@convention(thin) (Builtin.Int64) -> ()
+    // Note that the destroy_values do not occur in an order which coincides
+    // with stack disciplined dealloc_stacks.
+    destroy_value %first : $@noescape @callee_guaranteed () -> ()
+    destroy_value %second : $@noescape @callee_guaranteed () -> ()
+    %retval = tuple ()
+    return %retval : $()
+}
+
+// Verify that when inlining partial_apply_on_stack_nesting_violator, the stack
+// nesting of the on_stack closures is fixed.
+// CHECK-LABEL: sil @test_inline_stack_violating_ossa_func : {{.*}} {
+// CHECK:         [[PAABLE:%[^,]+]] = function_ref @paable
+// CHECK:         [[FIRST:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[PAABLE]]
+// CHECK:         [[SECOND:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[PAABLE]]
+// CHECK:         dealloc_stack [[SECOND]]
+// CHECK:         dealloc_stack [[FIRST]]
+// CHECK-LABEL: } // end sil function 'test_inline_stack_violating_ossa_func'
+sil @test_inline_stack_violating_ossa_func : $@convention(thin) () -> () {
+    %callee = function_ref @partial_apply_on_stack_nesting_violator : $@convention(thin) () -> ()
+    apply %callee() : $@convention(thin) () -> ()
+    %retval = tuple ()
+    return %retval : $()
+}
+

test/SILOptimizer/mandatory_inlining_ossa_to_non_ossa.sil

@@ -456,3 +456,36 @@ bb0(%0 : $Builtin.NativeObject):
   %9999 = tuple()
   return %9999 : $()
 }
+
+sil @paable : $@convention(thin) (Builtin.Int64) -> ()
+
+sil [transparent] [ossa] @partial_apply_on_stack_nesting_violator : $@convention(thin) () -> () {
+    %paable = function_ref @paable : $@convention(thin) (Builtin.Int64) -> ()
+    %one = integer_literal $Builtin.Int64, 1
+    %first = partial_apply [callee_guaranteed] [on_stack] %paable(%one) : $@convention(thin) (Builtin.Int64) -> ()
+    %two = integer_literal $Builtin.Int64, 2
+    %second = partial_apply [callee_guaranteed] [on_stack] %paable(%two) : $@convention(thin) (Builtin.Int64) -> ()
+    // Note that the destroy_values do not occur in an order which coincides
+    // with stack disciplined dealloc_stacks.
+    destroy_value %first : $@noescape @callee_guaranteed () -> ()
+    destroy_value %second : $@noescape @callee_guaranteed () -> ()
+    %retval = tuple ()
+    return %retval : $()
+}
+
+// Verify that when inlining partial_apply_on_stack_nesting_violator, the stack
+// nesting of the on_stack closures is fixed.
+// CHECK-LABEL: sil @test_inline_stack_violating_ossa_func : {{.*}} {
+// CHECK:         [[PAABLE:%[^,]+]] = function_ref @paable
+// CHECK:         [[FIRST:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[PAABLE]]
+// CHECK:         [[SECOND:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[PAABLE]]
+// CHECK:         dealloc_stack [[SECOND]]
+// CHECK:         dealloc_stack [[FIRST]]
+// CHECK-LABEL: } // end sil function 'test_inline_stack_violating_ossa_func'
+sil @test_inline_stack_violating_ossa_func : $@convention(thin) () -> () {
+    %callee = function_ref @partial_apply_on_stack_nesting_violator : $@convention(thin) () -> ()
+    apply %callee() : $@convention(thin) () -> ()
+    %retval = tuple ()
+    return %retval : $()
+}
+