Optimizer: reimplement the StackPromotion in swift and use EscapeInfo there

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

@@ -15,4 +15,5 @@ swift_compiler_sources(Optimizer
   RangeDumper.swift
   ReleaseDevirtualizer.swift
   RunUnitTests.swift
+  StackPromotion.swift
 )

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/StackPromotion.swift

@@ -0,0 +1,285 @@
+//===--- StackPromotion.swift - Stack promotion optimization --------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SIL
+
+/// Promotes heap allocated objects to the stack.
+///
+/// It handles `alloc_ref` and `alloc_ref_dynamic` instructions of native swift
+/// classes: if promoted, the `[stack]` attribute is set in the allocation
+/// instruction and a `dealloc_stack_ref` is inserted at the end of the object's
+/// lifetime.
+
+/// The main criteria for stack promotion is that the allocated object must not
+/// escape its function.
+///
+/// Example:
+///   %k = alloc_ref $Klass
+///   // .. some uses of %k
+///   destroy_value %k  // The end of %k's lifetime
+///
+/// is transformed to:
+///
+///   %k = alloc_ref [stack] $Klass
+///   // .. some uses of %k
+///   destroy_value %k
+///   dealloc_stack_ref %k
+///
+/// The destroy/release of the promoted object remains in the SIL, but is effectively
+/// a no-op, because a stack promoted object is initialized with an "immortal"
+/// reference count.
+/// Later optimizations can clean that up.
+let stackPromotion = FunctionPass(name: "stack-promotion", {
+  (function: Function, context: PassContext) in
+
+  var escapeInfo = EscapeInfo(calleeAnalysis: context.calleeAnalysis)
+  let deadEndBlocks = context.deadEndBlocks
+
+  var changed = false
+  for block in function.blocks {
+    for inst in block.instructions {
+      if let ar = inst as? AllocRefInstBase {
+        if deadEndBlocks.isDeadEnd(block) {
+          // 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: ar) {
+          break
+        }
+
+        if tryPromoteAlloc(ar, &escapeInfo, deadEndBlocks, context) {
+          changed = true
+        }
+      }
+    }
+  }
+  if changed {
+    // Make sure that all stack allocating instructions are nested correctly.
+    context.fixStackNesting(function: function)
+  }
+})
+
+private
+func tryPromoteAlloc(_ allocRef: AllocRefInstBase,
+                     _ escapeInfo: inout EscapeInfo,
+                     _ deadEndBlocks: DeadEndBlocksAnalysis,
+                     _ context: PassContext) -> Bool {
+  if allocRef.isObjC || allocRef.canAllocOnStack {
+    return false
+  }
+
+  // The most important check: does the object escape the current function?
+  if escapeInfo.isEscaping(object:allocRef) {
+    return false
+  }
+
+  // Try to find the top most dominator block which dominates all use points.
+  // * This block can be located "earlier" than the actual allocation block, in case the
+  //   promoted object is stored into an "outer" object, e.g.
+  //
+  //   bb0:           // outerDominatingBlock                 _
+  //     %o = alloc_ref $Outer                                 |
+  //     ...                                                   |
+  //   bb1:           // allocation block      _               |
+  //     %k = alloc_ref $Klass                  |              | "outer"
+  //     %f = ref_element_addr %o, #Outer.f     | "inner"      | liferange
+  //     store %k to %f                         | liferange    |
+  //     ...                                    |              |
+  //     destroy_value %o                      _|             _|
+  //
+  // * Finding the `outerDominatingBlock` is not guaranteed to work.
+  //   In this example, the top most dominator block is `bb0`, but `bb0` has no
+  //   use points in the outer liferange. We'll get `bb3` as outerDominatingBlock.
+  //   This is no problem because 1. it's an unusual case and 2. the `outerBlockRange`
+  //   is invalid in this case and we'll bail later.
+  //
+  //   bb0:                    // real top most dominating block
+  //     cond_br %c, bb1, bb2
+  //   bb1:
+  //     %o1 = alloc_ref $Outer
+  //     br bb3(%o1)
+  //   bb2:
+  //     %o2 = alloc_ref $Outer
+  //     br bb3(%o1)
+  //   bb3(%o):                // resulting outerDominatingBlock: wrong!
+  //     %k = alloc_ref $Klass
+  //     %f = ref_element_addr %o, #Outer.f
+  //     store %k to %f
+  //     destroy_value %o
+  //
+  let domTree = context.dominatorTree
+  let outerDominatingBlock = escapeInfo.getDominatingBlockOfAllUsePoints(allocRef, domTree: domTree)
+
+  // The "inner" liferange contains all use points which are dominated by the allocation block
+  var innerRange = InstructionRange(begin: allocRef, context)
+  defer { innerRange.deinitialize() }
+
+  // The "outer" liferange contains all use points.
+  var outerBlockRange = BasicBlockRange(begin: outerDominatingBlock, context)
+  defer { outerBlockRange.deinitialize() }
+
+  escapeInfo.computeInnerAndOuterLiferanges(&innerRange, &outerBlockRange, domTree: domTree)
+
+  precondition(innerRange.blockRange.isValid, "inner range should be valid because we did a dominance check")
+
+  if !outerBlockRange.isValid {
+    // This happens if we fail to find a correct outerDominatingBlock.
+    return false
+  }
+
+  // Check if there is a control flow edge from the inner to the outer liferange, which
+  // would mean that the promoted object can escape to the outer liferange.
+  // This can e.g. be the case if the inner liferange does not post dominate the outer range:
+  //                                                              _
+  //     %o = alloc_ref $Outer                                     |
+  //     cond_br %c, bb1, bb2                                      |
+  //   bb1:                                            _           |
+  //     %k = alloc_ref $Klass                          |          | outer
+  //     %f = ref_element_addr %o, #Outer.f             | inner    | range
+  //     store %k to %f                                 | range    |
+  //     br bb2       // branch from inner to outer    _|          |
+  //   bb2:                                                        |
+  //     destroy_value %o                                         _|
+  //
+  // Or if there is a loop with a back-edge from the inner to the outer range:
+  //                                                              _
+  //     %o = alloc_ref $Outer                                     |
+  //     br bb1                                                    |
+  //   bb1:                                            _           |
+  //     %k = alloc_ref $Klass                          |          | outer
+  //     %f = ref_element_addr %o, #Outer.f             | inner    | range
+  //     store %k to %f                                 | range    |
+  //     cond_br %c, bb1, bb2   // inner -> outer      _|          |
+  //   bb2:                                                        |
+  //     destroy_value %o                                         _|
+  //
+  if innerRange.blockRange.isControlFlowEdge(to: outerBlockRange) {
+    return false
+  }
+
+  // There shouldn't be any critical exit edges from the liferange, because that would mean
+  // that the promoted allocation is leaking.
+  // Just to be on the safe side, do a check and bail if we find critical exit edges: we
+  // cannot insert instructions on critical edges.
+  if innerRange.blockRange.containsCriticalExitEdges(deadEndBlocks: deadEndBlocks) {
+    return false
+  }
+
+  // Do the transformation!
+  // Insert `dealloc_stack_ref` instructions at the exit- and end-points of the inner liferange.
+  for exitInst in innerRange.exits {
+    if !deadEndBlocks.isDeadEnd(exitInst.block) {
+      let builder = Builder(at: exitInst, context)
+      _ = builder.createDeallocStackRef(allocRef)
+    }
+  }
+
+  for endInst in innerRange.ends {
+    Builder.insert(after: endInst, location: allocRef.location, context) {
+      (builder) in _ = builder.createDeallocStackRef(allocRef)
+    }
+  }
+
+  allocRef.setIsStackAllocatable(context)
+  return true
+}
+
+//===----------------------------------------------------------------------===//
+//                              utility functions
+//===----------------------------------------------------------------------===//
+
+private extension EscapeInfo {
+  mutating func getDominatingBlockOfAllUsePoints(_ value: SingleValueInstruction,
+                                                 domTree: DominatorTree) -> BasicBlock {
+    // The starting point
+    var dominatingBlock = value.block
+
+    _ = isEscaping(object: value, visitUse: { op, _, _ in
+        if let defBlock = op.value.definingBlock, defBlock.dominates(dominatingBlock, domTree) {
+          dominatingBlock = defBlock
+        }
+        return .continueWalking
+      })
+    return dominatingBlock
+  }
+
+  /// All uses which are dominated by the `innerInstRange`s begin-block are included
+  /// in both, the `innerInstRange` and the `outerBlockRange`.
+  /// All _not_ dominated uses are only included in the `outerBlockRange`.
+  mutating
+  func computeInnerAndOuterLiferanges(_ innerInstRange: inout InstructionRange,
+                                      _ outerBlockRange: inout BasicBlockRange,
+                                      domTree: DominatorTree) {
+    _ = isEscaping(object: innerInstRange.begin as! SingleValueInstruction, visitUse: { op, _, _ in
+        let user = op.instruction
+        if innerInstRange.blockRange.begin.dominates(user.block, domTree) {
+          innerInstRange.insert(user)
+        }
+        outerBlockRange.insert(user.block)
+
+        let val = op.value
+        if let defBlock = val.definingBlock {
+          // Also insert the operand's definition. Otherwise we would miss allocation
+          // instructions (for which the `visitUse` closure is not called).
+          outerBlockRange.insert(defBlock)
+
+          // We need to explicitly add predecessor blocks of phi-arguments becaues they
+          // are not necesesarily visited by `EscapeInfo.walkDown`.
+          // This is important for the special case where there is a back-edge from the
+          // inner range to the inner rage's begin-block:
+          //
+          //   bb0:                      // <- need to be in the outer range
+          //     br bb1(%some_init_val)
+          //   bb1(%arg):
+          //     %k = alloc_ref $Klass   // innerInstRange.begin
+          //     cond_br bb2, bb1(%k)    // back-edge to bb1 == innerInstRange.blockRange.begin
+          //
+          if val is BlockArgument {
+            outerBlockRange.insert(contentsOf: defBlock.predecessors)
+          }
+        }
+        return .continueWalking
+      })
+  }
+}
+
+private extension BasicBlockRange {
+  /// Returns true if there is a direct edge connecting this range with the `otherRange`.
+  func isControlFlowEdge(to otherRange: BasicBlockRange) -> Bool {
+    func isOnlyInOtherRange(_ block: BasicBlock) -> Bool {
+      return !inclusiveRangeContains(block) &&
+             otherRange.inclusiveRangeContains(block) && block != otherRange.begin
+    }
+
+    for lifeBlock in inclusiveRange {
+      precondition(otherRange.inclusiveRangeContains(lifeBlock), "range must be a subset of other range")
+      for succ in lifeBlock.successors {
+        if isOnlyInOtherRange(succ) {
+          return true
+        }
+        // The entry of the begin-block is conceptually not part of the range. We can check if
+        // it's part of the `otherRange` by checking the begin-block's predecessors.
+        if succ == begin && begin.predecessors.contains(where: { isOnlyInOtherRange($0) }) {
+          return true
+        }
+      }
+    }
+    return false
+  }
+
+  func containsCriticalExitEdges(deadEndBlocks: DeadEndBlocksAnalysis) -> Bool {
+    exits.contains { !deadEndBlocks.isDeadEnd($0) && !$0.hasSinglePredecessor }
+  }
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -43,6 +43,7 @@ private func registerSwiftPasses() {
   registerPass(escapeInfoDumper, { escapeInfoDumper.run($0) })
   registerPass(addressEscapeInfoDumper, { addressEscapeInfoDumper.run($0) })
   registerPass(computeEffects, { computeEffects.run($0) })
+  registerPass(stackPromotion, { stackPromotion.run($0) })
   registerPass(simplifyBeginCOWMutationPass, { simplifyBeginCOWMutationPass.run($0) })
   registerPass(simplifyGlobalValuePass, { simplifyGlobalValuePass.run($0) })
   registerPass(simplifyStrongRetainPass, { simplifyStrongRetainPass.run($0) })

include/swift/SILOptimizer/PassManager/Passes.def

@@ -396,7 +396,7 @@ PASS(SplitAllCriticalEdges, "split-critical-edges",
      "Split all Critical Edges in the SIL CFG")
 PASS(SplitNonCondBrCriticalEdges, "split-non-cond_br-critical-edges",
      "Split all Critical Edges not from SIL cond_br")
-PASS(StackPromotion, "stack-promotion",
+SWIFT_FUNCTION_PASS(StackPromotion, "stack-promotion",
      "Stack Promotion of Class Objects")
 PASS(StripDebugInfo, "strip-debug-info",
      "Strip Debug Information")

lib/SILOptimizer/Transforms/CMakeLists.txt

@@ -39,7 +39,6 @@ target_sources(swiftSILOptimizer PRIVATE
   SimplifyCFG.cpp
   Sink.cpp
   SpeculativeDevirtualizer.cpp
-  StackPromotion.cpp
   StringOptimization.cpp
   TempLValueOpt.cpp
   TempRValueElimination.cpp