[optimizer] Teach SIL optimizer that stack nesting should ignore nested stack allocations.

Author: gottesmm

lib/SIL/Verifier/SILVerifier.cpp

@@ -7052,6 +7052,22 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
           }
         }
       }
+
+      void handleScopeInst(SILInstruction &i) {
+        if (!ActiveOps.insert(&i).second) {
+          verificationFailure("operation was not ended before re-beginning it",
+                              &i, nullptr, nullptr);
+        }
+      }
+
+      void handleScopeEndingInst(SILInstruction &i) {
+        if (auto beginOp = i.getOperand(0)->getDefiningInstruction()) {
+          if (!ActiveOps.erase(beginOp)) {
+            verificationFailure("operation has already been ended", &i, nullptr,
+                                nullptr);
+          }
+        }
+      }
     };
 
     struct DeadEndRegionState {
@@ -7152,22 +7168,31 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
         }
 
         if (i.isAllocatingStack()) {
-          if (auto *BAI = dyn_cast<BeginApplyInst>(&i)) {
-            state.Stack.push_back(BAI->getCalleeAllocationResult());
+          // Nested stack allocations are pushed on the stack. Non-nested stack
+          // allocations are treated as active ops so that we can at least
+          // verify their joint post-dominance.
+          if (i.isStackAllocationNested()) {
+            state.Stack.push_back(i.getStackAllocation());
+
+            // Also track begin_apply as a scope instruction.
+            if (isa<BeginApplyInst>(i))
+              state.handleScopeInst(i);
           } else {
-            state.Stack.push_back(cast<SingleValueInstruction>(&i));
+            state.handleScopeInst(i);
           }
-          // Not "else if": begin_apply both allocates stack and begins an
-          // operation.
-        }
-        if (i.isDeallocatingStack()) {
+        } else if (i.isDeallocatingStack()) {
           SILValue op = i.getOperand(0);
           while (auto *mvi = dyn_cast<MoveValueInst>(op)) {
             op = mvi->getOperand();
           }
 
-          if (!state.Stack.empty() && op == state.Stack.back()) {
+          auto beginInst = op->getDefiningInstruction();
+          if (beginInst && !beginInst->isStackAllocationNested()) {
+            state.handleScopeEndingInst(i);
+          } else if (!state.Stack.empty() && op == state.Stack.back()) {
             state.Stack.pop_back();
+            if (beginInst && isa<BeginApplyInst>(beginInst))
+              state.handleScopeEndingInst(i);
           } else {
             verificationFailure("deallocating allocation that is not the top of the stack",
                                 &i, nullptr,
@@ -7178,20 +7203,12 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
             });
           }
         } else if (isScopeInst(&i)) {
-          bool notAlreadyPresent = state.ActiveOps.insert(&i).second;
-          require(notAlreadyPresent,
-                  "operation was not ended before re-beginning it");
+          state.handleScopeInst(i);
         } else if (isScopeEndingInst(&i)) {
-          if (auto beginOp = i.getOperand(0)->getDefiningInstruction()) {
-            bool present = state.ActiveOps.erase(beginOp);
-            require(present, "operation has already been ended");
-          }
+          state.handleScopeEndingInst(i);
         } else if (auto *endBorrow = dyn_cast<EndBorrowInst>(&i)) {
           if (isa<StoreBorrowInst>(endBorrow->getOperand())) {
-            if (auto beginOp = i.getOperand(0)->getDefiningInstruction()) {
-              bool present = state.ActiveOps.erase(beginOp);
-              require(present, "operation has already been ended");
-            }
+            state.handleScopeEndingInst(i);
           }
         } else if (auto gaci = dyn_cast<GetAsyncContinuationInstBase>(&i)) {
           require(!state.GotAsyncContinuation,

lib/SILOptimizer/Utils/StackNesting.cpp

@@ -579,7 +579,9 @@ StackNesting::Changes StackNesting::fixNesting(SILFunction *F) {
       // In the formal presentation, the state change is
       //   state = STATE_PUSH(state, alloc)
       if (I->isAllocatingStack()) {
-        state.allocations.push_back(I);
+        // Only handle nested stack allocations.
+        if (I->isStackAllocationNested() == StackAllocationIsNested)
+          state.allocations.push_back(I);
         continue;
       }
 
@@ -594,6 +596,11 @@ StackNesting::Changes StackNesting::fixNesting(SILFunction *F) {
       SILInstruction *dealloc = I;
       SILInstruction *alloc = getAllocForDealloc(dealloc);
 
+      // Since we only processed nested allocations above, ignore deallocations
+      // from non-nested allocations.
+      if (alloc->isStackAllocationNested() == StackAllocationIsNotNested)
+        continue;
+
 #ifndef NDEBUG
       if (state.allocations.empty()) {
         state.abortForUnknownAllocation(alloc, dealloc);
@@ -657,9 +664,10 @@ StackNesting::Changes StackNesting::fixNesting(SILFunction *F) {
 }
 
 namespace swift::test {
-static FunctionTest MyNewTest("stack_nesting_fixup",
-                              [](auto &function, auto &arguments, auto &test) {
-                                StackNesting::fixNesting(&function);
-                                function.print(llvm::outs());
-                              });
+static FunctionTest StackNestingTests("stack_nesting_fixup",
+                                      [](auto &function, auto &arguments,
+                                         auto &test) {
+                                        StackNesting::fixNesting(&function);
+                                        function.print(llvm::outs());
+                                      });
 } // end namespace swift::test

test/SILOptimizer/stack-nesting.sil

@@ -281,3 +281,31 @@ bb5:
   dealloc_stack %d
   br bb4
 }
+
+// We only re-order b and c. We leave a alone.
+// CHECK-LABEL: sil [ossa] @test_simple_non_nested
+// CHECK:         integer_literal $Builtin.Int32, 1
+// CHECK-NEXT:    [[A:%.*]] = alloc_stack [non_nested] $Int
+// CHECK-NEXT:    [[B:%.*]] = alloc_stack $Int
+// CHECK-NEXT:    [[C:%.*]] = alloc_stack $Int
+// CHECK-NEXT:    integer_literal $Builtin.Int32, 2
+// CHECK-NEXT:    dealloc_stack [[A]]
+// CHECK-NEXT:    dealloc_stack [[C]]
+// CHECK-NEXT:    dealloc_stack [[B]]
+// CHECK-NEXT:    integer_literal $Builtin.Int32, 3
+// CHECK-LABEL: // end sil function 'test_simple_non_nested'
+sil [ossa] @test_simple_non_nested : $@convention(thin) (Builtin.Int1) -> () {
+bb0(%cond: $Builtin.Int1):
+  specify_test "stack_nesting_fixup"
+  integer_literal $Builtin.Int32, 1
+  %a = alloc_stack [non_nested] $Int
+  %b = alloc_stack $Int
+  %c = alloc_stack $Int
+  integer_literal $Builtin.Int32, 2
+  dealloc_stack %a
+  dealloc_stack %b
+  dealloc_stack %c
+  integer_literal $Builtin.Int32, 3
+  %ret = tuple ()
+  return %ret : $()
+}