Merge pull request #66933 from eeckstein/stack-promotion

StackPromotion: support promoting allocations in dead-end control flow regions

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/StackPromotion.swift

@@ -46,17 +46,9 @@ let stackPromotion = FunctionPass(name: "stack-promotion") {
   var needFixStackNesting = false
   for inst in function.instructions {
     if let allocRef = inst as? AllocRefInstBase {
-      if deadEndBlocks.isDeadEnd(allocRef.parentBlock) {
-        // Don't stack promote any allocation inside a code region which ends up
-        // in a no-return block. Such allocations may missing their final release.
-        // We would insert the deallocation too early, which may result in a
-        // use-after-free problem.
-        continue
-      }
       if !context.continueWithNextSubpassRun(for: allocRef) {
         break
       }
-
       if tryPromoteAlloc(allocRef, deadEndBlocks, context) {
         needFixStackNesting = true
       }
@@ -75,6 +67,41 @@ private func tryPromoteAlloc(_ allocRef: AllocRefInstBase,
   if allocRef.isObjC || allocRef.canAllocOnStack {
     return false
   }
+
+  if deadEndBlocks.isDeadEnd(allocRef.parentBlock) {
+
+    // Allocations inside a code region which ends up in a no-return block may missing their
+    // final release. Therefore we extend their lifetime indefinitely, e.g.
+    //
+    //  %k = alloc_ref $Klass
+    //  ...
+    //  unreachable  // The end of %k's lifetime
+    //
+    // Also, such an allocation cannot escape the function, because the function does not
+    // return after the point of allocation. So we can stack-promote it unconditionally.
+    //
+    // There is one exception: if it's in a loop (within the dead-end region) we must not
+    // extend its lifetime. On the other hand we can be sure that its final release is not
+    // missing, because otherwise the object would be leaking. For example:
+    //
+    //  bb1:
+    //    %k = alloc_ref $Klass
+    //    ...                    // %k's lifetime must end somewhere here
+    //    cond_br %c, bb1, bb2
+    //  bb2:
+    //    unreachable
+    //
+    // Therefore, if the allocation is inside a loop, we can treat it like allocations in
+    // non dead-end regions.
+    if !isInLoop(block: allocRef.parentBlock, context),
+       // TODO: for some reason this doesn't work for aysnc functions.
+       //       Maybe a problem with co-routine splitting in LLVM?
+       !allocRef.parentFunction.isAsync {
+      allocRef.setIsStackAllocatable(context)
+      return true
+    }
+  }
+
   // The most important check: does the object escape the current function?
   if allocRef.isEscaping(context) {
     return false
@@ -293,3 +320,16 @@ private extension BasicBlockRange {
   }
 }
 
+private func isInLoop(block startBlock: BasicBlock, _ context: FunctionPassContext) -> Bool {
+  var worklist = BasicBlockWorklist(context)
+  defer { worklist.deinitialize() }
+
+  worklist.pushIfNotVisited(contentsOf: startBlock.successors)
+  while let block = worklist.pop() {
+    if block == startBlock {
+      return true
+    }
+    worklist.pushIfNotVisited(contentsOf: block.successors)
+  }
+  return false
+}

test/SILOptimizer/stack_promotion.sil

@@ -74,10 +74,14 @@ bb0:
   return %n1 : $XX
 }
 
-// CHECK-LABEL: sil @dont_promote_in_unreachable
-// CHECK: bb1:
-// CHECK-NEXT: alloc_ref $XX
-sil @dont_promote_in_unreachable : $@convention(thin) () -> () {
+// CHECK-LABEL: sil @promote_in_unreachable :
+// CHECK:       bb1:
+// CHECK-NEXT:    alloc_ref [stack] $XX
+// CHECK:         apply
+// CHECK:         apply
+// CHECK-NEXT:    unreachable
+// CHECK:       } // end sil function 'promote_in_unreachable'
+sil @promote_in_unreachable : $@convention(thin) () -> () {
 bb0:
   cond_br undef, bb1, bb2
 
@@ -97,6 +101,34 @@ bb2:
   return %r : $()
 }
 
+// CHECK-LABEL: sil @promote_in_unreachable_loop :
+// CHECK:       bb1:
+// CHECK-NEXT:    alloc_ref [stack] $XX
+// CHECK:         apply
+// CHECK:         apply
+// CHECK-NEXT:    strong_release
+// CHECK-NEXT:    dealloc_stack_ref
+// CHECK-NEXT:    cond_br
+// CHECK:       } // end sil function 'promote_in_unreachable_loop'
+sil @promote_in_unreachable_loop : $@convention(thin) () -> () {
+bb0:
+  br bb1
+
+bb1:
+  %o1 = alloc_ref $XX
+  %f1 = function_ref @xx_init : $@convention(thin) (@guaranteed XX) -> XX
+  %n1 = apply %f1(%o1) : $@convention(thin) (@guaranteed XX) -> XX
+  %l1 = ref_element_addr %n1 : $XX, #XX.x
+  %l2 = load %l1 : $*Int32
+  %f2 = function_ref @consume_int : $@convention(thin) (Int32) -> ()
+  %a = apply %f2(%l2) : $@convention(thin) (Int32) -> ()
+  strong_release %o1 : $XX
+  cond_br undef, bb1, bb2
+
+bb2:
+  unreachable
+}
+
 // CHECK-LABEL: sil @promote_nested
 // CHECK: [[X:%[0-9]+]] = alloc_ref [stack] $XX
 // CHECK: [[Y:%[0-9]+]] = alloc_ref [stack] $YY