Skip to content

[AutoDiff] Do not propagate same adjoint buffer multiple times #64963

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
Apr 6, 2023

Conversation

asl
Copy link
Contributor

@asl asl commented Apr 6, 2023

Adjoint buffers of projections (e.g. obtained via begin_access) are same as adjoint buffer of underlying struct value. As a result, when propagating adjoint values to pullback successor blocks we tend to produce lots of identical copies (essentially for every struct access and in every basic block) of adjoint buffers.

These copy_addrs instructions are then lowered down to plain loads and stores and while the redundant copies are usually optimized away by subsequent optimization passes, presence of such copies leads to elevated memory consumption and compilation time as one needs to track liveness of these values being copied.

Track the values being propagated and simply do not generate extra copies if the same value was already propagated.

One step towards #61773

@asl asl requested review from rxwei and dan-zheng April 6, 2023 01:06
@asl
Copy link
Contributor Author

asl commented Apr 6, 2023

Tagging @BradLarson

@asl
Copy link
Contributor Author

asl commented Apr 6, 2023

@swift-ci please test

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.

Nice fix!

(Duplicate copy_addr for adjoint buffers has been a longstanding control flow autodiff issue, huh?)

@asl
Copy link
Contributor Author

asl commented Apr 6, 2023

@dan-zheng This is only tip of the iceberg as we need to collect the cases that impact performance. Here essentially every struct access would generate new adjoint buffer copy for every basic block, so number of copies might scale quadratically as # struct accesses x # of bb's Later these copy_addrs would be exploded to multiple instructions leading to even more "fun" SIL / LLVM IR to optimize. Poor GVN :)

@asl asl force-pushed the adj-redundant-copies branch from 8af4f12 to 7e41029 Compare April 6, 2023 03:22
@asl
Copy link
Contributor Author

asl commented Apr 6, 2023

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Apr 6, 2023

@swift-ci please test macos

@asl asl merged commit 9ee8c73 into main Apr 6, 2023
@asl asl deleted the adj-redundant-copies branch April 6, 2023 23:49
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.

3 participants