[Autodiff] Adds logic to generate specialized functions in the closure-spec pass

Author: jkshtj

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -353,12 +353,41 @@ struct FunctionPassContext : MutatingContext {
     return String(taking: _bridged.mangleOutlinedVariable(function.bridged))
   }
 
+  func mangle(withClosureArgs closureArgs: [Value], closureArgIndices: [Int], from applySiteCallee: Function) -> String {
+    closureArgs.withBridgedValues { bridgedClosureArgsRef in
+      closureArgIndices.withBridgedArrayRef{bridgedClosureArgIndicesRef in 
+        String(taking: _bridged.mangleWithClosureArgs(
+          bridgedClosureArgsRef, 
+          bridgedClosureArgIndicesRef, 
+          applySiteCallee.bridged
+        ))
+      }
+    }
+  }
+
   func createGlobalVariable(name: String, type: Type, isPrivate: Bool) -> GlobalVariable {
     let gv = name._withBridgedStringRef {
       _bridged.createGlobalVariable($0, type.bridged, isPrivate)
     }
     return gv.globalVar
   }
+
+  /// Utility function that should be used by optimizations that generate new functions or specialized versions of
+  /// existing functions.
+  func createAndBuildSpecializedFunction(createFn: (FunctionPassContext) -> Function,
+                                         buildFn: (Function, FunctionPassContext) -> ()) -> Function 
+  {
+    let specializedFunction = createFn(self)
+    
+    let nestedFunctionPassContext = 
+        FunctionPassContext(_bridged: _bridged.initializeNestedPassContext(specializedFunction.bridged))
+
+    defer { _bridged.deinitializedNestedPassContext() }
+
+    buildFn(specializedFunction, nestedFunctionPassContext)
+
+    return specializedFunction
+  }
 }
 
 struct SimplifyContext : MutatingContext {
@@ -445,6 +474,13 @@ extension Builder {
               context.notifyInstructionChanged, context._bridged.asNotificationHandler())
   }
 
