Makes the swift-based closure-spec pass an experimental frontend feature

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
@@ -30,5 +31,4 @@ swift_compiler_sources(Optimizer
   SimplificationPasses.swift
   StackPromotion.swift
   StripObjectHeaders.swift
-  AutodiffClosureSpecialization.swift
 )

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AutodiffClosureSpecialization.swift renamed to SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -1,4 +1,4 @@
-//===--- AutodiffClosureSpecialization.swift ---------------------------===//
+//===--- ClosureSpecialization.swift ---------------------------===//
 //
 // This source file is part of the Swift.org open source project
 //
@@ -10,8 +10,14 @@
 //
 //===-----------------------------------------------------------------------===//
 
+/// This file contains the closure-specialization optimizations for general and differentiable Swift.
+
+/// General Closure Specialization
+/// ------------------------------------
+/// TODO: Add description when the functionality is added.
+
 /// AutoDiff Closure Specialization
-/// ----------------------
+/// -------------------------------
 /// This optimization performs closure specialization tailored for the patterns seen in Swift Autodiff. In principle,
 /// the optimization does the same thing as the existing closure specialization pass. However, it is tailored to the
 /// patterns of Swift Autodiff.
@@ -100,12 +106,32 @@ private func log(_ message: @autoclosure () -> String) {
 }
 
 // =========== Entry point =========== //
