Merge pull request #64005 from nate-chandler/cse/partial-apply-stack-discipline

[CSE] Stack nest at OSSA lowering after inlining.

Author: nate-chandler

lib/SILOptimizer/Transforms/CSE.cpp

@@ -37,6 +37,7 @@
 #include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
 #include "swift/SILOptimizer/Utils/SILInliner.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopedHashTable.h"
@@ -665,7 +666,7 @@ class CSE {
 
   bool processFunction(SILFunction &F, DominanceInfo *DT);
 
-  bool processLazyPropertyGetters();
+  bool processLazyPropertyGetters(SILFunction &F);
 
   bool canHandle(SILInstruction *Inst);
 
@@ -779,8 +780,9 @@ bool CSE::processFunction(SILFunction &Fm, DominanceInfo *DT) {
 
 /// Replace lazy property getters (which are dominated by the same getter)
 /// by a direct load of the value.
-bool CSE::processLazyPropertyGetters() {
+bool CSE::processLazyPropertyGetters(SILFunction &F) {
   bool changed = false;
+  bool invalidatedStackNesting = false;
   for (ApplyInst *ai : lazyPropertyGetters) {
     SILFunction *getter = ai->getReferencedFunctionOrNull();
     assert(getter && getter->isLazyPropertyGetter());
@@ -807,9 +809,19 @@ bool CSE::processLazyPropertyGetters() {
         builder.createUncheckedEnumData(sei->getLoc(), enumVal, someDecl, ty);
     builder.createBranch(sei->getLoc(), someDest, { ued });
     sei->eraseFromParent();
+    // 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.
+    if (getter->hasOwnership() && !ai->getFunction()->hasOwnership()) {
+      invalidatedStackNesting = true;
+    }
     changed = true;
     ++NumCSE;
   }
+  if (invalidatedStackNesting) {
+    StackNesting::fixNesting(&F);
+  }
   return changed;
 }
 
@@ -1475,7 +1487,7 @@ class SILCSE : public SILFunctionTransform {
 
     // Handle calls to lazy property getters, which are collected in
     // processFunction().
-    if (C.processLazyPropertyGetters()) {
+    if (C.processLazyPropertyGetters(*Fn)) {
       // Cleanup the dead blocks from the inlined lazy property getters.
       removeUnreachableBlocks(*Fn);
       invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);

test/SILOptimizer/cse.sil

@@ -1356,3 +1356,58 @@ bb0:
   %22 = tuple ()
   return %22 : $()
 }
+
+sil @paable : $@convention(thin) (Builtin.Int64) -> ()
+
+final class FC {
+  @_hasStorage @_hasInitialValue final var storage: Optional<Builtin.Int64>
+}
+
+sil @consume_int64 : $@convention(thin) (Builtin.Int64) -> ()
+
+sil [lazy_getter] [ossa] @partial_apply_on_stack_nesting_violator : $@convention(method) (@guaranteed FC) -> Builtin.Int64 {
+entry(%instance : @guaranteed $FC):
+    %storage_addr = ref_element_addr %instance : $FC, #FC.storage
+    %optional = load [trivial] %storage_addr : $*Optional<Builtin.Int64>
+    switch_enum %optional : $Optional<Builtin.Int64>, case #Optional.some!enumelt: some, case #Optional.none!enumelt: none
+none:
+    %one = integer_literal $Builtin.Int64, 1
+    %optone = enum $Optional<Builtin.Int64>, #Optional.some!enumelt, %one : $Builtin.Int64
+    store %optone to [trivial] %storage_addr : $*Optional<Builtin.Int64>
+    br exit(%one : $Builtin.Int64)
+
+some(%value : $Builtin.Int64):
+    %paable = function_ref @paable : $@convention(thin) (Builtin.Int64) -> ()
+    %first = partial_apply [callee_guaranteed] [on_stack] %paable(%value) : $@convention(thin) (Builtin.Int64) -> ()
+    %second = partial_apply [callee_guaranteed] [on_stack] %paable(%value) : $@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 () -> ()
+    br exit(%value : $Builtin.Int64)
+
+exit(%retval : $Builtin.Int64):
+    return %retval : $Builtin.Int64
+}
+
+// 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) (FC) -> () {
+entry(%instance : $FC):
+    %callee = function_ref @partial_apply_on_stack_nesting_violator : $@convention(method) (@guaranteed FC) -> Builtin.Int64
+    %consume = function_ref @consume_int64 : $@convention(thin) (Builtin.Int64) -> ()
+    %first = apply %callee(%instance) : $@convention(method) (@guaranteed FC) -> Builtin.Int64
+    apply %consume(%first) : $@convention(thin) (Builtin.Int64) -> ()
+    %second = apply %callee(%instance) : $@convention(method) (@guaranteed FC) -> Builtin.Int64
+    apply %consume(%second) : $@convention(thin) (Builtin.Int64) -> ()
+    %retval = tuple ()
+    return %retval : $()
+}
+