Skip to content

[AutoDiff] Speed up compile time by 3.5x by caching derivative function types. #29590

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 1, 2020

Conversation

rxwei
Copy link
Contributor

@rxwei rxwei commented Feb 1, 2020

Compiling packages that use AutoDiff is currently extremely slow. For example, compiling tensorflow/swift-apis takes ~5 minutes on a slower machine. This is largely due to calls to SILFunctionType::getAutoDiffDerivativeFunctionType in SIL verifier, which spends a long time computing generic signatures with GSB.

This PR improves this situation by caching computed derivative function types in ASTContext. This resulted in a 1.6x speedup in running the AutoDiff test suite, and a 3.5x speedup in compiling tensorflow/swift-apis.

Future investigations:

  • Cache generic signatures computed from getAutoDiffDerivativeFunctionGenericSignature.
  • Cache transpose function types computed from SILFunctionType::getAutoDiffTransposeFunctionType().

Here are the results building tensorflow/swift-apis with -j56.

Before:

swift build  411.45s user 31.20s system 334% cpu 2:12.52 total

After:

swift build  154.03s user 28.26s system 485% cpu 37.519 total

Resolves TF-1133.

@rxwei rxwei requested review from dan-zheng and marcrasi February 1, 2020 09:53
@rxwei
Copy link
Contributor Author

rxwei commented Feb 1, 2020

@swift-ci please test tensorflow

@rxwei
Copy link
Contributor Author

rxwei commented Feb 1, 2020

@swift-ci please test tensorflow linux gpu

Copy link
Contributor

@dan-zheng dan-zheng left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for identifying and fixing the issue!

@dan-zheng
Copy link
Contributor

Linux GPU CI failures are temporarily expected. Merging now so new toolchains can be built with the compilation time fix.

@dan-zheng dan-zheng merged commit ba93f64 into swiftlang:tensorflow Feb 1, 2020
@compnerd
Copy link
Member

compnerd commented Feb 1, 2020

rxwei added a commit to rxwei/swift that referenced this pull request Feb 1, 2020
@rxwei
Copy link
Contributor Author

rxwei commented Feb 1, 2020

My apologies. I just sent PR #29592 your way.

@rxwei rxwei deleted the TF-1133 branch February 1, 2020 23:40
compnerd added a commit that referenced this pull request Feb 2, 2020
[AutoDiff] Fix windows build failure introduced in #29590.
@marcrasi
Copy link

marcrasi commented Feb 3, 2020

Very nice drop on our graph of completime!!

j3sAFGwhQJB

@rxwei
Copy link
Contributor Author

rxwei commented Feb 3, 2020

Nice graph. This shows exactly when the slowdown occurred: It's when we switched to SILDifferentiabilityWitness for derivative registration.

dan-zheng added a commit to dan-zheng/swift that referenced this pull request Feb 20, 2020
dan-zheng added a commit that referenced this pull request Feb 20, 2020
Upstream #29590:
cache `SILFunctionType::getAutoDiffDerivativeFunctionType` results.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants