Merge pull request #84704 from eeckstein/closure-specialization

ClosureSpecialization: support for OSSA and a big overhaul

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -112,7 +112,7 @@ private func constantPropagateCaptures(of partialApply: PartialApplyInst, _ cont
     // Escaping closures consume their arguments. Therefore we need to destroy the removed argument values.
     addCompensatingDestroys(for: constArgs, context)
   }
-  let newArguments = nonConstArgs.map { $0.value }
+  let newArguments = Array(nonConstArgs.values)
   rewritePartialApply(partialApply, withSpecialized: specializedCallee, arguments: newArguments, context)
 }
 
@@ -195,8 +195,7 @@ private func cloneArgument(_ argumentOp: Operand,
   if partialApply.calleeArgumentConventions[calleeArgIdx].isGuaranteed {
     // If the original argument was passed as guaranteed, i.e. is _not_ destroyed in the closure, we have
     // to destroy the cloned argument at function exits.
-    for block in targetFunction.blocks where block.terminator.isFunctionExiting {
-      let builder = Builder(before: block.terminator, context)
+    Builder.insertCleanupAtFunctionExits(of: targetFunction, context) { builder in
       builder.emitDestroy(of: clonedArg)
     }
   }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ConstantCapturePropagation.swift

@@ -404,7 +404,7 @@ private extension Value {
       case is LoadInst:
         return true
       case is StructInst, is TupleInst:
-        worklist.pushIfNotVisited(contentsOf: (v as! Instruction).operands.lazy.map { $0.value })
+        worklist.pushIfNotVisited(contentsOf: (v as! Instruction).operands.values)
       case let ei as EnumInst:
         if let payload = ei.payload {
           worklist.pushIfNotVisited(payload)

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/InitializeStaticGlobals.swift

@@ -27,9 +27,9 @@ private extension CondBranchInst {
     }
     let builder = Builder(before: self, context)
     if conditionValue == 0 {
-      builder.createBranch(to: falseBlock, arguments: Array(falseOperands.map { $0.value }))
+      builder.createBranch(to: falseBlock, arguments: Array(falseOperands.values))
     } else {
-      builder.createBranch(to: trueBlock, arguments: Array(trueOperands.map { $0.value }))
+      builder.createBranch(to: trueBlock, arguments: Array(trueOperands.values))
     }
     context.erase(instruction: self)
   }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyCondBranch.swift

@@ -339,13 +339,9 @@ private struct StackProtectionOptimization {
 
     builder.createCopyAddr(from: argument, to: temporary, takeSource: true, initializeDest: true)
 
-    for block in function.blocks {
-      let terminator = block.terminator
-      if terminator.isFunctionExiting {
-        let exitBuilder = Builder(before: terminator, location: terminator.location.asAutoGenerated, context)
-        exitBuilder.createCopyAddr(from: temporary, to: argument, takeSource: true, initializeDest: true)
-        exitBuilder.createDeallocStack(temporary)
-      }
+    Builder.insertCleanupAtFunctionExits(of: function, context) { builder in
+      builder.createCopyAddr(from: temporary, to: argument, takeSource: true, initializeDest: true)
+      builder.createDeallocStack(temporary)
     }
     log("move addr protection in \(function.name): \(argument)")
 

SwiftCompilerSources/Sources/Optimizer/ModulePasses/StackProtection.swift

@@ -111,10 +111,34 @@ struct FunctionPassContext : MutatingContext {
     return String(taking: bridgedPassContext.mangleOutlinedVariable(function.bridged))
   }
 
-  func mangle(withClosureArguments closureArgs: [(argumentIndex: Int, argumentValue: Value)],
+  enum ClosureArgumentMangling {
+    case closure(SingleValueInstruction)
+
+    /// The argument specializes for the same closure as a previous argument, e.g.
+    /// ```
+    ///   %1 = partial_apply %closure
+    ///   apply %f(%1, %1)    // first argument: `.closure(%1)`
+    ///                       // second argument: `.previousArgumentIndex(0)`
+    case previousArgumentIndex(Int)
+  }
+
+  func mangle(withClosureArguments closureArgs: [(argumentIndex: Int, argumentValue: ClosureArgumentMangling)],
               from applySiteCallee: Function
   ) -> String {
-    closureArgs.withBridgedArrayRef{ bridgedClosureArgs in
+    let bridgedArgManglings = closureArgs.map {
+      switch $0.argumentValue {
+      case .closure(let closure):
+        return BridgedPassContext.ClosureArgMangling(argIdx: $0.argumentIndex,
+                                                     inst: Optional<Instruction>(closure).bridged,
+                                                     otherArgIdx: -1)
+      case .previousArgumentIndex(let idx):
+        return BridgedPassContext.ClosureArgMangling(argIdx: $0.argumentIndex,
+                                                     inst: OptionalBridgedInstruction(),
+                                                     otherArgIdx: idx)
+      }
+    }
+
+    return bridgedArgManglings.withBridgedArrayRef{ bridgedClosureArgs in
       String(taking: bridgedPassContext.mangleWithClosureArgs(bridgedClosureArgs, applySiteCallee.bridged))
     }
   }

SwiftCompilerSources/Sources/Optimizer/PassManager/FunctionPassContext.swift

@@ -105,7 +105,7 @@ private func registerSwiftPasses() {
   registerPass(tempRValueElimination, { tempRValueElimination.run($0) })
   registerPass(mandatoryTempRValueElimination, { mandatoryTempRValueElimination.run($0) })
   registerPass(tempLValueElimination, { tempLValueElimination.run($0) })
-  registerPass(generalClosureSpecialization, { generalClosureSpecialization.run($0) })
+  registerPass(closureSpecialization, { closureSpecialization.run($0) })
   registerPass(autodiffClosureSpecialization, { autodiffClosureSpecialization.run($0) })
   registerPass(loopInvariantCodeMotionPass, { loopInvariantCodeMotionPass.run($0) })
 

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -43,7 +43,6 @@ public func registerOptimizerTests() {
   registerFunctionTests(
     addressOwnershipLiveRangeTest,
     argumentConventionsTest,
-    getPullbackClosureInfoTest,
     interiorLivenessTest,
     lifetimeDependenceRootTest,
     lifetimeDependenceScopeTest,
@@ -52,8 +51,6 @@ public func registerOptimizerTests() {
     localVariableReachableUsesTest,
     localVariableReachingAssignmentsTest,
     rangeOverlapsPathTest,
-    rewrittenCallerBodyTest,
-    specializedFunctionSignatureAndBodyTest,
     variableIntroducerTest
   )
 

SwiftCompilerSources/Sources/Optimizer/Utilities/FunctionTest.swift

@@ -71,7 +71,7 @@ extension Value {
       }
       switch v {
       case let fw as ForwardingInstruction:
-        worklist.pushIfNotVisited(contentsOf: fw.definedOperands.lazy.map { $0.value })
+        worklist.pushIfNotVisited(contentsOf: fw.definedOperands.values)
       case let bf as BorrowedFromInst:
         worklist.pushIfNotVisited(bf.borrowedValue)
       case let bb as BeginBorrowInst:
@@ -82,7 +82,7 @@ extension Value {
         } else if let termResult = TerminatorResult(arg),
                let fw = termResult.terminator as? ForwardingInstruction
         {
-          worklist.pushIfNotVisited(contentsOf: fw.definedOperands.lazy.map { $0.value })
+          worklist.pushIfNotVisited(contentsOf: fw.definedOperands.values)
         }
       default:
         continue
@@ -262,6 +262,18 @@ extension Builder {
     }
   }
 
+  static func insertCleanupAtFunctionExits(
+    of function: Function,
+    _ context: some MutatingContext,
+    insert: (Builder) -> ()
+  ) {
+    for exitBlock in function.blocks where exitBlock.terminator.isFunctionExiting {
+      let terminator = exitBlock.terminator
+      let builder = Builder(before: terminator, location: terminator.location.asCleanup, context)
+      insert(builder)
+    }
+  }
+
   func destroyCapturedArgs(for paiOnStack: PartialApplyInst) {
     precondition(paiOnStack.isOnStack, "Function must only be called for `partial_apply`s on stack!")
     self.bridged.destroyCapturedArgs(paiOnStack.bridged)

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -170,7 +170,7 @@ public struct Phi {
   }
 
   public var incomingValues: LazyMapSequence<LazyMapSequence<PredecessorList, Operand>, Value> {
-    incomingOperands.lazy.map { $0.value }
+    incomingOperands.values
   }
 
   public var isReborrow: Bool { value.isReborrow }