Merge pull request #82728 from eeckstein/fix-allocbox-to-stack

MandatoryAllocBoxToStack: also handle new specialized functions and re-enable the pass again

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/DataStructures/Worklist.swift

@@ -125,7 +125,7 @@ struct FunctionWorklist {
     }
   }
 
-  mutating func pushIfNotVisited(contentsOf functions: [Function]) {
+  mutating func pushIfNotVisited<S: Sequence>(contentsOf functions: S) where S.Element == Function {
     for f in functions {
       pushIfNotVisited(f)
     }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -65,12 +65,16 @@ let allocBoxToStack = FunctionPass(name: "allocbox-to-stack") {
 let mandatoryAllocBoxToStack = ModulePass(name: "mandatory-allocbox-to-stack") {
   (moduleContext: ModulePassContext) in
 
+  var worklist = FunctionWorklist()
+  worklist.pushIfNotVisited(contentsOf: moduleContext.functions)
+
   var originalsOfSpecializedFunctions = FunctionWorklist()
 
-  for function in moduleContext.functions {
+  while let function = worklist.pop() {
     moduleContext.transform(function: function) { context in
       let specFns = tryConvertBoxesToStack(in: function, context)
-      originalsOfSpecializedFunctions.pushIfNotVisited(contentsOf: specFns)
+      worklist.pushIfNotVisited(contentsOf: specFns.specializedFunctions)
+      originalsOfSpecializedFunctions.pushIfNotVisited(contentsOf: specFns.originalFunctions)
     }
   }
 
@@ -82,7 +86,7 @@ let mandatoryAllocBoxToStack = ModulePass(name: "mandatory-allocbox-to-stack") {
 /// Converts all non-escaping `alloc_box` to `alloc_stack` and specializes called functions if a
 /// box is passed to a function.
 /// Returns the list of original functions for which a specialization has been created.
-private func tryConvertBoxesToStack(in function: Function, _ context: FunctionPassContext) -> [Function] {
+private func tryConvertBoxesToStack(in function: Function, _ context: FunctionPassContext) -> FunctionSpecializations {
   var promotableBoxes = Array<(AllocBoxInst, Flags)>()
   var functionsToSpecialize = FunctionSpecializations()
 
@@ -101,7 +105,7 @@ private func tryConvertBoxesToStack(in function: Function, _ context: FunctionPa
     function.fixStackNesting(context)
   }
 
-  return functionsToSpecialize.originalFunctions
+  return functionsToSpecialize
 }
 
 private func findPromotableBoxes(in function: Function,
@@ -187,6 +191,7 @@ private struct FunctionSpecializations {
   private var originalToSpecialized = Dictionary<Function, Function>()
 
   var originalFunctions: [Function] { originals.functions }
+  var specializedFunctions: [Function] { originals.functions.lazy.map { originalToSpecialized[$0]! } }
 
   mutating func add(promotableArguments: [FunctionArgument]) {
     for arg in promotableArguments {

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -116,13 +116,12 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   // Select access kind after capture promotion and before stack promotion.
   // This guarantees that stack-promotable boxes have [static] enforcement.
   P.addAccessEnforcementSelection();
-/* Temporarily disabled: rdar://154686063, rdar://154713388
+
 #ifdef SWIFT_ENABLE_SWIFT_IN_SWIFT
   P.addMandatoryAllocBoxToStack();
 #else
-*/
   P.addLegacyAllocBoxToStack();
-//#endif
+#endif
   P.addNoReturnFolding();
   P.addBooleanLiteralFolding();
   addDefiniteInitialization(P);
@@ -418,8 +417,7 @@ void addHighLevelLoopOptPasses(SILPassPipelinePlan &P) {
 void addFunctionPasses(SILPassPipelinePlan &P,
                        OptimizationLevelKind OpLevel) {
   // Promote box allocations to stack allocations.
-  // TODO: change this to add addAllocBoxToStack once rdar://154686063, rdar://154713388 is fixed
-  P.addLegacyAllocBoxToStack();
+  P.addAllocBoxToStack();
 
   if (P.getOptions().DestroyHoisting == DestroyHoistingOption::On) {
     P.addDestroyAddrHoisting();

test/IRGen/generic_tuples.swift

@@ -27,9 +27,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/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -1198,3 +1198,36 @@ bb0(%0 : $*T, %1 : $*T, %2 : $*T):
   %15 = tuple ()
   return %15 : $()
 }
+
+// CHECK-LABEL: sil [ossa] @alloc_box_in_specialized_callee :
+// CHECK-NOT:     alloc_box
+// CHECK-LABEL: } // end sil function 'alloc_box_in_specialized_callee'
+sil [ossa] @alloc_box_in_specialized_callee : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  %1 = alloc_box ${ var Int }
+  %2 = project_box %1, 0
+  store %0 to [trivial] %2
+  %4 = function_ref @callee_with_allocbox : $@convention(thin) (@guaranteed { var Int }) -> ()
+  %5 = apply %4(%1) : $@convention(thin) (@guaranteed { var Int }) -> ()
+  destroy_value %1
+  %r = tuple ()
+  return %r
+}
+
+// CHECK-LABEL: sil shared [ossa] @$s20callee_with_allocboxTf0s_n :
+// CHECK-NOT:     alloc_box
+// CHECK-LABEL: } // end sil function '$s20callee_with_allocboxTf0s_n'
+
+// CHECK-LABEL: sil [ossa] @callee_with_allocbox :
+// CHECK-NOT:     alloc_box
+// CHECK-LABEL: } // end sil function 'callee_with_allocbox'
+sil [ossa] @callee_with_allocbox : $@convention(thin) (@guaranteed { var Int }) -> () {
+bb0(%0 : @guaranteed ${ var Int }):
+  %1 = alloc_box ${ var Int }
+  %2 = project_box %1 : ${ var Int }, 0
+  %3 = project_box %0 : ${ var Int }, 0
+  copy_addr %3 to %2
+  destroy_value %1
+  %r = tuple ()
+  return %r
+}

test/SILOptimizer/allocboxtostack_localapply.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
+}
+

test/SILOptimizer/definite_init_protocol_init.swift

@@ -99,7 +99,8 @@ struct AddressOnlyStruct : TriviallyConstructible {
 // CHECK-NEXT: apply [[FN]]<AddressOnlyStruct>([[SELF_BOX]], %1, [[METATYPE]])
 // CHECK-NEXT: copy_addr [take] [[SELF_BOX]] to [init] [[SELF]]
 // CHECK-NEXT: dealloc_stack [[SELF_BOX]]
-// CHECK-NEXT: copy_addr [take] [[SELF]] to [init] %0
+// CHECK-NEXT: copy_addr [[SELF]] to [init] %0
+// CHECK-NEXT: destroy_addr [[SELF]]
 // CHECK-NEXT: dealloc_stack [[SELF]]
 // CHECK-NEXT: [[RESULT:%.*]] = tuple ()
 // CHECK-NEXT: return [[RESULT]]