Merge pull request #73596 from jkshtj/specialize

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

Author: jkshtj

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -353,12 +353,49 @@ 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
   }
+
+  func createFunctionForClosureSpecialization(from applySiteCallee: Function, withName specializedFunctionName: String, 
+                                              withParams specializedParameters: [ParameterInfo], 
+                                              withSerialization isSerialized: Bool) -> Function 
+  {
+    return specializedFunctionName._withBridgedStringRef { nameRef in
+      let bridgedParamInfos = specializedParameters.map { $0._bridged }
+
+      return bridgedParamInfos.withUnsafeBufferPointer { paramBuf in
+        _bridged.ClosureSpecializer_createEmptyFunctionWithSpecializedSignature(nameRef, paramBuf.baseAddress, 
+                                                                                paramBuf.count, 
+                                                                                applySiteCallee.bridged, 
+                                                                                isSerialized).function
+      }
+    }
+  }
+
+  func buildSpecializedFunction(specializedFunction: Function, buildFn: (Function, FunctionPassContext) -> ()) {
+    let nestedFunctionPassContext = 
+        FunctionPassContext(_bridged: _bridged.initializeNestedPassContext(specializedFunction.bridged))
+
+      defer { _bridged.deinitializedNestedPassContext() }
+
+      buildFn(specializedFunction, nestedFunctionPassContext)
+  }
 }
 
 struct SimplifyContext : MutatingContext {
@@ -445,6 +482,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), location: function.location,
+              context.notifyInstructionChanged, context._bridged.asNotificationHandler())
+  }
+
   init(staticInitializerOf global: GlobalVariable, _ context: some MutatingContext) {
     self.init(insertAt: .staticInitializer(global),
               location: Location.artificialUnreachableLocation,
@@ -629,7 +673,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,49 @@
+//===--- SpecializationCloner.swift --------------------------------------------==//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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 let bridged: BridgedSpecializationCloner
+  let context: FunctionPassContext
+
+  init(emptySpecializedFunction: Function, _ context: FunctionPassContext) {
+    self.bridged = BridgedSpecializationCloner(emptySpecializedFunction.bridged)
+    self.context = context
+  }
+  
+  var cloned: Function {
+    bridged.getCloned().function
+  }
+
+  var entryBlock: BasicBlock {
+    if cloned.blocks.isEmpty {
+      return cloned.appendNewBlock(context)
+    } else {
+      return cloned.entryBlock
+    }
+  }
+
+  func getClonedBlock(for originalBlock: BasicBlock) -> BasicBlock {
+    bridged.getClonedBasicBlock(originalBlock.bridged).block
+  }
+
+  func cloneFunctionBody(from originalFunction: Function, entryBlockArguments: [Value]) {
+    entryBlockArguments.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,6 +19,7 @@ public struct Builder {
   public enum InsertionPoint {
     case before(Instruction)
     case atEndOf(BasicBlock)
+    case atStartOf(Function)
     case staticInitializer(GlobalVariable)
   }
 
@@ -35,6 +36,9 @@ public struct Builder {
     case .atEndOf(let block):
       return BridgedBuilder(insertAt: .endOfBlock, insertionObj: block.bridged.obj,
                             loc: location.bridged)
+    case .atStartOf(let function):
+      return BridgedBuilder(insertAt: .startOfFunction, insertionObj: function.bridged.obj,
+                            loc: location.bridged)
     case .staticInitializer(let global):
       return BridgedBuilder(insertAt: .intoGlobal, insertionObj: global.bridged.obj,
                             loc: location.bridged)
@@ -290,6 +294,20 @@ public struct Builder {
     return notifyNew(tttf.getAs(ThinToThickFunctionInst.self))
   }
 
+  public func createPartialApply(
+    function: Value,
+    substitutionMap: SubstitutionMap, 
+    capturedArguments: [Value], 
+    calleeConvention: ArgumentConvention, 
+    hasUnknownResultIsolation: Bool, 
+    isOnStack: Bool
+  ) -> PartialApplyInst {
+    return capturedArguments.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 +452,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

@@ -21,6 +21,10 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
     return StringRef(bridged: bridged.getName())
   }
 
+  public var location: Location {
+    return Location(bridged: bridged.getLocation())
+  }
+
   final public var description: String {
     return String(taking: bridged.getDebugDescription())
   }
@@ -93,8 +97,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 +142,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