Files
swift-mirror/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AsyncDemotion.swift
2024-08-08 22:22:39 -07:00

724 lines
22 KiB
Swift

//===--- AsyncDemotion.swift -----------------------------------------------==//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2023 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
private struct Stats: CustomStringConvertible {
enum Kind {
case asyncFns // Functions analyzed.
case demotionCandidateFn // Functions that qualified for demotion.
case demotionPerformedFn // Functions that were demoted.
case asyncCallsCount // Calls analyzed.
case asyncKnownCallsCount // Calls to known async functions.
case asyncCallsReplaced // Calls replaced with call to sync demotion.
case asyncCallExecutorMismatch // Calls that could not be replaced due to mismatch.
case hopCount // Hop-to-executors analyzed.
case hopsRemoved // Hop-to-executors removed.
}
private var data: [Kind: Int] = [:]
mutating func tick(_ k: Kind, _ increment: Int = 1) {
data[k, default: 0] += increment
}
/// Produces a JSON description of all statistics
var description: String {
var s = "{\n"
var entries = data.map { " \"\($0.0)\": \($0.1)" }
entries.sort() // for consistency of output
s += entries.joined(separator: ",\n")
s += "\n}"
return s
}
}
fileprivate var stats: Stats = Stats()
enum Executor: CustomStringConvertible, Hashable {
case selfActor // self argument of this function
case nonisolated // the global, concurrent executor
case globalActor(NominalTypeDecl) // a global actor of the given type
case unknown(HashableValue)
var description: String {
switch self {
case .selfActor: return "self"
case .nonisolated: return "nonisolated"
case let .globalActor(nom): return "globalActor(\(nom.name))"
case let .unknown(value): return "unknown(\(value))"
}
}
init(_ op: Operand) {
self = .init(op.value)
}
init(_ val: Value) {
var finder = FindExecutor()
self = finder.search(val)
}
/// Same-function executor equality only.
static func == (_ a: Executor, _ b: Executor) -> Bool {
return same(duringCall: nil, from: a, to: b)
}
/// Returns true if the executors are the same.
///
/// If `duringCall` is `nil`, then the executor comparison
/// is assumed to be within the same function. Otherwise, the
/// callsite for the transfer of control is used to do matching
/// across functions.
static func same(duringCall: ApplySite?,
from: Executor,
to: Executor) -> Bool {
switch (from, to) {
case (.nonisolated, .nonisolated): return true
case let (.unknown(f), .unknown(t)): return f == t
case (.selfActor, .selfActor):
guard duringCall != nil else { return true }
fatalError("implement selfActor comparison logic")
case let (.globalActor(ty1), .globalActor(ty2)): return ty1 == ty2
default: return false
}
}
// Finds the executor corresponding to a hop instruction's operand.
struct FindExecutor : ValueUseDefWalker {
private var def: Executor?
internal var walkUpCache = WalkerCache<SmallProjectionPath>()
mutating func search(_ val: Value) -> Executor {
// assert(val.type == $Optional<Builtin.Executor>)
_ = walkUp(value: val, path: SmallProjectionPath(.anyValueFields))
guard let result = def else {
fatalError("could not find def for value:\n \(val)")
}
return result
}
internal mutating func rootDef(value: Value,
path: SmallProjectionPath) -> WalkResult {
// dump("FindExecutor.rootDef -> \(value)")
if let arg = value as? FunctionArgument {
def = arg.isSelf ? .selfActor : .unknown(value.hashable)
} else if let enm = value as? EnumInst {
assert(enm.payload == nil, "expected empty payload case for a def")
def = .nonisolated
} else if value.type.isNominal {
let nom = value.type.nominal
def = nom.isGlobalActor ? .globalActor(nom) : .unknown(value.hashable)
} else {
def = .unknown(value.hashable)
}
return .abortWalk
}
}
}
// A node represents a function in the demotion graph. It encapsulates all of
// the state of the function as it is analyzed and transformed.
private struct Node: CustomStringConvertible {
enum DemotionStatus: Equatable, CustomStringConvertible {
case prohibited(ProhibitionReason) // demotion is never permitted
case awaiting // demotion might be possible
case demoted(Function)
enum ProhibitionReason {
case ambiguousExecutor
case suspends
case external
}
var isDemoted: Bool {
switch self {
case .demoted(_): return true
default: return false
}
}
var description: String {
switch self {
case let .prohibited(reason): return "prohibited(\(reason))"
case .awaiting: return "awaiting"
case let .demoted(fn): return "demoted(\(fn.name))"
}
}
}
let function: Function
let executors: Set<Executor>
private(set) var demotionState: DemotionStatus
private(set) var asyncCalls: [ApplySite]
private(set) var hops: [Instruction]
init(_ analysis: AnalysisResult) {
let data = analysis.data
function = data.function
executors = data.executors
demotionState = analysis.demotionStatus
asyncCalls = data.knownAsyncCalls
hops = data.hops
assert(executors.count <= 1 || demotionState != .awaiting,
"could only be awaiting demotion if <1 executor")
}
var description: String {
return """
DemotionStatus {
function: \(function.name),
demotionState: \(demotionState),
executors: \(executors),
asyncCalls: \(asyncCalls.map { $0.referencedFunction!.name.description })
}
"""
}
var isDemoted: Bool { demotionState.isDemoted }
// Effectively the out-edges of this node.
// FIXME: incrementally update this as the async calls are rewritten
func getUniqueCallees() -> Set<Function> {
return .init(asyncCalls.map { $0.referencedFunction! })
}
/// Rewrites async calls into synchronous calls, if possible.
///
/// Returns true iff this function had a call rewritten.
mutating
func rewriteCalls(to callee: Node) -> Bool {
let oldEntry = callee.function
guard case let .demoted(newEntry) = callee.demotionState else {
fatalError("rewriting calls when callee is not demoted?")
}
// It is only safe to rewrite calls in this function of the target
// if the executors are compatible.
assert(callee.executors.count <= 1)
guard executors.count <= 1 else {
// TODO: once we have precise information about which call-sites will be
// encountered while on a specific executor, we can be less conservative.
return false
}
let before = asyncCalls.count
assert(before > 0, "no async calls to rewrite?")
asyncCalls.removeAll { site in
let target = site.referencedFunction!
guard target == oldEntry else { return false }
// When converting a call from async-to-sync, we must ensure that prior to
// entering the sync entry-point, that we're on an executor compatible with
// the original async callee.
switch (executors.first, callee.executors.first) {
// If both require an executor, they must be the same
case let (.some(from), .some(to)):
if Executor.same(duringCall: site, from: from, to: to) {
break
} else {
dump("executor mismatch:\n\tfrom:\(from)\n\tto:\(to)", .med)
stats.tick(.asyncCallExecutorMismatch)
return false
}
// The callee requires an executor, but as the caller we don't know
// which executor we run on!
case (.none, .some(_)):
stats.tick(.asyncCallExecutorMismatch)
return false
// Callee has no preference.
case (_, .none): break
}
dump("📝 \(function.name) | rewriting \(oldEntry.name) --> \(newEntry.name) @ \(site)")
// FIXME: actually implement the rewriting
return true
}
let callsReplaced = before - asyncCalls.count
stats.tick(.asyncCallsReplaced, callsReplaced)
return callsReplaced > 0
}
/// Attempts a simple, local demotion of the this function.
///
/// Returns true if demotion was performed.
mutating func tryDemotion(_ moduleContext: ModulePassContext) -> Bool {
assert(!demotionState.isDemoted, "already demoted!")
// Is it valid to demote this function?
guard case .awaiting = demotionState else { return false }
let uniqueCallees = getUniqueCallees()
// We can demote if there are only self-recursive async calls.
guard uniqueCallees.count <= 1 else { return false }
// If there is a callee, then it must be ourself.
// Otherwise we need non-local analysis.
if let callee = uniqueCallees.first,
callee != function {
return false
}
/// Demotion involves a few steps:
/// 1. Remove all `hop_to_executor`'s from the async body.
///
/// TODO: the below
///
/// - Create an empty non-async `Function` with an otherwise identical
/// signature.
/// - Transfer the body of the async function to the non-async function.
///
/// - Fill in the body of the async function with an executor hop followed
/// by a tail-call to the non-async function.
/// - Replace self-recursive async calls with non-async calls.
// dump("-------> \(moduleContext.mangleAsyncRemoved(from: function))\n")
moduleContext.transform(function: function) { context in
while let hop = hops.popLast() {
context.erase(instruction: hop)
stats.tick(.hopsRemoved)
}
}
// FIXME: temporary
demotionState = .demoted(function)
stats.tick(.demotionPerformedFn)
dump("✨ successful demotion: \(self)")
return true
}
}
/// The DemotionGraph is a subgraph of the entire program's call-graph that
/// only deals with functions that are candidates of async demotion. Two
/// connected nodes of the graph represent a possible demotion chain.
private struct DemotionGraph: CustomStringConvertible {
typealias StatusMap = [Function : Node]
private let context: ModulePassContext
// A map from function F to the demotion status of F.
private var status: StatusMap = [:]
// A map from function F to all functions that are known to call F.
private var knownCallers: [Function: Set<Function>] = [:]
// The set of demotable functions.
private var demotable: Set<Function> = .init()
init(_ context: ModulePassContext) { self.context = context}
mutating func addNode(_ ds: Node) {
let empty = Set<Function>()
assert(status.index(forKey: ds.function) == nil,
"adding function twice!")
status[ds.function] = ds
if case .awaiting = ds.demotionState {
demotable.insert(ds.function)
}
// cache which functions 'fn' calls.
for apply in ds.asyncCalls {
let callee = apply.referencedFunction!
knownCallers[callee, default: empty].insert(ds.function)
}
}
/// Generates a GraphViz-compatible representation of the graph.
func asGraphViz(_ title: String = "DemotionGraph") -> String {
var nameMap: [String: String] = [:]
var freshID: Int = 0
func nameOf(_ fn: Function) -> String {
let full = fn.name.string
if let shortName = nameMap[full] {
return shortName
}
let newShortName = "\(freshID)"
freshID += 1
nameMap[full] = newShortName
return newShortName
}
func colorOf(_ fn: Function) -> String {
switch status[fn]!.demotionState {
case .prohibited(.external): return "black"
case .prohibited(.suspends): return "red"
case .prohibited(.ambiguousExecutor): return "hotpink"
case .awaiting: return "darkviolet"
case .demoted: return "green"
}
}
var output = "\ndigraph \(title) {\n"
for fn in status.keys {
output += " \(nameOf(fn)) [color=\"\(colorOf(fn))\"] // \(fn.name)\n"
for callee in status[fn]!.getUniqueCallees() {
output += " \(nameOf(fn)) -> \(nameOf(callee))\n"
}
}
output += "}\n"
return output
}
var description : String {
return "\(status.values)\n\n" + asGraphViz()
}
/// Entry-point into the main optimization process.
mutating func optimize() {
/** Do an initial chain-demotion pass over all demotable functions while
building the SCC graph.
This will clean-up the call-graph so that chains of non-SCC function
calls are removed. The benefit is that a chain of demotable functions
that leave an SCC will be removed. For example, if we have a graph like:
```
digraph G {
// SCC 1
a -> b
b -> c
c -> a
// SCC 2
d -> d
x -> a
c -> d
// edges that will be removed by the initial chain-demotion
c -> y
y -> z
}
```
Then this initial pass over the nodes will demote `z` and `y`, removing
the edge between them. But more importantly, it will also remove the
edge `c -> y`, which is an edge that leaves SCC 1!
*/
for fn in demotable {
guard !status[fn]!.isDemoted else {
demotable.remove(fn)
continue
}
if status[fn]!.tryDemotion(context) {
demotable.remove(fn)
rewriteCallers(of: fn)
continue
}
}
/// Demote SCC's until we no longer can make any demotions.
///
/// Demotion an SCC is quite similar to demoting a single function. An SCC
/// can be demoted iff:
/// 1. There does not exist any edges leaving the SCC.
/// 2. The executors of all functions in the SCC are mutually compatible.
///
// TODO: implement
}
mutating func rewriteCallers(of demotedFn: Function) {
/// Rewriting a function's async calls can cause a chain of
/// demotions through the call graph. We can track our progress
/// through the chain using a simple stack, relying on the removal
/// of edges in `knownCallers` to ensure we don't loop forever.
var worklist = [demotedFn]
while let fn = worklist.popLast() {
let state = status[fn]!
assert(state.demotionState.isDemoted)
guard let callers = knownCallers.removeValue(forKey: fn) else {
continue // no callers to rewrite
}
for caller in callers {
let changed = status[caller]!.rewriteCalls(to: state)
// Skip if we re-wrote the self-recursive calls.
guard caller != demotedFn else { continue }
// If rewriting changed the caller, chain a demotion attempt.
if changed && status[caller]!.tryDemotion(context) {
worklist.append(caller)
}
}
}
}
} // end struct
/// Embodies the result of analyzing a Function to determine whether
/// it is a valid candidate for async demotion. If it is not a valid
/// candidate, then the reason for rejection is indicated.
private enum AnalysisResult: CustomStringConvertible {
typealias Data = (function: Function,
knownAsyncCalls: [ApplySite],
executors: Set<Executor>,
hops: [Instruction])
case demotionCandidate(Data)
case externalFunction(Data)
case multipleExecutors(Data)
case unknownAsyncOp(Data, Instruction) // an unanalyzed suspending op
var data: Data {
switch self {
case let .demotionCandidate(data): return data
case let .externalFunction(data): return data
case let .multipleExecutors(data): return data
case let .unknownAsyncOp(data, _): return data
}
}
var demotionStatus: Node.DemotionStatus {
switch self {
case .demotionCandidate: return .awaiting
case .unknownAsyncOp: return .prohibited(.suspends)
case .multipleExecutors: return .prohibited(.ambiguousExecutor)
case .externalFunction: return .prohibited(.external)
}
}
init(_ data: Data,
_ unknownAsyncOp: Instruction?) {
assert(data.knownAsyncCalls.allSatisfy {
$0.callee.type.isAsyncFunction
&& $0.parentFunction == data.function })
assert(data.hops.allSatisfy { $0.parentFunction == data.function })
// Can't demote or analyze external async functions.
if data.function.blocks.isEmpty {
assert(data.function.isAsync)
self = .externalFunction(data)
// Demotion to non-async is only ever possible if the function hops to at
// most one executor.
} else if data.executors.count > 1 {
self = .multipleExecutors(data)
// Next, if there are any unknown async operations we also cannot demote.
} else if let op = unknownAsyncOp {
self = .unknownAsyncOp(data, op)
} else {
self = .demotionCandidate(data)
stats.tick(.demotionCandidateFn)
}
}
var description: String {
switch self {
case let .multipleExecutors(base):
return "\(base.0.name) -- ❌ multiple executors \(base.2)"
case let .unknownAsyncOp(base, inst):
return "\(base.0.name) -- ❌ unknown async operation: \(inst) | \(base.2)"
case let .externalFunction(base):
return "\(base.0.name) -- ❌ external function"
case let .demotionCandidate(base):
return "\(base.0.name) -- ✅ demotion candidate \(base.2)"
}
}
}
extension Optional {
mutating func setIfUnset(_ value: @autoclosure () -> Wrapped) {
if case .none = self {
self = .some(value())
}
}
}
/// Analyzes the kinds of possible suspension points and determines whether
/// the function is locally a candidate for async demotion.
///
/// Returns the initial demotion status for this function if it is async.
private func analyzeForDemotion(_ function: Function) -> Node? {
// must be async
guard function.isAsync else { return nil }
stats.tick(.asyncFns)
// We want to collect all of the ApplySites in this analysis,
// even if the async function isn't a candidate for demotion itself,
// because it may be a caller of a function that can be demoted.
var executors = Set<Executor>()
var knownAsyncApplySites: [ApplySite] = []
var hops: [Instruction] = []
var unknownAsyncOp: Instruction? = nil
for inst in function.instructions {
// Track hops in a low-resolution way.
//
// FIXME: would be better to associate each call-site with the executor
// it would happen under, so demoted entry-points can be called from
// multi-executor async functions.
if let hop = inst as? HopToExecutorInst {
executors.insert(Executor(hop.operand))
hops.append(hop)
stats.tick(.hopCount)
} else if let apply = inst as? ApplySite {
// Only record async call-sites.
guard apply.isAsync else { continue }
stats.tick(.asyncCallsCount)
// Treat an unknown async callee as outside of our analysis
if apply.referencedFunction == nil {
unknownAsyncOp.setIfUnset(inst)
} else {
knownAsyncApplySites.append(apply)
stats.tick(.asyncKnownCallsCount)
}
} else if inst.maySuspend {
// assert(inst is GetAsyncContinuationInst ||
// inst is GetAsyncContinuationAddrInst ||
// inst is AwaitAsyncContinuationInst, "unexpected instruction")
unknownAsyncOp.setIfUnset(inst)
}
}
let data: AnalysisResult.Data =
(function: function,
knownAsyncCalls: knownAsyncApplySites,
executors: executors,
hops: hops)
let result = AnalysisResult(data, unknownAsyncOp)
dump(result)
return Node(result)
}
/// An async demotion is a conversion of an async function like this:
///
/// ```
/// func f(_ args: Args) async -> Result { Body }
/// ```
///
/// into this:
///
/// ```
/// func f(_ args: Args) async -> Result {
/// hop_to_executor %E
/// return f_specialized(args)
/// }
///
/// func f_specialized(_ args: Args) async -> Result {
/// Body, but with all hop_to_executor instructions removed and
/// async calls replaced with calls to their demoted specializations.
/// }
/// ```
///
/// The demotion is valid only if an async function satisfies all of the
/// following requirements:
/// 1. It hops to at most one kind of executor.
/// 2. It does not perform a continuation capture.
/// 3. It only makes async function calls to callees who themselves can be
/// async demoted, with respect to either the same executor or no executor.
///
/// Then the function can be async demoted.
///
/// Requirement 3 is the tricky part of this analysis and makes it
/// interprocedural. In particular, this pass handles cycles and other kinds of
/// mutually recursive function groups.
let asyncDemotion = ModulePass(name: "async-demotion") {
(moduleContext: ModulePassContext) in
dump("\n\n\n-------------------------------\n\n\n")
defer { dump("\n\n\n-------------------------------\n\n\n") }
var graph = DemotionGraph(moduleContext)
for function in moduleContext.functions {
// Perform local analysis of each function to generate initial conditions
// of async demotion candidacy.
guard let initialState = analyzeForDemotion(function) else { continue }
// Using the local analysis, build the subgraph of the full call-graph
// that we actually want to analyze. This subgraph only deals with
// async functions that are possible candidates of demotion.
graph.addNode(initialState)
}
dump("\n\n" + graph.asGraphViz("BeforeAsyncDemotion"),
.high)
graph.optimize()
dump(stats, .high)
dump("\n\n" + graph.asGraphViz("AfterAsyncDemotion"),
.high)
}
enum LogPriority: UInt8 {
case low = 0
case med = 1
case high = 2
}
let ENABLE_LOGGING: LogPriority? = nil
/// Debug logging
private func dump<T: CustomStringConvertible>(
_ t: @autoclosure () -> T,
_ level: LogPriority = .low) {
guard let minLevel = ENABLE_LOGGING else { return }
if level.rawValue >= minLevel.rawValue {
print(t())
}
}