+  /// Creates a builder which inserts instructions into an empty function, using the location of the function itself.
+  init(atStartOf function: Function, _ context: some MutatingContext) {
+    context.verifyIsTransforming(function: function)
+    self.init(insertAt: .atStartOf(function), context.notifyInstructionChanged, 
+              context._bridged.asNotificationHandler())
+  }
+
   init(staticInitializerOf global: GlobalVariable, _ context: some MutatingContext) {
     self.init(insertAt: .staticInitializer(global),
               location: Location.artificialUnreachableLocation,
@@ -629,7 +665,6 @@ extension Function {
     bridged.setIsPerformanceConstraint(isPerformanceConstraint)
   }
 
-
   func fixStackNesting(_ context: FunctionPassContext) {
     context._bridged.fixStackNesting(bridged)
   }

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -10,6 +10,7 @@ swift_compiler_sources(Optimizer
   AddressUtils.swift
   BorrowedFromUpdater.swift
   BorrowUtils.swift
+  SpecializationCloner.swift
   DiagnosticEngine.swift
   Devirtualization.swift
   EscapeUtils.swift

SwiftCompilerSources/Sources/Optimizer/Utilities/CMakeLists.txt

@@ -178,6 +178,11 @@ extension Builder {
       insertFunc(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)
+  }
 }
 
 extension Value {
@@ -398,23 +403,6 @@ extension LoadInst {
   }
 }
 
-extension PartialApplyInst {
-  var isPartialApplyOfReabstractionThunk: Bool {
-    // A partial_apply of a reabstraction thunk either has a single capture
-    // (a function) or two captures (function and dynamic Self type).
-    if self.numArguments == 1 || self.numArguments == 2, 
-       let fun = self.referencedFunction,
-       fun.isReabstractionThunk,
-       self.arguments[0].type.isFunction,
-       self.arguments[0].type.isReferenceCounted(in: self.parentFunction) || self.callee.type.isThickFunction
-    {
-      return true
-    }
-    
-    return false
-  }
-}
-
 extension FunctionPassContext {
   /// Returns true if any blocks were removed.
   func removeDeadBlocks(in function: Function) -> Bool {

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -0,0 +1,57 @@
+//===--- SpecializationCloner.swift --------------------------------------------==//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 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 OptimizerBridging
+import SIL
+
+/// Utility cloner type that can be used by optimizations that generate new functions or specialized versions of
+/// existing functions. 
+struct SpecializationCloner {
+  private var _context: FunctionPassContext
+  private var _bridged: BridgedSpecializationCloner
+
+  init(emptySpecializedFunction: Function, _ context: FunctionPassContext) {
+    self._context = context
+    self._bridged = BridgedSpecializationCloner(emptySpecializedFunction.bridged)
+  }
+
+  public var context: FunctionPassContext {
+    self._context
+  }
+
+  public var bridged: BridgedSpecializationCloner {
+    self._bridged
+  }
+  
+  public var cloned: Function {
+    bridged.getCloned().function
+  }
+
+  public var entryBlock: BasicBlock {
+    if cloned.blocks.isEmpty {
+      cloned.appendNewBlock(context)
+    } else {
+      cloned.entryBlock
+    }
+  }
+
+  public func getClonedBlock(for originalBlock: BasicBlock) -> BasicBlock {
+    bridged.getClonedBasicBlock(originalBlock.bridged).block
+  }
+
+  public func cloneFunctionBody(from originalFunction: Function, entryBlockArgs: [Value]) {
+    entryBlockArgs.withBridgedValues { bridgedEntryBlockArgs in
+      bridged.cloneFunctionBody(originalFunction.bridged, self.entryBlock.bridged, bridgedEntryBlockArgs)
+    }
+  }
+
+}

SwiftCompilerSources/Sources/Optimizer/Utilities/SpecializationCloner.swift

@@ -165,7 +165,8 @@ public func registerOptimizerTests() {
     linearLivenessTest,
     parseTestSpecificationTest,
     variableIntroducerTest,
-    gatherCallSitesTest
+    gatherCallSitesTest,
+    specializedFunctionSignatureAndBodyTest
   )
 
   // Finally register the thunk they all call through.

SwiftCompilerSources/Sources/Optimizer/Utilities/Test.swift

@@ -69,6 +69,15 @@ final public class FunctionArgument : Argument {
   public var resultDependence: LifetimeDependenceConvention? {
     parentFunction.argumentConventions[resultDependsOn: index]
   }
+
+  /// Copies the following flags from `arg`:
+  /// 1. noImplicitCopy
+  /// 2. lifetimeAnnotation
+  /// 3. closureCapture
+  /// 4. parameterPack
+  public func copyFlags(from arg: FunctionArgument) {
+    bridged.copyFlags(arg.bridged)
+  }
 }
 
 public struct Phi {
@@ -445,6 +454,17 @@ public enum ArgumentConvention : CustomStringConvertible {
     }
   }
 
+  public var isConsumed: Bool {
+    switch self {
+    case .indirectIn, .directOwned, .packOwned:
+      return true
+    case .indirectInGuaranteed, .directGuaranteed, .packGuaranteed,
+          .indirectInout, .indirectInoutAliasable, .indirectOut,
+          .packOut, .packInout, .directUnowned:
+      return false
+    }
+  }
+
   public var isExclusiveIndirect: Bool {
     switch self {
     case .indirectIn,

SwiftCompilerSources/Sources/SIL/Argument.swift

@@ -122,6 +122,8 @@ public struct InstructionList : CollectionLikeSequence, IteratorProtocol {
 
   public var first: Instruction? { currentInstruction }
 
+  public var last: Instruction? { reversed().first }
+
   public func reversed() -> ReverseInstructionList {
     if let inst = currentInstruction {
       let lastInst = inst.bridged.getLastInstOfParent().instruction

SwiftCompilerSources/Sources/SIL/BasicBlock.swift

@@ -19,25 +19,28 @@ public struct Builder {
   public enum InsertionPoint {
     case before(Instruction)
     case atEndOf(BasicBlock)
+    case atStartOf(Function)
     case staticInitializer(GlobalVariable)
   }
 
   let insertAt: InsertionPoint
-  let location: Location
+  let location: Location?
   private let notificationHandler: BridgedChangeNotificationHandler
   private let notifyNewInstruction: (Instruction) -> ()
 
   public var bridged: BridgedBuilder {
     switch insertAt {
     case .before(let inst):
       return BridgedBuilder(insertAt: .beforeInst, insertionObj: inst.bridged.obj,
-                            loc: location.bridged)
+                            loc: location!.bridged)
     case .atEndOf(let block):
       return BridgedBuilder(insertAt: .endOfBlock, insertionObj: block.bridged.obj,
-                            loc: location.bridged)
+                            loc: location!.bridged)
+    case .atStartOf(let function):
+      return BridgedBuilder(insertAt: .startOfFunction, function: function.bridged)
     case .staticInitializer(let global):
       return BridgedBuilder(insertAt: .intoGlobal, insertionObj: global.bridged.obj,
-                            loc: location.bridged)
+                            loc: location!.bridged)
     }
   }
 
@@ -53,9 +56,26 @@ public struct Builder {
     return instruction
   }
 
+  public init(insertAt: InsertionPoint, 
+              _ notifyNewInstruction: @escaping (Instruction) -> (),
+              _ notificationHandler: BridgedChangeNotificationHandler) {
+    guard case let .atStartOf(_) = insertAt else {
+      fatalError("Initializer must only be used to initialize builders that insert at the start of functions.")
+    }
+    
+    self.insertAt = insertAt
+    self.location = nil;
+    self.notifyNewInstruction = notifyNewInstruction
+    self.notificationHandler = notificationHandler
+  }
+
   public init(insertAt: InsertionPoint, location: Location,
               _ notifyNewInstruction: @escaping (Instruction) -> (),
               _ notificationHandler: BridgedChangeNotificationHandler) {
+    if case let .atStartOf(_) = insertAt {
+      fatalError("Initializer must not be used to initialize builders that insert at the start of functions.")
+    }
+
     self.insertAt = insertAt
     self.location = location;
     self.notifyNewInstruction = notifyNewInstruction
@@ -147,6 +167,18 @@ public struct Builder {
     return notifyNew(endInit.getAs(EndInitLetRefInst.self))
   }
 
+  @discardableResult
+  public func createRetainValue(operand: Value) -> RetainValueInst {
+    let retain = bridged.createRetainValue(operand.bridged)
+    return notifyNew(retain.getAs(RetainValueInst.self))
+  }
+
+  @discardableResult
+  public func createReleaseValue(operand: Value) -> ReleaseValueInst {
+    let release = bridged.createReleaseValue(operand.bridged)
+    return notifyNew(release.getAs(ReleaseValueInst.self))
+  }
+
   @discardableResult
   public func createStrongRetain(operand: Value) -> StrongRetainInst {
     let retain = bridged.createStrongRetain(operand.bridged)
@@ -290,6 +322,20 @@ public struct Builder {
     return notifyNew(tttf.getAs(ThinToThickFunctionInst.self))
   }
 
+  public func createPartialApply(
+    forFunction function: Value,
+    substitutionMap: SubstitutionMap, 
+    capturedArgs: [Value], 
+    calleeConvention: ArgumentConvention, 
+    hasUnknownResultIsolation: Bool, 
+    isOnStack: Bool
+  ) -> PartialApplyInst {
+    return capturedArgs.withBridgedValues { capturedArgsRef in
+      let pai = bridged.createPartialApply(function.bridged, capturedArgsRef, calleeConvention.bridged, substitutionMap.bridged, hasUnknownResultIsolation, isOnStack)
+      return notifyNew(pai.getAs(PartialApplyInst.self))
+    }
+  }
+
   @discardableResult
   public func createSwitchEnum(enum enumVal: Value,
                                cases: [(Int, BasicBlock)],
@@ -434,4 +480,14 @@ public struct Builder {
       let endAccess = bridged.createEndAccess(beginAccess.bridged)
       return notifyNew(endAccess.getAs(EndAccessInst.self))
   }
+
+  public func createConvertFunction(originalFunction: Value, resultType: Type, withoutActuallyEscaping: Bool) -> ConvertFunctionInst {
+    let convertFunction = bridged.createConvertFunction(originalFunction.bridged, resultType.bridged, withoutActuallyEscaping)
+    return notifyNew(convertFunction.getAs(ConvertFunctionInst.self))
+  }
+
+  public func createConvertEscapeToNoEscape(originalFunction: Value, resultType: Type, isLifetimeGuaranteed: Bool) -> ConvertEscapeToNoEscapeInst {
+    let convertFunction = bridged.createConvertEscapeToNoEscape(originalFunction.bridged, resultType.bridged, isLifetimeGuaranteed)
+    return notifyNew(convertFunction.getAs(ConvertEscapeToNoEscapeInst.self))
+  }
 }

SwiftCompilerSources/Sources/SIL/Builder.swift

@@ -93,8 +93,6 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
 
   public var isAsync: Bool { bridged.isAsync() }
 
-  public var isReabstractionThunk: Bool { bridged.isReabstractionThunk() }
-
   /// True if this is a `[global_init]` function.
   ///
   /// Such a function is typically a global addressor which calls the global's
@@ -140,7 +138,7 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
     case noThunk, thunk, reabstractionThunk, signatureOptimizedThunk
   }
 
-  var thunkKind: ThunkKind {
+  public var thunkKind: ThunkKind {
     switch bridged.isThunk() {
     case .IsNotThunk:                return .noThunk
     case .IsThunk:                   return .thunk