Skip to content

Commit d78f064

Browse files
eecksteineeckstein
committed
re-implement the StackPromotion pass in swift
It uses the new EscapeInfo.
1 parent da4ae95 commit d78f064

File tree

7 files changed

+730
-25
lines changed

7 files changed

+730
-25
lines changed

‎SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,5 @@ swift_compiler_sources(Optimizer
1515
RangeDumper.swift
1616
ReleaseDevirtualizer.swift
1717
RunUnitTests.swift
18+
StackPromotion.swift
1819
)

‎SwiftCompilerSources/Sources/Optimizer/FunctionPasses/StackPromotion.swift

Lines changed: 282 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,282 @@
1+
//===--- StackPromotion.swift - Stack promotion optimization --------------===//
2+
//
3+
// This source file is part of the Swift.org open source project
4+
//
5+
// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
import SIL
14+
15+
/// Promotes heap allocated objects to the stack.
16+
///
17+
/// It handles `alloc_ref` and `alloc_ref_dynamic` instructions of native swift
18+
/// classes: if promoted, the `[stack]` attribute is set in the allocation
19+
/// instruction and a `dealloc_stack_ref` is inserted at the end of the object's
20+
/// lifetime.
21+
22+
/// The main criteria for stack promotion is that the allocated object must not
23+
/// escape its function.
24+
///
25+
/// Example:
26+
/// %k = alloc_ref $Klass
27+
/// // .. some uses of %k
28+
/// destroy_value %k // The end of %k's lifetime
29+
///
30+
/// is transformed to:
31+
///
32+
/// %k = alloc_ref [stack] $Klass
33+
/// // .. some uses of %k
34+
/// destroy_value %k
35+
/// dealloc_stack_ref %k
36+
///
37+
/// The destroy/release of the promoted object remains in the SIL, but is effectively
38+
/// a no-op, because a stack promoted object is initialized with an "immortal"
39+
/// reference count.
40+
/// Later optimizations can clean that up.
41+
let stackPromotion = FunctionPass(name: "stack-promotion", {
42+
(function: Function, context: PassContext) in
43+
44+
var escapeInfo = EscapeInfo(calleeAnalysis: context.calleeAnalysis)
45+
let deadEndBlocks = context.deadEndBlocks
46+
47+
var changed = false
48+
for block in function.blocks {
49+
for inst in block.instructions {
50+
if let ar = inst as? AllocRefInstBase {
51+
if deadEndBlocks.isDeadEnd(block) {
52+
// Don't stack promote any allocation inside a code region which ends up
53+
// in a no-return block. Such allocations may missing their final release.
54+
// We would insert the deallocation too early, which may result in a
55+
// use-after-free problem.
56+
continue
57+
}
58+
59+
if tryPromoteAlloc(ar, &escapeInfo, deadEndBlocks, context) {
60+
changed = true
61+
}
62+
}
63+
}
64+
}
65+
if changed {
66+
// Make sure that all stack allocating instructions are nested correctly.
67+
context.fixStackNesting(function: function)
68+
}
69+
})
70+
71+
private
72+
func tryPromoteAlloc(_ allocRef: AllocRefInstBase,
73+
_ escapeInfo: inout EscapeInfo,
74+
_ deadEndBlocks: DeadEndBlocksAnalysis,
75+
_ context: PassContext) -> Bool {
76+
if allocRef.isObjC || allocRef.canAllocOnStack {
77+
return false
78+
}
79+
80+
// The most important check: does the object escape the current function?
81+
if escapeInfo.isEscaping(object:allocRef) {
82+
return false
83+
}
84+
85+
// Try to find the top most dominator block which dominates all use points.
86+
// * This block can be located "earlier" than the actual allocation block, in case the
87+
// promoted object is stored into an "outer" object, e.g.
88+
//
89+
// bb0: // outerDominatingBlock _
90+
// %o = alloc_ref $Outer |
91+
// ... |
92+
// bb1: // allocation block _ |
93+
// %k = alloc_ref $Klass | | "outer"
94+
// %f = ref_element_addr %o, #Outer.f | "inner" | liferange
95+
// store %k to %f | liferange |
96+
// ... | |
97+
// destroy_value %o _| _|
98+
//
99+
// * Finding the `outerDominatingBlock` is not guaranteed to work.
100+
// In this example, the top most dominator block is `bb0`, but `bb0` has no
101+
// use points in the outer liferange. We'll get `bb3` as outerDominatingBlock.
102+
// This is no problem because 1. it's an unusual case and 2. the `outerBlockRange`
103+
// is invalid in this case and we'll bail later.
104+
//
105+
// bb0: // real top most dominating block
106+
// cond_br %c, bb1, bb2
107+
// bb1:
108+
// %o1 = alloc_ref $Outer
109+
// br bb3(%o1)
110+
// bb2:
111+
// %o2 = alloc_ref $Outer
112+
// br bb3(%o1)
113+
// bb3(%o): // resulting outerDominatingBlock: wrong!
114+
// %k = alloc_ref $Klass
115+
// %f = ref_element_addr %o, #Outer.f
116+
// store %k to %f
117+
// destroy_value %o
118+
//
119+
let domTree = context.dominatorTree
120+
let outerDominatingBlock = escapeInfo.getDominatingBlockOfAllUsePoints(allocRef, domTree: domTree)
121+
122+
// The "inner" liferange contains all use points which are dominated by the allocation block
123+
var innerRange = InstructionRange(begin: allocRef, context)
124+
defer { innerRange.deinitialize() }
125+
126+
// The "outer" liferange contains all use points.
127+
var outerBlockRange = BasicBlockRange(begin: outerDominatingBlock, context)
128+
defer { outerBlockRange.deinitialize() }
129+
130+
escapeInfo.computeInnerAndOuterLiferanges(&innerRange, &outerBlockRange, domTree: domTree)
131+
132+
precondition(innerRange.blockRange.isValid, "inner range should be valid because we did a dominance check")
133+
134+
if !outerBlockRange.isValid {
135+
// This happens if we fail to find a correct outerDominatingBlock.
136+
return false
137+
}
138+
139+
// Check if there is a control flow edge from the inner to the outer liferange, which
140+
// would mean that the promoted object can escape to the outer liferange.
141+
// This can e.g. be the case if the inner liferange does not post dominate the outer range:
142+
// _
143+
// %o = alloc_ref $Outer |
144+
// cond_br %c, bb1, bb2 |
145+
// bb1: _ |
146+
// %k = alloc_ref $Klass | | outer
147+
// %f = ref_element_addr %o, #Outer.f | inner | range
148+
// store %k to %f | range |
149+
// br bb2 // branch from inner to outer _| |
150+
// bb2: |
151+
// destroy_value %o _|
152+
//
153+
// Or if there is a loop with a back-edge from the inner to the outer range:
154+
// _
155+
// %o = alloc_ref $Outer |
156+
// br bb1 |
157+
// bb1: _ |
158+
// %k = alloc_ref $Klass | | outer
159+
// %f = ref_element_addr %o, #Outer.f | inner | range
160+
// store %k to %f | range |
161+
// cond_br %c, bb1, bb2 // inner -> outer _| |
162+
// bb2: |
163+
// destroy_value %o _|
164+
//
165+
if innerRange.blockRange.isControlFlowEdge(to: outerBlockRange) {
166+
return false
167+
}
168+
169+
// There shouldn't be any critical exit edges from the liferange, because that would mean
170+
// that the promoted allocation is leaking.
171+
// Just to be on the safe side, do a check and bail if we find critical exit edges: we
172+
// cannot insert instructions on critical edges.
173+
if innerRange.blockRange.containsCriticalExitEdges(deadEndBlocks: deadEndBlocks) {
174+
return false
175+
}
176+
177+
// Do the transformation!
178+
// Insert `dealloc_stack_ref` instructions at the exit- and end-points of the inner liferange.
179+
for exitInst in innerRange.exits {
180+
if !deadEndBlocks.isDeadEnd(exitInst.block) {
181+
let builder = Builder(at: exitInst, context)
182+
_ = builder.createDeallocStackRef(allocRef)
183+
}
184+
}
185+
186+
for endInst in innerRange.ends {
187+
Builder.insert(after: endInst, location: allocRef.location, context) {
188+
(builder) in _ = builder.createDeallocStackRef(allocRef)
189+
}
190+
}
191+
192+
allocRef.setIsStackAllocatable(context)
193+
return true
194+
}
195+
196+
//===----------------------------------------------------------------------===//
197+
// utility functions
198+
//===----------------------------------------------------------------------===//
199+
200+
private extension EscapeInfo {
201+
mutating func getDominatingBlockOfAllUsePoints(_ value: SingleValueInstruction,
202+
domTree: DominatorTree) -> BasicBlock {
203+
// The starting point
204+
var dominatingBlock = value.block
205+
206+
_ = isEscaping(object: value, visitUse: { op, _, _ in
207+
if let defBlock = op.value.definingBlock, defBlock.dominates(dominatingBlock, domTree) {
208+
dominatingBlock = defBlock
209+
}
210+
return .continueWalking
211+
})
212+
return dominatingBlock
213+
}
214+
215+
/// All uses which are dominated by the `innerInstRange`s begin-block are included
216+
/// in both, the `innerInstRange` and the `outerBlockRange`.
217+
/// All _not_ dominated uses are only included in the `outerBlockRange`.
218+
mutating
219+
func computeInnerAndOuterLiferanges(_ innerInstRange: inout InstructionRange,
220+
_ outerBlockRange: inout BasicBlockRange,
221+
domTree: DominatorTree) {
222+
_ = isEscaping(object: innerInstRange.begin as! SingleValueInstruction, visitUse: { op, _, _ in
223+
let user = op.instruction
224+
if innerInstRange.blockRange.begin.dominates(user.block, domTree) {
225+
innerInstRange.insert(user)
226+
}
227+
outerBlockRange.insert(user.block)
228+
229+
let val = op.value
230+
if let defBlock = val.definingBlock {
231+
// Also insert the operand's definition. Otherwise we would miss allocation
232+
// instructions (for which the `visitUse` closure is not called).
233+
outerBlockRange.insert(defBlock)
234+
235+
// We need to explicitly add predecessor blocks of phi-arguments becaues they
236+
// are not necesesarily visited by `EscapeInfo.walkDown`.
237+
// This is important for the special case where there is a back-edge from the
238+
// inner range to the inner rage's begin-block:
239+
//
240+
// bb0: // <- need to be in the outer range
241+
// br bb1(%some_init_val)
242+
// bb1(%arg):
243+
// %k = alloc_ref $Klass // innerInstRange.begin
244+
// cond_br bb2, bb1(%k) // back-edge to bb1 == innerInstRange.blockRange.begin
245+
//
246+
if val is BlockArgument {
247+
outerBlockRange.insert(contentsOf: defBlock.predecessors)
248+
}
249+
}
250+
return .continueWalking
251+
})
252+
}
253+
}
254+
255+
private extension BasicBlockRange {
256+
/// Returns true if there is a direct edge connecting this range with the `otherRange`.
257+
func isControlFlowEdge(to otherRange: BasicBlockRange) -> Bool {
258+
func isOnlyInOtherRange(_ block: BasicBlock) -> Bool {
259+
return !inclusiveRangeContains(block) &&
260+
otherRange.inclusiveRangeContains(block) && block != otherRange.begin
261+
}
262+
263+
for lifeBlock in inclusiveRange {
264+
precondition(otherRange.inclusiveRangeContains(lifeBlock), "range must be a subset of other range")
265+
for succ in lifeBlock.successors {
266+
if isOnlyInOtherRange(succ) {
267+
return true
268+
}
269+
// The entry of the begin-block is conceptually not part of the range. We can check if
270+
// it's part of the `otherRange` by checking the begin-block's predecessors.
271+
if succ == begin && begin.predecessors.contains(where: { isOnlyInOtherRange($0) }) {
272+
return true
273+
}
274+
}
275+
}
276+
return false
277+
}
278+
279+
func containsCriticalExitEdges(deadEndBlocks: DeadEndBlocksAnalysis) -> Bool {
280+
exits.contains { !deadEndBlocks.isDeadEnd($0) && !$0.hasSinglePredecessor }
281+
}
282+
}

‎SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ private func registerSwiftPasses() {
4949
registerPass(escapeInfoDumper, { escapeInfoDumper.run($0) })
5050
registerPass(addressEscapeInfoDumper, { addressEscapeInfoDumper.run($0) })
5151
registerPass(computeEffects, { computeEffects.run($0) })
52+
registerPass(stackPromotion, { stackPromotion.run($0) })
5253
registerPass(simplifyBeginCOWMutationPass, { simplifyBeginCOWMutationPass.run($0) })
5354
registerPass(simplifyGlobalValuePass, { simplifyGlobalValuePass.run($0) })
5455
registerPass(simplifyStrongRetainPass, { simplifyStrongRetainPass.run($0) })

‎include/swift/SILOptimizer/PassManager/Passes.def

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -396,7 +396,7 @@ PASS(SplitAllCriticalEdges, "split-critical-edges",
396396
"Split all Critical Edges in the SIL CFG")
397397
PASS(SplitNonCondBrCriticalEdges, "split-non-cond_br-critical-edges",
398398
"Split all Critical Edges not from SIL cond_br")
399-
PASS(StackPromotion, "stack-promotion",
399+
SWIFT_FUNCTION_PASS_WITH_LEGACY(StackPromotion, "stack-promotion",
400400
"Stack Promotion of Class Objects")
401401
PASS(StripDebugInfo, "strip-debug-info",
402402
"Strip Debug Information")

‎lib/SILOptimizer/Transforms/StackPromotion.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,6 @@ bool StackPromotion::tryPromoteAlloc(AllocRefInstBase *ARI, EscapeAnalysis *EA,
158158

159159
} // end anonymous namespace
160160

161-
SILTransform *swift::createStackPromotion() {
161+
SILTransform *swift::createLegacyStackPromotion() {
162162
return new StackPromotion();
163163
}

0 commit comments

Comments
 (0)