Merge pull request #71775 from jkshtj/main

[Autodiff] Adds part of the Autodiff specific closure-specialization optimization pass

Author: jkshtj

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

@@ -12,6 +12,7 @@ swift_compiler_sources(Optimizer
   AsyncDemotion.swift
   BooleanLiteralFolding.swift
   CleanupDebugSteps.swift
+  ClosureSpecialization.swift
   ComputeEscapeEffects.swift
   ComputeSideEffects.swift
   DeadStoreElimination.swift

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -93,6 +93,9 @@ private func registerSwiftPasses() {
   registerPass(lifetimeDependenceDiagnosticsPass, { lifetimeDependenceDiagnosticsPass.run($0) })
   registerPass(lifetimeDependenceInsertionPass, { lifetimeDependenceInsertionPass.run($0) })
   registerPass(lifetimeDependenceScopeFixupPass, { lifetimeDependenceScopeFixupPass.run($0) })
+  registerPass(generalClosureSpecialization, { generalClosureSpecialization.run($0) })
+  registerPass(autodiffClosureSpecialization, { autodiffClosureSpecialization.run($0) })
+
   // Instruction passes
   registerForSILCombine(BeginCOWMutationInst.self, { run(BeginCOWMutationInst.self, $0) })
   registerForSILCombine(GlobalValueInst.self,      { run(GlobalValueInst.self, $0) })

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -398,6 +398,23 @@ 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 {
@@ -540,6 +557,10 @@ extension Function {
     }
     return nil
   }
+
+  var mayBindDynamicSelf: Bool {
+    self.bridged.mayBindDynamicSelf()
+  }
 }
 
 extension FullApplySite {

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

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

SwiftCompilerSources/Sources/Optimizer/Utilities/Test.swift

@@ -108,6 +108,10 @@ public protocol ApplySite : Instruction {
 extension ApplySite {
   public var callee: Value { operands[ApplyOperandConventions.calleeIndex].value }
 
+  public var hasSubstitutions: Bool {
+    return substitutionMap.hasAnySubstitutableParams
+  }
+
   public var isAsync: Bool {
     return callee.type.isAsyncFunction
   }
@@ -126,7 +130,15 @@ extension ApplySite {
     return false
   }
 
-  /// Returns the subset of operands that are argument operands.
+  public var isCalleeNoReturn: Bool {
+    bridged.ApplySite_isCalleeNoReturn()  
+  }
+
+  public var isCalleeTrapNoReturn: Bool {
+    referencedFunction?.isTrapNoReturn ?? false
+  }
+
+  /// Returns the subset of operands which are argument operands.
   ///
   /// This does not include the callee function operand.
   public var argumentOperands: OperandArray {

SwiftCompilerSources/Sources/SIL/ApplySite.swift

@@ -57,6 +57,26 @@ final public class BasicBlock : CustomStringConvertible, HasShortDescription, Ha
     successors.count == 1 ? successors[0] : nil
   }
 
+  /// All function exiting blocks except for ones with an `unreachable` terminator,
+  /// not immediately preceded by an apply of a no-return function.
+  public var isReachableExitBlock: Bool {
+    switch terminator {
+      case let termInst where termInst.isFunctionExiting:
+        return true
+      case is UnreachableInst:
+        if let instBeforeUnreachable = terminator.previous,
+           let ai = instBeforeUnreachable as? ApplyInst,
+           ai.isCalleeNoReturn && !ai.isCalleeTrapNoReturn
+        {
+          return true
+        }
+
+        return false
+      default:
+        return false
+    }
+  }
+
   /// The index of the basic block in its function.
   /// This has O(n) complexity. Only use it for debugging
   public var index: Int {

SwiftCompilerSources/Sources/SIL/BasicBlock.swift

@@ -31,6 +31,12 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
     hasher.combine(ObjectIdentifier(self))
   }
 
+  public var isTrapNoReturn: Bool { bridged.isTrapNoReturn() }
+
+  public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() }
+
+  public var specializationLevel: Int { bridged.specializationLevel() }
+  
   public var hasOwnership: Bool { bridged.hasOwnership() }
 
   public var hasLoweredAddresses: Bool { bridged.hasLoweredAddresses() }
@@ -59,6 +65,10 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
     entryBlock.arguments.lazy.map { $0 as! FunctionArgument }
   }
 
+  public func argument(at index: Int) -> FunctionArgument {
+    entryBlock.arguments[index] as! FunctionArgument
+  }
+
   /// All instructions of all blocks.
   public var instructions: LazySequence<FlattenSequence<LazyMapSequence<BasicBlockList, InstructionList>>> {
     blocks.lazy.flatMap { $0.instructions }
@@ -83,6 +93,8 @@ 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

SwiftCompilerSources/Sources/SIL/Function.swift

@@ -893,7 +893,16 @@ class ConvertFunctionInst : SingleValueInstruction, UnaryInstruction {
 }
 
 final public
-class ThinToThickFunctionInst : SingleValueInstruction, UnaryInstruction {}
+class ThinToThickFunctionInst : SingleValueInstruction, UnaryInstruction {
+  public var callee: Value { operand.value }
+
+  public var referencedFunction: Function? {
+    if let fri = callee as? FunctionRefInst {
+      return fri.referencedFunction
+    }
+    return nil
+  }
+}
 
 final public class ThickToObjCMetatypeInst : SingleValueInstruction {}
 final public class ObjCToThickMetatypeInst : SingleValueInstruction {}

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -25,6 +25,8 @@ public struct SubstitutionMap {
 
   public var isEmpty: Bool { bridged.isEmpty() }
 
+  public var hasAnySubstitutableParams: Bool { bridged.hasAnySubstitutableParams() }
+
   public var replacementTypes: OptionalTypeArray {
     let types = BridgedTypeArray.fromReplacementTypes(bridged)
     return OptionalTypeArray(bridged: types)