Skip to content

[AutoDiff] Run AutoDiff closure spec pass for all VJPs #81548

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from

Conversation

kovdan01
Copy link
Contributor

@kovdan01 kovdan01 commented May 16, 2025

Previously, AutoDiff closure specialization pass was triggered only on VJPs containing single basic block. However, the pass logic allows running on arbitrary VJPs. This PR enables the pass for all VJPs unconditionally. So, if the pullback corresponding to multiple-BB VJP accepts some closures directly as arguments, these closures might become specialized by the pass. Closures passed via payload of branch tracing enum are not specialized - this is subject for future changes.

The PR contains several commits.

  1. The thing named "call site" in the code is partial_apply of pullback corresponding to the VJP. This might appear only once, so we drop support for multiple "call sites".
  2. Enhance existing SILOptimizer tests for the pass.
  3. Add validation-tests for single basic block case.
  4. The change itself - delete check against single basic block.
  5. Add validation-tests for multiple basic block case.
  6. Add SILOptimizer tests for multiple basic block case.

@kovdan01 kovdan01 requested a review from asl May 16, 2025 00:12
@kovdan01
Copy link
Contributor Author

Tagging @JaapWijnen

@clackary Could you please run your tests against the toolchain built from mainline with this PR applied?

@kovdan01 kovdan01 marked this pull request as ready for review May 16, 2025 00:15
@kovdan01 kovdan01 requested a review from eeckstein as a code owner May 16, 2025 00:15
@kovdan01
Copy link
Contributor Author

@swift-ci please test

@kovdan01 kovdan01 force-pushed the autodiff-closure-specialization-run-unconditionally branch 2 times, most recently from 0b44392 to 5beb934 Compare May 16, 2025 12:51
@kovdan01
Copy link
Contributor Author

@swift-ci please test

@kovdan01 kovdan01 force-pushed the autodiff-closure-specialization-run-unconditionally branch from 5beb934 to 72a4006 Compare May 18, 2025 23:53
@kovdan01
Copy link
Contributor Author

@swift-ci please test

@kovdan01 kovdan01 force-pushed the autodiff-closure-specialization-run-unconditionally branch from 72a4006 to 3f84308 Compare May 19, 2025 07:43
@kovdan01
Copy link
Contributor Author

@swift-ci please test

@kovdan01
Copy link
Contributor Author

preset=buildbot,tools=RA,stdlib=DA
@swift-ci please test with preset

@clackary
Copy link

@kovdan01 I cherrypicked this onto swift-DEVELOPMENT-SNAPSHOT-2025-05-14-a tag, built a toolchain, and can compile all our internal projects that import _Differentiation in release mode. Looks good from that perspective.

@kovdan01
Copy link
Contributor Author

@kovdan01 I cherrypicked this onto swift-DEVELOPMENT-SNAPSHOT-2025-05-14-a tag, built a toolchain, and can compile all our internal projects that import _Differentiation in release mode. Looks good from that perspective.

@clackary Thanks! Could you please also run tests so we ensure that nothing becomes silently miscompiled?

@clackary
Copy link

@clackary Thanks! Could you please also run tests so we ensure that nothing becomes silently miscompiled?

Ah, I should have specified. My tests include both compilation and running unit tests for those projects that are configured for it. We don't have unit tests configured to run for every single project, but most of the autodiff projects are set up this way.

@kovdan01
Copy link
Contributor Author

@asl FYI: Swift Test Linux Platform with preset looks unrelated to the changes introduced in this PR. Other build jobs are successful.

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.

2 participants