[Autodiff] Adds logic to rewrite call-sites using functions specialized by the closure-spec optimization #73688

jkshtj · 2024-05-16T22:41:11Z

This PR adds the final part of the initial autodiff specific closure-specialization optimization implementation.

The previous parts of the optimization have been implemented in:

jkshtj · 2024-05-16T22:41:24Z

eeckstein · 2024-05-17T16:17:13Z

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

-  if !function.isAutodiffVJP {
-    return
+
+  if function.isOptimizable,


This is not needed because the pass manager doesn't invoke function passes for functions which shouldn't be optimized

eeckstein · 2024-05-17T16:19:15Z

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

+     !function.isAvailableExternally,
+     function.isAutodiffVJP,
+     function.hasSingleBlock
+  {


coding style: It's better to turn this into an early-exit if (or guard):

guard ... else { return }

eeckstein · 2024-05-17T16:24:06Z

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

+          var (specializedFunction, alreadyExists, invalidatedStackNestingWhileCreatingSpecializedFunc) = 
+            getOrCreateSpecializedFunction(basedOn: callSite, context)
+
+          invalidatedStackNesting = invalidatedStackNesting || invalidatedStackNestingWhileCreatingSpecializedFunc


You can use context.notifyInvalidatedStackNesting/context.needFixStackNesting instead. Then you don't need to pass around the flag

eeckstein · 2024-05-17T16:27:13Z

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

+            .forEach { deadClosures.insert($0) }
+        }
+
+        for deadClosure in deadClosures {


Please don't iterate over a Set. The order is not deterministic.
You can use an InstructionWorklist instead

eeckstein · 2024-05-17T16:49:29Z

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

+          function.fixStackNesting(context)
+        }
+
+        context.notifyFunctionBodyChanged()


You don't need this. Notifications are sent automatically if instructions are changed/inserted or deleted

eeckstein · 2024-05-17T17:04:12Z

SwiftCompilerSources/Sources/SIL/Function.swift

@@ -39,6 +39,8 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash

  public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() }

+  public var isOptimizable: Bool { bridged.isOptimizable() }


Beside that you don't need this, I would call it shouldOptimize and add a comment that this is false if the function has an @_optimize(none) attribute

Not needed after your suggested changes.

eeckstein · 2024-05-17T17:06:34Z

SwiftCompilerSources/Sources/SIL/Function.swift

@@ -64,6 +66,10 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
  public var blocks : BasicBlockList { BasicBlockList(first: bridged.getFirstBlock().block) }

  public var entryBlock: BasicBlock { blocks.first! }
+
+  public var hasSingleBlock: Bool {


You don't need this. Just use function.blocks.singleElement

eeckstein · 2024-05-17T17:07:23Z

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1634,6 +1634,9 @@ void BridgedChangeNotificationHandler::notifyChanges(Kind changeKind) const {
  case Kind::effectsChanged:
    invocation->notifyChanges(SILAnalysis::InvalidationKind::Effects);
    break;
+  case Kind::functionBodyChanged:


As I said in my other comment, you don't need this

eeckstein · 2024-05-17T17:10:26Z

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -735,6 +735,17 @@ static void addMidLevelFunctionPipeline(SILPassPipelinePlan &P) {
  P.addLoopUnroll();
 }

+static void addAutodiffClosureSpecializationPassPipeline(SILPassPipelinePlan &P) {
+  P.startPipeline("AutodiffClosureSpecialization");
+  P.addDeadStoreElimination();


Adding such a number of new passes in the pipeline is not completely unproblematic. We have to be careful to not increase compile time.
Is it really needed to add those passes? If yes, can you check the impact on compile time, e.g. by measuring the compile time of the stdlib core module.

I was modelling the pipeline after the existing closure-specialization pipeline and kept only the basic optimization passes that I knew could affect autodiff.

That said, those other passes are not strictly needed. I've removed them.

eeckstein · 2024-05-17T17:12:24Z

lib/SILOptimizer/PassManager/PassPipeline.cpp

+
+  P.addAutodiffClosureSpecialization();
+
+  P.addDeadFunctionAndGlobalElimination();


dead-function-elimination is especially problematic because it's a module pass. Module passes split the function pass pipeline and you can get a completely different optimization behavior. (Unfortunately this is not really visible and documented). Isn't it good enough that the later function-elimination pass cleans up dead functions?

Addressed in a previous comment. Removed extra optimization passes.

jkshtj · 2024-05-20T22:42:46Z

@swift-ci please test

jkshtj · 2024-05-20T23:08:26Z

@swift-ci please test

jkshtj · 2024-05-21T05:25:52Z

@swift-ci please test macos

eeckstein

lgtm!

eeckstein · 2024-05-21T07:09:21Z

SwiftCompilerSources/Sources/SIL/Function.swift

@@ -38,7 +38,7 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
  public var isTrapNoReturn: Bool { bridged.isTrapNoReturn() }

  public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() }
-
+  


accidental change?

Looks like it. Fixed now.

eeckstein · 2024-05-21T07:09:29Z

lib/SILOptimizer/Transforms/PerformanceInliner.cpp

@@ -215,7 +215,6 @@ class SILPerformanceInliner {
  bool isAutoDiffLinearMapWithControlFlow(FullApplySite AI);

  bool isTupleWithAllocsOrPartialApplies(SILValue retVal);
-


accidental change?

Looks like it. Fixed now.

jkshtj · 2024-05-21T15:09:02Z

@swift-ci please test macos

…ed by the closure-spec optimization

…he closure-spec optimization pass The OSSA elimination pass has not yet been moved below all high level function passes. Until that work has been completed the Autodiff closure-spec optimization pass cannot solely support OSSA instructions.

jkshtj · 2024-05-21T19:21:56Z

@swift-ci please test

The inlining benefit for VJPs and pullbacks seems to be a bit different on macos and linux. This difference seems to be arising due to certain function calls such as, `sin`, `cos` etc. always being inlined on macos but not on linux. Until I figure out a better way, I'm modifying these tests to avoid matching the exact value of the inlining benefit, which causes them to fail on Linux.

jkshtj · 2024-05-21T23:05:50Z

@swift-ci please test

jkshtj · 2024-05-22T02:20:17Z

@swift-ci please test macos

jkshtj · 2024-05-22T02:24:54Z

@swift-ci please test linux

jkshtj · 2024-05-22T14:57:39Z

@swift-ci please test linux

jkshtj · 2024-05-23T01:52:09Z

@swift-ci please test

…d" test failures on Linux builds

jkshtj · 2024-05-23T05:38:13Z

@swift-ci please test

jkshtj · 2024-05-23T15:12:37Z

@swift-ci please test macos

jkshtj · 2024-05-23T15:48:52Z

@shahmishal The macos CI builds seem to have been a bit flaky recently. I've had some rare successful builds but mostly failures. The logs seem to be scattered with messages like these:

Test Case '-[SKCoreTests.TaskSchedulerTests testTasksWithElevatedPrioritiesGetExecutedFirst]' started.
/Users/ec2-user/jenkins/workspace/swift-PR-macos/branch-main/sourcekit-lsp/Tests/SKCoreTests/TaskSchedulerTests.swift:89: error: -[SKCoreTests.TaskSchedulerTests testTasksWithElevatedPrioritiesGetExecutedFirst] : XCTAssertTrue failed - [Set([high(9)]), Set([])] did not fulfill predicate
Test Case '-[SKCoreTests.TaskSchedulerTests testTasksWithElevatedPrioritiesGetExecutedFirst]' failed (3.794 seconds).
Test Suite 'TaskSchedulerTests' failed at 2024-05-23 10:39:31.836.
	 Executed 1 test, with 1 failure (0 unexpected) in 3.794 (3.794) seconds
Test Suite 'SourceKitLSPPackageTests.xctest' failed at 2024-05-23 10:39:31.836.
	 Executed 1 test, with 1 failure (0 unexpected) in 3.794 (3.794) seconds
Test Suite 'Selected tests' failed at 2024-05-23 10:39:31.836.
	 Executed 1 test, with 1 failure (0 unexpected) in 3.794 (3.795) seconds

warning: --multiroot-data-file is an *unsupported* option which can be removed at any time; do not rely on it
warning: 'swift-format': found 1 file(s) which are unhandled; explicitly declare them as resources or exclude from the target
    /Users/ec2-user/jenkins/workspace/swift-PR-macos/branch-main/swift-format/Sources/SwiftFormat/CMakeLists.txt
warning: 'swift-format': found 1 file(s) which are unhandled; explicitly declare them as resources or exclude from the target
    /Users/ec2-user/jenkins/workspace/swift-PR-macos/branch-main/swift-format/Sources/SwiftFormat/CMakeLists.txt
warning: 'swift-format': found 1 file(s) which are unhandled; explicitly declare them as resources or exclude from the target
    /Users/ec2-user/jenkins/workspace/swift-PR-macos/branch-main/swift-format/Sources/swift-format/CMakeLists.txt
warning: 'swift-format': found 1 file(s) which are unhandled; explicitly declare them as resources or exclude from the target
    /Users/ec2-user/jenkins/workspace/swift-PR-macos/branch-main/swift-format/Sources/swift-format/CMakeLists.txt
warning: failed to retrieve search paths with pkg-config; maybe pkg-config is not installed
warning: couldn't find pc file for sqlite3

[org.swift.sourcekit-lsp:default] error 2024-05-23 10:54:15.7830 +0000
/Users/ec2-user/jenkins/workspace/swift-PR-macos/branch-main/build/buildbot_incremental/toolchain-macosx-x86_64/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/Info.plist doesn't exist in file system

Do you know what's going on here?

jkshtj requested a review from eeckstein as a code owner May 16, 2024 22:41

eeckstein reviewed May 17, 2024

View reviewed changes

jkshtj force-pushed the main branch from 571a796 to 51b9628 Compare May 20, 2024 22:42

jkshtj force-pushed the main branch from 51b9628 to e5e129a Compare May 20, 2024 23:08

eeckstein approved these changes May 21, 2024

View reviewed changes

jkshtj added 2 commits May 21, 2024 12:02

[Autodiff] Adds logic to rewrite call-sites using functions specializ…

12faf79

…ed by the closure-spec optimization

jkshtj force-pushed the main branch from e5e129a to c288d95 Compare May 21, 2024 19:19

jkshtj force-pushed the main branch from 993f7c3 to 8595e64 Compare May 23, 2024 01:51

[Autodiff] Fixes bugs in closure-spec opt that were causing "optimize…

487648a

…d" test failures on Linux builds

jkshtj force-pushed the main branch from 8595e64 to 487648a Compare May 23, 2024 05:37

jkshtj merged commit 64da348 into apple:main May 23, 2024
5 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Autodiff] Adds logic to rewrite call-sites using functions specialized by the closure-spec optimization #73688

[Autodiff] Adds logic to rewrite call-sites using functions specialized by the closure-spec optimization #73688

jkshtj commented May 16, 2024

jkshtj commented May 16, 2024

eeckstein May 17, 2024

eeckstein May 17, 2024

eeckstein May 17, 2024

eeckstein May 17, 2024

eeckstein May 17, 2024

eeckstein May 17, 2024

jkshtj May 20, 2024

eeckstein May 17, 2024

eeckstein May 17, 2024

eeckstein May 17, 2024

jkshtj May 20, 2024

eeckstein May 17, 2024

jkshtj May 20, 2024

jkshtj commented May 20, 2024

jkshtj commented May 20, 2024

jkshtj commented May 21, 2024

eeckstein left a comment

eeckstein May 21, 2024

jkshtj May 21, 2024

eeckstein May 21, 2024

jkshtj May 21, 2024

jkshtj commented May 21, 2024

jkshtj commented May 21, 2024

jkshtj commented May 21, 2024

jkshtj commented May 22, 2024

jkshtj commented May 22, 2024

jkshtj commented May 22, 2024

jkshtj commented May 23, 2024

jkshtj commented May 23, 2024

jkshtj commented May 23, 2024

jkshtj commented May 23, 2024

		@@ -39,6 +39,8 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash

		public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() }

		public var isOptimizable: Bool { bridged.isOptimizable() }


		P.addAutodiffClosureSpecialization();

		P.addDeadFunctionAndGlobalElimination();

		@@ -38,7 +38,7 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
		public var isTrapNoReturn: Bool { bridged.isTrapNoReturn() }

		public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() }

		@@ -215,7 +215,6 @@ class SILPerformanceInliner {
		bool isAutoDiffLinearMapWithControlFlow(FullApplySite AI);

		bool isTupleWithAllocsOrPartialApplies(SILValue retVal);

[Autodiff] Adds logic to rewrite call-sites using functions specialized by the closure-spec optimization #73688

[Autodiff] Adds logic to rewrite call-sites using functions specialized by the closure-spec optimization #73688

Conversation

jkshtj commented May 16, 2024

jkshtj commented May 16, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkshtj commented May 20, 2024

jkshtj commented May 20, 2024

jkshtj commented May 21, 2024

eeckstein left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkshtj commented May 21, 2024

jkshtj commented May 21, 2024

jkshtj commented May 21, 2024

jkshtj commented May 22, 2024

jkshtj commented May 22, 2024

jkshtj commented May 22, 2024

jkshtj commented May 23, 2024

jkshtj commented May 23, 2024

jkshtj commented May 23, 2024

jkshtj commented May 23, 2024