-let autodiffClosureSpecialization = FunctionPass(name: "autodiff-closure-specialize") {
+let generalClosureSpecialization = FunctionPass(name: "experimental-swift-based-closure-specialization") {
+  (function: Function, context: FunctionPassContext) in
+  // TODO: Implement general closure specialization optimization
+  print("NOT IMPLEMENTED")
+}
+
+let autodiffClosureSpecialization = FunctionPass(name: "autodiff-closure-specialization") {
   (function: Function, context: FunctionPassContext) in
   // TODO: Pass is a WIP and current implementation is incomplete
   if !function.isAutodiffVJP {
     return
   }
+
+  print("Specializing closures in function: \(function.name)")
+  print("===============================================")
+  var callSites = gatherCallSites(in: function, context)
+
+  callSites.forEach { callSite in
+    print("PartialApply call site: \(callSite.applySite)")
+    print("Passed in closures: ")
+    for index in callSite.closureArgDescriptors.indices {
+      var closureArgDescriptor = callSite.closureArgDescriptors[index]
+      print("\(index+1). \(closureArgDescriptor.closureInfo.closure)")
+    }
+  }
+  print("\n")
 }
 
 // =========== Top-level functions ========== //
@@ -276,8 +302,8 @@ private func handleNonApplies(for rootClosure: SingleValueInstruction,
       if !pai.isPullbackInResultOfAutodiffVJP,
           pai.isPartialApplyOfReabstractionThunk,
           pai.isSupportedClosure,
-          pai.callee.type.isNoEscapeFunction,
-          pai.callee.type.isThickFunction 
+          pai.arguments[0].type.isNoEscapeFunction,
+          pai.arguments[0].type.isThickFunction
       {
         rootClosureConversionsAndReabstractions.pushIfNotVisited(contentsOf: pai.uses)
         possibleMarkDependenceBases.insert(pai)
@@ -338,6 +364,7 @@ private func handleApplies(for rootClosure: SingleValueInstruction, callSiteMap:
       continue
     }
 
+    // TODO: Handling generic closures may be possible but is not yet implemented
     if pai.hasSubstitutions || !pai.calleeIsDynamicFunctionRef || !pai.isPullbackInResultOfAutodiffVJP {
       continue
     }
@@ -507,7 +534,7 @@ private func markConvertedAndReabstractedClosuresAsUsed(rootClosure: Value, conv
       convertedAndReabstractedClosures.insert(pai)
       return 
         markConvertedAndReabstractedClosuresAsUsed(rootClosure: rootClosure, 
-                                                   convertedAndReabstractedClosure: pai.callee, 
+                                                   convertedAndReabstractedClosure: pai.arguments[0], 
                                                    convertedAndReabstractedClosures: &convertedAndReabstractedClosures)
     case let cvt as ConvertFunctionInst:
       convertedAndReabstractedClosures.insert(cvt)

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -93,6 +93,7 @@ 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

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -405,7 +405,8 @@ extension PartialApplyInst {
     if self.numArguments == 1 || self.numArguments == 2, 
        let fun = self.referencedFunction,
        fun.isReabstractionThunk,
-       self.callee.type.isReferenceCounted(in: self.parentFunction) || self.callee.type.isThickFunction
+       self.arguments[0].type.isFunction,
+       self.arguments[0].type.isReferenceCounted(in: self.parentFunction) || self.callee.type.isThickFunction
     {
       return true
     }

include/swift/AST/SILOptions.h

@@ -290,6 +290,10 @@ class SILOptions {
   /// Are we building in embedded Swift + -no-allocations?
   bool NoAllocations = false;
 
+  /// Should we use the experimental Swift based closure-specialization
+  /// optimization pass instead of the existing C++ one.
+  bool EnableExperimentalSwiftBasedClosureSpecialization = false;
+
   /// The name of the file to which the backend should save optimization
   /// records.
   std::string OptRecordFile;

include/swift/Option/FrontendOptions.td

@@ -359,6 +359,10 @@ def enable_experimental_async_top_level :
 // HIDDEN FLAGS
 let Flags = [FrontendOption, NoDriverOption, HelpHidden] in {
 
+def enable_experimental_swift_based_closure_specialization : 
+  Flag<["-"], "experimental-swift-based-closure-specialization">,
+  HelpText<"Use the experimental Swift based closure-specialization optimization pass instead of the existing C++ one">; 
+
 def checked_async_objc_bridging : Joined<["-"], "checked-async-objc-bridging=">,
   HelpText<"Control whether checked continuations are used when bridging "
            "async calls from Swift to ObjC: 'on', 'off' ">;

include/swift/SILOptimizer/PassManager/Passes.def

@@ -146,8 +146,13 @@ PASS(CapturePropagation, "capture-prop",
      "Captured Constant Propagation")
 PASS(ClosureSpecializer, "closure-specialize",
      "Closure Specialization on Constant Function Arguments")
-SWIFT_FUNCTION_PASS(AutodiffClosureSpecialization, "autodiff-closure-specialize",
+
+// NOTE - ExperimentalSwiftBasedClosureSpecialization and AutodiffClosureSpecialization are a WIP
+SWIFT_FUNCTION_PASS(ExperimentalSwiftBasedClosureSpecialization, "experimental-swift-based-closure-specialization",
+     "General closure-specialization pass written in Swift")
+SWIFT_FUNCTION_PASS(AutodiffClosureSpecialization, "autodiff-closure-specialization",
      "Autodiff specific closure-specialization pass")
+
 PASS(ClosureLifetimeFixup, "closure-lifetime-fixup",
      "Closure Lifetime Fixup")
 PASS(CodeSinking, "code-sinking",

lib/Frontend/CompilerInvocation.cpp

@@ -2609,6 +2609,9 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
 
   Opts.NoAllocations = Args.hasArg(OPT_no_allocations);
 
+  Opts.EnableExperimentalSwiftBasedClosureSpecialization = 
+      Args.hasArg(OPT_enable_experimental_swift_based_closure_specialization);
+
   return false;
 }
 

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -768,7 +768,11 @@ static void addMidLevelFunctionPipeline(SILPassPipelinePlan &P) {
   P.addCapturePropagation();
 
   // Specialize closure.
-  P.addClosureSpecializer();
+  if (P.getOptions().EnableExperimentalSwiftBasedClosureSpecialization) {
+    P.addExperimentalSwiftBasedClosureSpecialization();
+  } else {
+    P.addClosureSpecializer();
+  }
 
   // Do the second stack promotion on low-level SIL.
   P.addStackPromotion();