[Autodiff] Moves bridging code accesses in closure-spec opt behind APIs

Addresses some other surfacial feedback as well.

Author: jkshtj

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -185,27 +185,17 @@ private func getOrCreateSpecializedFunction(basedOn callSite: CallSite, _ contex
   let applySiteCallee = callSite.applyCallee
   let specializedParameters = applySiteCallee.convention.getSpecializedParameters(basedOn: callSite)
 
-  let createFn = { (functionPassContext: FunctionPassContext) in 
-    specializedFunctionName._withBridgedStringRef { nameRef in
-      let bridgedParamInfos = specializedParameters.map { $0._bridged }
-
-      return bridgedParamInfos.withUnsafeBufferPointer { paramBuf in
-        functionPassContext
-        ._bridged
-        .ClosureSpecializer_createEmptyFunctionWithSpecializedSignature(nameRef, paramBuf.baseAddress, paramBuf.count, 
-                                                                        applySiteCallee.bridged, 
-                                                                        applySiteCallee.isSerialized)
-        .function
-      }
-    }
-  }
+  let specializedFunction = 
+    context.createFunctionForClosureSpecialization(from: applySiteCallee, withName: specializedFunctionName, 
+                                                   withParams: specializedParameters, 
+                                                   withSerialization: applySiteCallee.isSerialized)
 
-  let buildFn = { (emptySpecializedFunction, functionPassContext) in 
-    let closureSpecCloner = SpecializationCloner(emptySpecializedFunction: emptySpecializedFunction, functionPassContext)
-    closureSpecCloner.cloneAndSpecializeFunctionBody(using: callSite)
-  }
+  context.buildSpecializedFunction(specializedFunction: specializedFunction,
+                                   buildFn: { (emptySpecializedFunction, functionPassContext) in 
+                                      let closureSpecCloner = SpecializationCloner(emptySpecializedFunction: emptySpecializedFunction, functionPassContext)
+                                      closureSpecCloner.cloneAndSpecializeFunctionBody(using: callSite)
+                                   })
 
-  let specializedFunction = context.createAndBuildSpecializedFunction(createFn: createFn, buildFn: buildFn)
   return (specializedFunction, false)
 }
 
@@ -479,7 +469,6 @@ private func handleApplies(for rootClosure: SingleValueInstruction, callSiteMap:
       continue
     }
 
