Merge pull request #73688 from jkshtj/main

jkshtj · web-flow · commit 64da348ee911 · 2024-05-23T14:21:31.000-07:00
[Autodiff] Adds logic to rewrite call-sites using functions specialized by the closure-spec optimization
diff --git a/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift b/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift
diff --git a/SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift b/SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift
@@ -60,6 +60,10 @@ extension Context {
       _bridged.lookupFunction($0).function
     }
   }
+
+  func notifyNewFunction(function: Function, derivedFrom: Function) {
+    _bridged.addFunctionToPassManagerWorklist(function.bridged, derivedFrom.bridged)
+  }
 }
 
 /// A context which allows mutation of a function's SIL.
@@ -357,7 +361,7 @@ struct FunctionPassContext : MutatingContext {
     return String(taking: _bridged.mangleOutlinedVariable(function.bridged))
   }
 
-  func mangle(withClosureArgs closureArgs: [Value], closureArgIndices: [Int], from applySiteCallee: Function) -> String {
+  func mangle(withClosureArguments closureArgs: [Value], closureArgIndices: [Int], from applySiteCallee: Function) -> String {
     closureArgs.withBridgedValues { bridgedClosureArgsRef in
       closureArgIndices.withBridgedArrayRef{bridgedClosureArgIndicesRef in 
         String(taking: _bridged.mangleWithClosureArgs(
@@ -392,13 +396,13 @@ struct FunctionPassContext : MutatingContext {
     }
   }
 
-  func buildSpecializedFunction(specializedFunction: Function, buildFn: (Function, FunctionPassContext) -> ()) {
+  func buildSpecializedFunction<T>(specializedFunction: Function, buildFn: (Function, FunctionPassContext) -> T) -> T {
     let nestedFunctionPassContext = 
         FunctionPassContext(_bridged: _bridged.initializeNestedPassContext(specializedFunction.bridged))
 
       defer { _bridged.deinitializedNestedPassContext() }
 
-      buildFn(specializedFunction, nestedFunctionPassContext)
+      return buildFn(specializedFunction, nestedFunctionPassContext)
   }
 }
 
diff --git a/SwiftCompilerSources/Sources/Optimizer/Utilities/Test.swift b/SwiftCompilerSources/Sources/Optimizer/Utilities/Test.swift
@@ -167,7 +167,8 @@ public func registerOptimizerTests() {
     parseTestSpecificationTest,
     variableIntroducerTest,
     gatherCallSitesTest,
-    specializedFunctionSignatureAndBodyTest
+    specializedFunctionSignatureAndBodyTest,
+    rewrittenCallerBodyTest
   )
 
   // Finally register the thunk they all call through.
diff --git a/SwiftCompilerSources/Sources/SIL/Builder.swift b/SwiftCompilerSources/Sources/SIL/Builder.swift
@@ -151,6 +151,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)
diff --git a/SwiftCompilerSources/Sources/SIL/Operand.swift b/SwiftCompilerSources/Sources/SIL/Operand.swift
@@ -62,6 +62,10 @@ public struct OperandArray : RandomAccessCollection, CustomReflectable {
     self.count = count
   }
 
+  static public var empty: OperandArray {
+    OperandArray(base: OptionalBridgedOperand(bridged: nil), count: 0)
+  }
+
   public var startIndex: Int { return 0 }
   public var endIndex: Int { return count }
   
diff --git a/include/swift/SIL/SILBridging.h b/include/swift/SIL/SILBridging.h
@@ -1210,6 +1210,10 @@ struct BridgedBuilder{
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createBeginDeallocRef(BridgedValue reference,
                                                                               BridgedValue allocation) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createEndInitLetRef(BridgedValue op) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction
+  createRetainValue(BridgedValue op) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction
+  createReleaseValue(BridgedValue op) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createStrongRetain(BridgedValue op) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createStrongRelease(BridgedValue op) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUnownedRetain(BridgedValue op) const;
diff --git a/include/swift/SIL/SILBridgingImpl.h b/include/swift/SIL/SILBridgingImpl.h
@@ -19,6 +19,7 @@
 #ifndef SWIFT_SIL_SILBRIDGING_IMPL_H
 #define SWIFT_SIL_SILBRIDGING_IMPL_H
 
+#include "SILBridging.h"
 #include "swift/AST/Builtins.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/SubstitutionMap.h"
@@ -1589,6 +1590,18 @@ BridgedInstruction BridgedBuilder::createEndInitLetRef(BridgedValue op) const {
   return {unbridged().createEndInitLetRef(regularLoc(), op.getSILValue())};
 }
 
+BridgedInstruction BridgedBuilder::createRetainValue(BridgedValue op) const {
+  auto b = unbridged();
+  return {b.createRetainValue(regularLoc(), op.getSILValue(),
+                              b.getDefaultAtomicity())};
+}
+
+BridgedInstruction BridgedBuilder::createReleaseValue(BridgedValue op) const {
+  auto b = unbridged();
+  return {b.createReleaseValue(regularLoc(), op.getSILValue(),
+                               b.getDefaultAtomicity())};
+}
+
 BridgedInstruction BridgedBuilder::createStrongRetain(BridgedValue op) const {
   auto b = unbridged();
   return {b.createStrongRetain(regularLoc(), op.getSILValue(), b.getDefaultAtomicity())};
diff --git a/include/swift/SILOptimizer/OptimizerBridging.h b/include/swift/SILOptimizer/OptimizerBridging.h
@@ -337,6 +337,9 @@ struct BridgedPassContext {
   BRIDGED_INLINE bool continueWithNextSubpassRun(OptionalBridgedInstruction inst) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedPassContext initializeNestedPassContext(BridgedFunction newFunction) const;
   BRIDGED_INLINE void deinitializedNestedPassContext() const;
+  BRIDGED_INLINE void
+  addFunctionToPassManagerWorklist(BridgedFunction newFunction,
+                                   BridgedFunction oldFunction) const;
 
   // SSAUpdater
 
diff --git a/include/swift/SILOptimizer/OptimizerBridgingImpl.h b/include/swift/SILOptimizer/OptimizerBridgingImpl.h
@@ -454,6 +454,13 @@ void BridgedPassContext::SSAUpdater_initialize(
                                    BridgedValue::castToOwnership(ownership));
 }
 
+void BridgedPassContext::addFunctionToPassManagerWorklist(
+    BridgedFunction newFunction, BridgedFunction oldFunction) const {
+  swift::SILPassManager *pm = invocation->getPassManager();
+  pm->addFunctionToWorklist(newFunction.getFunction(),
+                            oldFunction.getFunction());
+}
+
 void BridgedPassContext::SSAUpdater_addAvailableValue(BridgedBasicBlock block, BridgedValue value) const {
   invocation->getSSAUpdater()->addAvailableValue(block.unbridged(),
                                                  value.getSILValue());
diff --git a/lib/SILOptimizer/PassManager/PassPipeline.cpp b/lib/SILOptimizer/PassManager/PassPipeline.cpp
@@ -1008,10 +1008,14 @@ SILPassPipelinePlan::getPerformancePassPipeline(const SILOptions &Options) {
   if (Options.StopOptimizationAfterSerialization)
     return P;
 
+  P.addAutodiffClosureSpecialization();
+
   // After serialization run the function pass pipeline to iteratively lower
   // high-level constructs like @_semantics calls.
   addMidLevelFunctionPipeline(P);
 
+  P.addAutodiffClosureSpecialization();
+
   // Perform optimizations that specialize.
   addClosureSpecializePassPipeline(P);
 
diff --git a/lib/SILOptimizer/Transforms/PerformanceInliner.cpp b/lib/SILOptimizer/Transforms/PerformanceInliner.cpp
@@ -387,6 +387,70 @@ bool SILPerformanceInliner::isTupleWithAllocsOrPartialApplies(SILValue val) {
   return false;
 }
 
+// Uses a function's mangled name to determine if it is an Autodiff VJP
+// function.
+//
+// TODO: VJPs of differentiable functions with custom silgen names are not
+// recognized as VJPs by this function. However, this is not a hard limitation
+// and can be fixed.
+bool isFunctionAutodiffVJP(SILFunction *callee) {
+  swift::Demangle::Context Ctx;
+  if (auto *Root = Ctx.demangleSymbolAsNode(callee->getName())) {
+    if (auto *node = Root->findByKind(
+            swift::Demangle::Node::Kind::AutoDiffFunctionKind, 3)) {
+      if (node->hasIndex()) {
+        auto index = (char)node->getIndex();
+        auto ADFunctionKind = swift::Demangle::AutoDiffFunctionKind(index);
+        if (ADFunctionKind == swift::Demangle::AutoDiffFunctionKind::VJP) {
+          return true;
+        }
+      }
+    }
+  }
+
+  return false;
+}
+
+bool isProfitableToInlineAutodiffVJP(SILFunction *vjp, SILFunction *caller,
+                                     InlineSelection whatToInline,
+                                     StringRef stageName) {
+  bool isLowLevelFunctionPassPipeline = stageName == "LowLevel,Function";
+  auto isHighLevelFunctionPassPipeline =
+      stageName == "HighLevel,Function+EarlyLoopOpt";
+  auto calleeHasControlFlow = vjp->size() > 1;
+  auto isCallerVJP = isFunctionAutodiffVJP(caller);
+  auto callerHasControlFlow = caller->size() > 1;
+
+  // If the pass is being run as part of the low-level function pass pipeline,
+  // the autodiff closure-spec optimization is done doing its work. Therefore,
+  // all VJPs should be considered for inlining.
+  if (isLowLevelFunctionPassPipeline) {
+    return true;
+  }
+
+  // If callee has control-flow it will definitely not be handled by the
+  // Autodiff closure-spec optimization. Therefore, we should consider it for
+  // inlining.
+  if (calleeHasControlFlow) {
+    return true;
+  }
+
+  // If this is the EarlyPerfInline pass we want to have the Autodiff
+  // closure-spec optimization pass optimize VJPs in isolation before they are
+  // inlined into other VJPs.
+  if (isHighLevelFunctionPassPipeline) {
+    return false;
+  }
+
+  // If this is not the EarlyPerfInline pass, VJPs should only be inlined into
+  // other VJPs that do not contain any control-flow.
+  if (!isCallerVJP || (isCallerVJP && callerHasControlFlow)) {
+    return false;
+  }
+
+  return true;
+}
+
 bool SILPerformanceInliner::isProfitableToInline(
     FullApplySite AI, Weight CallerWeight, ConstantTracker &callerTracker,
     int &NumCallerBlocks,
@@ -395,6 +459,12 @@ bool SILPerformanceInliner::isProfitableToInline(
   assert(Callee);
   bool IsGeneric = AI.hasSubstitutions();
 
+  if (isFunctionAutodiffVJP(Callee) &&
+      !isProfitableToInlineAutodiffVJP(Callee, AI.getFunction(), WhatToInline,
+                                       this->pm->getStageName())) {
+    return false;
+  }
+
   // Start with a base benefit.
   int BaseBenefit = isa<BeginApplyInst>(AI) ? RemovedCoroutineCallBenefit
                                             : RemovedCallBenefit;
diff --git a/test/AutoDiff/SILOptimizer/closure_specialization.sil b/test/AutoDiff/SILOptimizer/closure_specialization.sil
diff --git a/test/AutoDiff/SILOptimizer/pullback_inlining.swift b/test/AutoDiff/SILOptimizer/pullback_inlining.swift
diff --git a/test/AutoDiff/SILOptimizer/vjp_inlining.swift b/test/AutoDiff/SILOptimizer/vjp_inlining.swift

Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,10 @@ extension Context {`
`60`	`60`	`_bridged.lookupFunction($0).function`
`61`	`61`	`}`
`62`	`62`	`}`
	`63`	`+`
	`64`	`+ func notifyNewFunction(function: Function, derivedFrom: Function) {`
	`65`	`+ _bridged.addFunctionToPassManagerWorklist(function.bridged, derivedFrom.bridged)`
	`66`	`+ }`
`63`	`67`	`}`
`64`	`68`
`65`	`69`	`/// A context which allows mutation of a function's SIL.`
`@@ -357,7 +361,7 @@ struct FunctionPassContext : MutatingContext {`
`357`	`361`	`return String(taking: _bridged.mangleOutlinedVariable(function.bridged))`
`358`	`362`	`}`
`359`	`363`
`360`		`- func mangle(withClosureArgs closureArgs: [Value], closureArgIndices: [Int], from applySiteCallee: Function) -> String {`
	`364`	`+ func mangle(withClosureArguments closureArgs: [Value], closureArgIndices: [Int], from applySiteCallee: Function) -> String {`
`361`	`365`	`closureArgs.withBridgedValues { bridgedClosureArgsRef in`
`362`	`366`	`closureArgIndices.withBridgedArrayRef{bridgedClosureArgIndicesRef in`
`363`	`367`	`String(taking: _bridged.mangleWithClosureArgs(`
`@@ -392,13 +396,13 @@ struct FunctionPassContext : MutatingContext {`
`392`	`396`	`}`
`393`	`397`	`}`
`394`	`398`
`395`		`- func buildSpecializedFunction(specializedFunction: Function, buildFn: (Function, FunctionPassContext) -> ()) {`
	`399`	`+ func buildSpecializedFunction<T>(specializedFunction: Function, buildFn: (Function, FunctionPassContext) -> T) -> T {`
`396`	`400`	`let nestedFunctionPassContext =`
`397`	`401`	`FunctionPassContext(_bridged: _bridged.initializeNestedPassContext(specializedFunction.bridged))`
`398`	`402`
`399`	`403`	`defer { _bridged.deinitializedNestedPassContext() }`
`400`	`404`
`401`		`- buildFn(specializedFunction, nestedFunctionPassContext)`
	`405`	`+ return buildFn(specializedFunction, nestedFunctionPassContext)`
`402`	`406`	`}`
`403`	`407`	`}`
`404`	`408`
Original file line number	Diff line number	Diff line change
`@@ -167,7 +167,8 @@ public func registerOptimizerTests() {`
`167`	`167`	`parseTestSpecificationTest,`
`168`	`168`	`variableIntroducerTest,`
`169`	`169`	`gatherCallSitesTest,`
`170`		`- specializedFunctionSignatureAndBodyTest`
	`170`	`+ specializedFunctionSignatureAndBodyTest,`
	`171`	`+ rewrittenCallerBodyTest`
`171`	`172`	`)`
`172`	`173`
`173`	`174`	`// Finally register the thunk they all call through.`
Original file line number	Diff line number	Diff line change
`@@ -62,6 +62,10 @@ public struct OperandArray : RandomAccessCollection, CustomReflectable {`
`62`	`62`	`self.count = count`
`63`	`63`	`}`
`64`	`64`
	`65`	`+ static public var empty: OperandArray {`
	`66`	`+ OperandArray(base: OptionalBridgedOperand(bridged: nil), count: 0)`
	`67`	`+ }`
	`68`	`+`
`65`	`69`	`public var startIndex: Int { return 0 }`
`66`	`70`	`public var endIndex: Int { return count }`
`67`	`71`