Optimizer: re-implement and improve the AllocBoxToStack pass

This pass replaces `alloc_box` with `alloc_stack` if the box is not escaping.
The original implementation had some limitations. It could not handle cases of local functions which are called multiple times or even recursively, e.g.

```
public func foo() -> Int {
  var i = 1

  func localFunction() { i += 1 }

  localFunction()
  localFunction()
  return i
}

```

The new implementation (done in Swift) fixes this problem with a new algorithm.
It's not only more powerful, but also simpler: the new pass has less than half lines of code than the old pass.

The pass is invoked in the mandatory pipeline and later in the optimizer pipeline.
The new implementation provides a module-pass for the mandatory pipeline (whereas the "regular" pass is a function pass).
This is required because the mandatory pass needs to remove originals of specialized closures, which cannot be done from a function-pass.
In the old implementation this was done with a hack by adding a semantic attribute and deleting the function later in the pipeline.

I still kept the sources of the old pass for being able to bootstrap the compiler without a host compiler.

rdar://142756547

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -7,6 +7,7 @@
 # See http://swift.org/CONTRIBUTORS.txt for Swift project authors
 
 swift_compiler_sources(Optimizer
+  AllocBoxToStack.swift
   AssumeSingleThreaded.swift
   AsyncDemotion.swift
   BooleanLiteralFolding.swift

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

@@ -63,13 +63,15 @@ private func registerForSILCombine<InstType: SILCombineSimplifiable>(
 
 private func registerSwiftPasses() {
   // Module passes
+  registerPass(mandatoryAllocBoxToStack, { mandatoryAllocBoxToStack.run($0) })
   registerPass(mandatoryPerformanceOptimizations, { mandatoryPerformanceOptimizations.run($0) })
   registerPass(diagnoseUnknownConstValues, { diagnoseUnknownConstValues.run($0)})
   registerPass(readOnlyGlobalVariablesPass, { readOnlyGlobalVariablesPass.run($0) })
   registerPass(stackProtection, { stackProtection.run($0) })
   registerPass(embeddedSwiftDiagnostics, { embeddedSwiftDiagnostics.run($0) })
 
   // Function passes
+  registerPass(allocBoxToStack, { allocBoxToStack.run($0) })
   registerPass(asyncDemotion, { asyncDemotion.run($0) })
   registerPass(booleanLiteralFolding, { booleanLiteralFolding.run($0) })
   registerPass(letPropertyLowering, { letPropertyLowering.run($0) })

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -59,6 +59,8 @@
 #endif
 
 
+PASS(AllocBoxToStack, "allocbox-to-stack",
+     "stack promotion of box objects")
 PASS(CopyToBorrowOptimization, "copy-to-borrow-optimization",
      "Convert load [copy] instructions to load_borrow and remove copies of borrowed values")
 PASS(AliasInfoDumper, "dump-alias-info",
@@ -153,6 +155,8 @@ PASS(ExperimentalSwiftBasedClosureSpecialization, "experimental-swift-based-clos
 PASS(AutodiffClosureSpecialization, "autodiff-closure-specialization",
      "Autodiff specific closure-specialization pass")
 
+MODULE_PASS(MandatoryAllocBoxToStack, "mandatory-allocbox-to-stack",
+     "Mandatory stack promotion of box objects")
 MODULE_PASS(AsyncDemotion, "async-demotion",
      "Convert async functions to be synchronous")
 MODULE_PASS(RunUnitTests, "run-unit-tests",
@@ -214,7 +218,7 @@ LEGACY_PASS(AccessMarkerElimination, "access-marker-elim",
      "Access Marker Elimination.")
 LEGACY_PASS(AddressLowering, "address-lowering",
      "SIL Address Lowering")
-LEGACY_PASS(AllocBoxToStack, "allocbox-to-stack",
+LEGACY_PASS(LegacyAllocBoxToStack, "legacy-allocbox-to-stack",
      "Stack Promotion of Box Objects")
 LEGACY_PASS(ArrayCountPropagation, "array-count-propagation",
      "Array Count Propagation")

include/swift/SILOptimizer/PassManager/Passes.def

@@ -117,15 +117,18 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   // This guarantees that stack-promotable boxes have [static] enforcement.
   P.addAccessEnforcementSelection();
 
-  P.addAllocBoxToStack();
+#ifdef SWIFT_ENABLE_SWIFT_IN_SWIFT
+  P.addMandatoryAllocBoxToStack();
+#else
+  P.addLegacyAllocBoxToStack();
+#endif
   P.addNoReturnFolding();
   P.addBooleanLiteralFolding();
   addDefiniteInitialization(P);
 
   P.addAddressLowering();
 
-  // Before we run later semantic optimizations, eliminate simple functions that
-  // we specialized to ensure that we do not emit diagnostics twice.
+  // TODO: remove this once CapturePromotion deletes specialized functions itself.
   P.addDiagnosticDeadFunctionElimination();
 
   P.addFlowIsolation();

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -1315,6 +1315,6 @@ class AllocBoxToStack : public SILFunctionTransform {
 };
 } // end anonymous namespace
 
-SILTransform *swift::createAllocBoxToStack() {
+SILTransform *swift::createLegacyAllocBoxToStack() {
   return new AllocBoxToStack();
 }

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -798,7 +798,7 @@ func testActiveOptional(_ x: Float) -> Float {
 // CHECK: [NONE]   // function_ref _diagnoseUnexpectedNilOptional(_filenameStart:_filenameLength:_filenameIsASCII:_line:_isImplicitUnwrap:)
 // CHECK: [NONE]   %24 = apply %23(%17, %18, %19, %20, %22) : $@convention(thin) (Builtin.RawPointer, Builtin.Word, Builtin.Int1, Builtin.Word, Builtin.Int1) -> ()
 // CHECK: bb2:
-// CHECK: [ACTIVE] %26 = argument of bb2 : $Float
+// CHECK: [ACTIVE] {{%[0-9]+}} = argument of bb2 : $Float
 
 enum DirectEnum: Differentiable & AdditiveArithmetic {
   case case0

test/AutoDiff/SILOptimizer/activity_analysis.swift

@@ -3,7 +3,7 @@
 // Make sure that optimization passes don't choke on storage types for generic tuples
 // RUN: %target-swift-frontend -module-name generic_tuples -emit-ir -O %s
 
-// REQUIRES: CPU=x86_64
+// REQUIRES: PTRSIZE=64
 
 // CHECK-DAG: [[OPAQUE:%swift.opaque]] = type opaque
 // CHECK-DAG: [[TUPLE_TYPE:%swift.tuple_type]] = type { %swift.type, i64, ptr, [0 x %swift.tuple_element_type] }
@@ -26,9 +26,10 @@ func dup<T>(_ x: T) -> (T, T) { var x = x; return (x,x) }
 //   Copy 'x' into the first result.
 // CHECK-NEXT: call ptr [[WITNESS]](ptr noalias %0, ptr noalias [[X_ALLOCA]], ptr %T)
 //   Copy 'x' into the second element.
-// CHECK-NEXT: [[WITNESS_ADDR:%.*]] = getelementptr inbounds ptr, ptr [[VWT]], i32 4
-// CHECK-NEXT: [[WITNESS:%.*]] = load ptr, ptr [[WITNESS_ADDR]], align 8
 // CHECK-NEXT: call ptr [[WITNESS]](ptr noalias %1, ptr noalias [[X_ALLOCA]], ptr %T)
+// CHECK-NEXT: [[WITNESS_ADDR:%.*]] = getelementptr inbounds ptr, ptr [[VWT]], i32 1
+// CHECK-NEXT: [[DESTROYWITNESS:%.*]] = load ptr, ptr [[WITNESS_ADDR]], align 8
+// CHECK-NEXT: call void [[DESTROYWITNESS]](ptr noalias [[X_ALLOCA]],
 
 struct S {}
 

test/IRGen/generic_tuples.swift

@@ -125,10 +125,9 @@ bb1(%0 : $Int):
 
 sil @callee : $@convention(thin) (@inout Int) -> ()
 
-// CHECK-LABEL: sil @inout_nocapture
-sil @inout_nocapture : $@convention(thin) () -> Int {
+// CHECK-LABEL: sil @escaping_pointer
+sil @escaping_pointer : $@convention(thin) () -> Int {
 bb0:
-  // CHECK: alloc_stack
   %1 = alloc_box ${ var Int }
   %1a = project_box %1 : ${ var Int }, 0
   %6 = function_ref @callee : $@convention(thin) (@inout Int) -> ()
@@ -398,7 +397,7 @@ bb0(%0 : $Int):
 }
 
 // CHECK-LABEL: sil private @$s6struct8useStack1tySi_tFSiycfU_Tf0s_n
-// CHECK-LABEL: sil private [_semantics "sil.optimizer.delete_if_unused"] @$s6struct8useStack1tySi_tFSiycfU_
+// CHECK-LABEL: sil private @$s6struct8useStack1tySi_tFSiycfU_
 // struct.(useStack (t : Swift.Int) -> ()).(closure #1)
 sil private @$s6struct8useStack1tySi_tFSiycfU_ : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Int>) -> Int {
 bb0(%0 : $<τ_0_0> { var τ_0_0 } <Int>):
@@ -651,7 +650,7 @@ bb0(%0 : $Int, %1 : $*S<T>):
   return %9 : $Bool
 }
 
-// CHECK-LABEL: sil shared [_semantics "sil.optimizer.delete_if_unused"] @closure1
+// CHECK-LABEL: sil shared @closure1
 sil shared @closure1 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
 // CHECK: bb0
 bb0(%0 : $Int, %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):
@@ -673,7 +672,7 @@ bb0(%0 : $Int, %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):
 // CHECK: return
 // CHECK-NOT: bb1
 
-// CHECK-LABEL: sil shared [_semantics "sil.optimizer.delete_if_unused"] @closure2
+// CHECK-LABEL: sil shared @closure2
 sil shared @closure2 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
 // CHECK: bb0
 bb0(%0 : $Int, %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):

test/SILOptimizer/allocbox_to_stack.sil

@@ -54,6 +54,7 @@ bb0(%0 : $Int):
 // CHECK:         [[STACK:%[^,]+]] = alloc_stack
 // CHECK:         cond_br undef, [[DIE:bb[0-9]+]]
 // CHECK:       [[DIE]]:
+// CHECK-NEXT:    dealloc_stack [[STACK]]
 // CHECK-NEXT:    unreachable
 // CHECK-LABEL: } // end sil function 'keep_dead_end'
 sil [ossa] @keep_dead_end : $@convention(thin) () -> () {

test/SILOptimizer/allocbox_to_stack_lifetime.sil

@@ -28,8 +28,7 @@ bb0(%0 : @closureCapture @guaranteed ${ var NonTrivialStruct }):
 // CHECK-NEXT:   [[STACK2:%.*]] = alloc_stack $NonTrivialStruct
 // CHECK-NEXT:   [[MARKED2:%.*]] = mark_unresolved_non_copyable_value [consumable_and_assignable] [[STACK2]]
 // CHECK-NEXT:   store [[LOADED_VALUE]] to [init] [[MARKED2]]
-// CHECK-NEXT:   // function_ref
-// CHECK-NEXT:   [[FUNC:%.*]] = function_ref @$s26host_markmustcheck_closureTf0s_n : $@convention(thin) (@inout_aliasable NonTrivialStruct) -> ()
+// CHECK:        [[FUNC:%.*]] = function_ref @$s26host_markmustcheck_closureTf0s_n : $@convention(thin) (@inout_aliasable NonTrivialStruct) -> ()
 // CHECK-NEXT:   [[PA:%.*]] = partial_apply [callee_guaranteed] [[FUNC]]([[MARKED2]]) : $@convention(thin) (@inout_aliasable NonTrivialStruct) -> ()
 // CHECK-NEXT:   apply [[PA]]()
 // CHECK-NEXT:   destroy_value [[PA]]

test/SILOptimizer/allocbox_to_stack_noncopyable.sil

@@ -1,4 +1,5 @@
-// RUN: %target-sil-opt -sil-print-types -enable-sil-verify-all %s -allocbox-to-stack -enable-copy-propagation=requested-passes-only | %FileCheck %s
+// RUN: %target-sil-opt -sil-print-types %s -allocbox-to-stack | %FileCheck %s --check-prefix=CHECK --check-prefix=OPT
+// RUN: %target-sil-opt -sil-print-types %s -mandatory-allocbox-to-stack | %FileCheck %s --check-prefix=CHECK --check-prefix=MANDATORY
 
 sil_stage raw
 
@@ -123,10 +124,9 @@ bb1(%0 : $Int):
 
 sil [ossa] @callee : $@convention(thin) (@inout Int) -> ()
 
-// CHECK-LABEL: sil [ossa] @inout_nocapture
-sil [ossa] @inout_nocapture : $@convention(thin) () -> Int {
+// CHECK-LABEL: sil [ossa] @escaping_pointer
+sil [ossa] @escaping_pointer : $@convention(thin) () -> Int {
 bb0:
-  // CHECK: alloc_stack
   %1 = alloc_box ${ var Int }
   %1a = project_box %1 : ${ var Int }, 0
   %6 = function_ref @callee : $@convention(thin) (@inout Int) -> ()
@@ -394,8 +394,8 @@ bb0(%0 : $Int):
 }
 
 // CHECK-LABEL: sil private [transparent] [ossa] @$s6struct8useStack1tySi_tFSiycfU_Tf0s_n
-// CHECK-LABEL: sil private [transparent] [_semantics "sil.optimizer.delete_if_unused"] [ossa] @$s6struct8useStack1tySi_tFSiycfU_
-// struct.(useStack (t : Swift.Int) -> ()).(closure #1)
+// OPT-LABEL: sil private [transparent] [ossa] @$s6struct8useStack1tySi_tFSiycfU_ :
+// MANDATORY-NOT: sil private [transparent] [ossa] @$s6struct8useStack1tySi_tFSiycfU_ :
 sil private [transparent] [ossa] @$s6struct8useStack1tySi_tFSiycfU_ : $@convention(thin) (@owned <τ_0_0> { var τ_0_0 } <Int>) -> Int {
 bb0(%0 : @owned $<τ_0_0> { var τ_0_0 } <Int>):
   %1 = project_box %0 : $<τ_0_0> { var τ_0_0 } <Int>, 0
@@ -751,7 +751,8 @@ bb0(%0 : $Int, %1 : $*S<T>):
   return %9 : $Bool
 }
 
-// CHECK-LABEL: sil shared [_semantics "sil.optimizer.delete_if_unused"] [ossa] @closure1
+// OPT-LABEL:     sil shared [ossa] @closure1 :
+// MANDATORY-NOT: sil shared [ossa] @closure1 :
 sil shared [ossa] @closure1 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
 // CHECK: bb0
 bb0(%0 : $Int, %1 : @owned $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):
@@ -773,9 +774,9 @@ bb0(%0 : $Int, %1 : @owned $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>)
 // CHECK: return
 // CHECK-NOT: bb1
 
-// CHECK-LABEL: sil shared [_semantics "sil.optimizer.delete_if_unused"] [ossa] @closure2
+// OPT-LABEL: sil shared [ossa] @closure2 :
+// MANDATORY-NOT: sil shared [ossa] @closure2 :
 sil shared [ossa] @closure2 : $@convention(thin) <T where T : Count> (Int, @owned <τ_0_0 : Count> { var S<τ_0_0> } <T>) -> Bool {
-// CHECK: bb0
 bb0(%0 : $Int, %1 : @owned $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>):
   %2 = project_box %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>, 0
   %3 = function_ref @binary : $@convention(thin) (Int, Int) -> Bool
@@ -791,7 +792,6 @@ bb0(%0 : $Int, %1 : @owned $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>)
   destroy_addr %4 : $*S<T>
   dealloc_stack %4 : $*S<T>
   destroy_value %1 : $<τ_0_0 where τ_0_0 : Count> { var S<τ_0_0> } <T>
-// CHECK: return
   return %8 : $Bool
 }
 
@@ -1126,22 +1126,17 @@ bb2:
   return %r : $()
 }
 
-// Test that
-//  begin_access [read], copy_addr, end_access, destroy_addr
-// is folded into
-//  begin_access [deinit], copy_addr [take], end_access
-//
 // CHECK-LABEL: sil [ossa] @deinit_access : $@convention(thin) <T> (@in T) -> (@out T, @out T) {
 // CHECK: bb0(%0 : $*T, %1 : $*T, %2 : $*T):
 // CHECK: [[STK:%.*]] = alloc_stack $T, var, name "x"
 // CHECK: copy_addr %2 to [init] [[STK]] : $*T
 // CHECK: [[READ:%.*]] = begin_access [read] [static] [[STK]] : $*T
 // CHECK: copy_addr [[READ]] to [init] %0 : $*T
 // CHECK: end_access [[READ]] : $*T
-// CHECK: [[READ:%.*]] = begin_access [deinit] [static] [[STK]] : $*T
-// CHECK: copy_addr [take] [[READ]] to [init] %1 : $*T
+// CHECK: [[READ:%.*]] = begin_access [read] [static] [[STK]] : $*T
+// CHECK: copy_addr [[READ]] to [init] %1 : $*T
 // CHECK: end_access [[READ]] : $*T
-// CHECK-NOT: destroy_addr
+// CHECK: destroy_addr [[STK]]
 // CHECK: dealloc_stack [[STK]] : $*T
 // CHECK: destroy_addr %2 : $*T
 // CHECK: return %{{.*}} : $()

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -257,7 +257,7 @@ bb0(%0 : ${ var Int64 }):
 }
 
 // CHECK-LABEL: sil [noinline] @$testrecur :
-// CHECK: alloc_box ${ var Int64 }, var, name "x"
+// CHECK:         alloc_stack $Int64, var, name "x"
 // CHECK-LABEL: } // end sil function '$testrecur'
 sil [noinline] @$testrecur : $@convention(thin) () -> () {
 bb0:
@@ -275,6 +275,7 @@ bb0:
   return %10 : $()
 }
 
+// CHECK-LABEL: sil private [noinline] @$s13$testrecurbarTf0s_n : $@convention(thin) (@inout_aliasable Int64) -> () {
 sil private [noinline] @$testrecurbar : $@convention(thin) (@guaranteed { var Int64 }) -> () {
 bb0(%0 : ${ var Int64 }):
   %1 = project_box %0 : ${ var Int64 }, 0
@@ -292,6 +293,7 @@ bb0(%0 : ${ var Int64 }):
   return %13 : $()
 }
 
+// CHECK-LABEL: sil private [noinline] @$s13$testrecurbasTf0s_n : $@convention(thin) (@inout_aliasable Int64) -> () {
 sil private [noinline] @$testrecurbas : $@convention(thin) (@guaranteed { var Int64 }) -> () {
 bb0(%0 : ${ var Int64 }):
   %1 = project_box %0 : ${ var Int64 }, 0

test/SILOptimizer/allocboxtostack_escape.swift

@@ -103,7 +103,7 @@ public func testboxescapes() -> (() -> ()) {
 }
 
 // CHECK-LABEL: sil [noinline] @$s26allocboxtostack_localapply9testrecurSiyF :
-// CHECK: alloc_box ${ var Int }, var, name "x"
+// CHECK:          alloc_stack [var_decl] $Int, var, name "x"
 // CHECK-LABEL: } // end sil function '$s26allocboxtostack_localapply9testrecurSiyF'
 @inline(never)
 public func testrecur() -> Int {
@@ -146,14 +146,9 @@ public func testdfs1() -> Int {
   return common()
 }
 
-// Test to make sure we don't optimize the case when we have an inner common function call for multiple boxes.
-// We don't optimize this case now, because we don't have additional logic to correctly construct AppliesToSpecialize
-// Order of function calls constructed in PromotedOperands: bar innercommon local1 bas innercommon local2
-// AppliesToSpecialize should have the order: bar bas innercommon local1 local2
-// Since we don't maintain any tree like data structure with more info on the call tree, this is not possible to construct today 
 // CHECK-LABEL: sil [noinline] @$s26allocboxtostack_localapply8testdfs2SiyF :
-// CHECK: alloc_box ${ var Int }, var, name "x"
-// CHECK: alloc_box ${ var Int }, var, name "y"
+// CHECK:          alloc_stack [var_decl] $Int, var, name "x"
+// CHECK:          alloc_stack [var_decl] $Int, var, name "y"
 // CHECK-LABEL:} // end sil function '$s26allocboxtostack_localapply8testdfs2SiyF'
 @inline(never)
 public func testdfs2() -> Int {
@@ -182,3 +177,20 @@ public func testdfs2() -> Int {
   return local1() + local2()
 }
 
+// CHECK-LABEL: sil @$s26allocboxtostack_localapply15call2localfuncsSiyF :
+// CHECK-NOT:     alloc_box
+// CHECK-LABEL:} // end sil function '$s26allocboxtostack_localapply15call2localfuncsSiyF'
+public func call2localfuncs() -> Int {
+    var a1 = 1
+    
+    @inline(never)
+    func innerFunction() {
+        a1 += 1
+    }
+
+    innerFunction()
+    innerFunction()
+
+    return a1
+}
+