-    // Mark the converted/reabstracted closures as used.
     if haveUsedReabstraction {
       markConvertedAndReabstractedClosuresAsUsed(rootClosure: rootClosure, convertedAndReabstractedClosure: use.value, 
                                                  convertedAndReabstractedClosures: &convertedAndReabstractedClosures)
@@ -589,7 +578,7 @@ private extension SpecializationCloner {
 
     let (allSpecializedEntryBlockArgs, closureArgIndexToAllClonedReleasableClosures) = cloneAllClosures(at: callSite)
 
-    self.cloneFunctionBody(from: callSite.applyCallee, entryBlockArgs: allSpecializedEntryBlockArgs)
+    self.cloneFunctionBody(from: callSite.applyCallee, entryBlockArguments: allSpecializedEntryBlockArgs)
 
     self.insertCleanupCodeForClonedReleasableClosures(
       from: callSite, closureArgIndexToAllClonedReleasableClosures: closureArgIndexToAllClonedReleasableClosures)
@@ -672,9 +661,7 @@ private extension SpecializationCloner {
 
     let (finalClonedReabstractedClosure, releasableClonedReabstractedClosures) = 
       builder.cloneRootClosureReabstractions(rootClosure: closureArgDesc.closure, clonedRootClosure: clonedRootClosure,
-                                             reabstractedClosure: callSite
-                                                                    .appliedArgForClosure(
-                                                                      at: closureArgDesc.closureArgIndex)!,
+                                             reabstractedClosure: callSite.appliedArgForClosure(at: closureArgDesc.closureArgIndex)!,
                                              origToClonedValueMap: origToClonedValueMap,
                                              self.context)
 
@@ -715,7 +702,7 @@ private extension SpecializationCloner {
       // Insert a `destroy_value`, for all releasable closures, in all reachable exit BBs if the closure was passed as a
       // guaranteed parameter or its type was noescape+thick. This is b/c the closure was passed at +0 originally and we
       // need to balance the initial increment of the newly created closure(s).
-      if closureArgDesc.isClosureGuaranteed || closureArgDesc.isClosureTrivialNoEscape,
+      if closureArgDesc.isClosureGuaranteed || closureArgDesc.parameterInfo.isTrivialNoescapeClosure,
          !allClonedReleasableClosures.isEmpty
       {
         for exitBlock in closureArgDesc.reachableExitBBs {
@@ -745,10 +732,10 @@ private extension SpecializationCloner {
 private extension [HashableValue: Value] {
   subscript(key: Value) -> Value? {
     get {
-      self[HashableValue(key)]
+      self[key.hashable]
     }
     set {
-      self[HashableValue(key)] = newValue
+      self[key.hashable] = newValue
     }
   }
 }
@@ -760,8 +747,8 @@ private extension Builder {
     let function = self.createFunctionRef(closureArgDesc.callee)
 
     if let pai = closureArgDesc.closure as? PartialApplyInst {
-      return self.createPartialApply(forFunction: function, substitutionMap: SubstitutionMap(), 
-                                     capturedArgs: capturedArgs, calleeConvention: pai.calleeConvention,
+      return self.createPartialApply(function: function, substitutionMap: SubstitutionMap(), 
+                                     capturedArguments: capturedArgs, calleeConvention: pai.calleeConvention,
                                      hasUnknownResultIsolation: pai.hasUnknownResultIsolation, 
                                      isOnStack: pai.isOnStack)
     } else {
@@ -806,8 +793,8 @@ private extension Builder {
           }
 
           let fri = self.createFunctionRef(function)
-          let reabstracted = self.createPartialApply(forFunction: fri, substitutionMap: SubstitutionMap(), 
-                                                     capturedArgs: [toBeReabstracted], 
+          let reabstracted = self.createPartialApply(function: fri, substitutionMap: SubstitutionMap(), 
+                                                     capturedArguments: [toBeReabstracted], 
                                                      calleeConvention: pai.calleeConvention, 
                                                      hasUnknownResultIsolation: pai.hasUnknownResultIsolation, 
                                                      isOnStack: pai.isOnStack)
@@ -895,7 +882,7 @@ private extension FunctionConvention {
 private extension ParameterInfo {
   func withSpecializedConvention(isArgTypeTrivial: Bool) -> Self {
     let specializedParamConvention =
-      if self.hasAllowedIndirectConvForClosureSpec {
+      if self.convention.isAllowedIndirectConvForClosureSpec {
         self.convention
       } else {
         isArgTypeTrivial ? ArgumentConvention.directUnowned : ArgumentConvention.directOwned
@@ -905,13 +892,8 @@ private extension ParameterInfo {
                          hasLoweredAddresses: self.hasLoweredAddresses)
   }
 
-  var hasAllowedIndirectConvForClosureSpec: Bool {
-    switch convention {
-    case .indirectInout, .indirectInoutAliasable:
-      return true
-    default:
-      return false
-    }
+  var isTrivialNoescapeClosure: Bool {
+    self.type.SILFunctionType_isTrivialNoescape()
   }
 }
 
@@ -995,21 +977,6 @@ private extension Function {
 }
 
 // ===================== Utility Types ===================== //
-private enum HashableValue: Hashable {
-  case Argument(FunctionArgument)
-  case Instruction(SingleValueInstruction)
-
-  init(_ value: Value) {
-    if let instr = value as? SingleValueInstruction {
-      self = .Instruction(instr)
-    } else if let arg = value as? FunctionArgument {
-      self = .Argument(arg)
-    } else {
-      fatalError("Invalid hashable value: \(value)")
-    }
-  }
-}
-
 private struct OrderedDict<Key: Hashable, Value> {
   private var valueIndexDict: [Key: Int] = [:]
   private var entryList: [(Key, Value)] = []
@@ -1077,10 +1044,6 @@ private struct ClosureArgDescriptor {
     closureInfo.closure
   }
 
-  var isPartialApply: Bool {
-    closure is PartialApplyInst
-  }
-
   var isPartialApplyOnStack: Bool {
     if let pai = closure as? PartialApplyInst {
       return pai.isOnStack
@@ -1116,7 +1079,7 @@ private struct ClosureArgDescriptor {
     }
   }
 
-  var arguments: (some Sequence<Value>)? {
+  var arguments: LazyMapSequence<OperandArray, Value>? {
     if let pai = closure as? PartialApplyInst {
       return pai.arguments
     }
@@ -1131,14 +1094,6 @@ private struct ClosureArgDescriptor {
     closureParamInfo.convention.isConsumed
   }
 
-  var isClosureTrivialNoEscape: Bool {
-    closureParamInfo.type.SILFunctionType_isTrivialNoescape()
-  }
-
-  var parentFunction: Function {
-    closure.parentFunction
-  }
-
   var reachableExitBBs: [BasicBlock] {
     closure.parentFunction.blocks.filter { $0.isReachableExitBlock }
   }
@@ -1161,14 +1116,6 @@ private struct CallSite {
     applySite.referencedFunction!
   }
 
-  var isCalleeSerialized: Bool {
-    applyCallee.isSerialized
-  }
-
-  var firstClosureArgDesc: ClosureArgDescriptor? {
-    closureArgDescriptors.first
-  }
-
   func hasClosureArg(at index: Int) -> Bool {
     closureArgDescriptors.contains { $0.closureArgumentIndex == index }
   }
@@ -1177,10 +1124,6 @@ private struct CallSite {
     closureArgDescriptors.first { $0.closureArgumentIndex == index }
   }
 
-  func closureArg(at index: Int) -> SingleValueInstruction? {
-    closureArgDesc(at: index)?.closure
-  }
-
   func appliedArgForClosure(at index: Int) -> Value? {
     if let closureArgDesc = closureArgDesc(at: index) {
       return applySite.arguments[closureArgDesc.closureArgIndex - applySite.unappliedArgumentCount]
@@ -1189,14 +1132,6 @@ private struct CallSite {
     return nil
   }
 
-  func closureCallee(at index: Int) -> Function? {
-    closureArgDesc(at: index)?.callee
-  }
-
-  func closureLoc(at index: Int) -> Location? {
-    closureArgDesc(at: index)?.location
-  }
-
   func specializedCalleeName(_ context: FunctionPassContext) -> String {
     let closureArgs = Array(self.closureArgDescriptors.map { $0.closure })
     let closureIndices = Array(self.closureArgDescriptors.map { $0.closureArgIndex })

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -372,21 +372,29 @@ struct FunctionPassContext : MutatingContext {
     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 
+  func createFunctionForClosureSpecialization(from applySiteCallee: Function, withName specializedFunctionName: String, 
+                                              withParams specializedParameters: [ParameterInfo], 
+                                              withSerialization isSerialized: Bool) -> Function 
   {
-    let specializedFunction = createFn(self)
-    
+    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)
+      defer { _bridged.deinitializedNestedPassContext() }
 
-    return specializedFunction
+      buildFn(specializedFunction, nestedFunctionPassContext)
   }
 }
 
@@ -477,8 +485,8 @@ extension Builder {
   /// 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())
+    self.init(insertAt: .atStartOf(function), location: function.location,
+              context.notifyInstructionChanged, context._bridged.asNotificationHandler())
   }
 
   init(staticInitializerOf global: GlobalVariable, _ context: some MutatingContext) {

SwiftCompilerSources/Sources/Optimizer/Utilities/SpecializationCloner.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// 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
@@ -16,40 +16,32 @@ 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
+  private let bridged: BridgedSpecializationCloner
+  let context: FunctionPassContext
 
   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
+    self.bridged = BridgedSpecializationCloner(emptySpecializedFunction.bridged)
+    self.context = context
   }
 
-  public var cloned: Function {
+  var cloned: Function {
     bridged.getCloned().function
   }
 
-  public var entryBlock: BasicBlock {
+  var entryBlock: BasicBlock {
     if cloned.blocks.isEmpty {
       cloned.appendNewBlock(context)
     } else {
       cloned.entryBlock
     }
   }
 
-  public func getClonedBlock(for originalBlock: BasicBlock) -> BasicBlock {
+  func getClonedBlock(for originalBlock: BasicBlock) -> BasicBlock {
     bridged.getClonedBasicBlock(originalBlock.bridged).block
   }
 
-  public func cloneFunctionBody(from originalFunction: Function, entryBlockArgs: [Value]) {
-    entryBlockArgs.withBridgedValues { bridgedEntryBlockArgs in
+  func cloneFunctionBody(from originalFunction: Function, entryBlockArguments: [Value]) {
+    entryBlockArguments.withBridgedValues { bridgedEntryBlockArgs in
       bridged.cloneFunctionBody(originalFunction.bridged, self.entryBlock.bridged, bridgedEntryBlockArgs)
     }
   }