Skip to content

[AutoDiff] Move Differentiation before OwnershipModelEliminator. #26157

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 2 commits into from
Aug 18, 2019

Conversation

rxwei
Copy link
Contributor

@rxwei rxwei commented Jul 16, 2019

Moves Differentiation before OwnershipModelEliminator. Now Differentiation happens right after DefiniteInitialization.

Changes

Activity analysis

  • Propagate usefulness and variedness through store_borrow instructions.
  • Fix a bug where standard usefulness does not get propagated if the instruction goes through the mayReadFromMemory() case. Remove special handling mayReadFromMemory() and propagate usefulness through all buffer operands instead.
  • Handle destructure_tuple instead of tuple_extract in array literal initialization pattern matcher.

autodiff_function canonicalization

  • Rewrite argument cloning logic as copyParameterArgumentsForApply, used in reapplyFunctionConversions and the curry thunk cloning logic in ADContext::promoteToDifferentiableFunction.

VJPEmitter

  • Each @guaranteed trampoline argument needs to have a lifetime-ending use past its destination argument's lifetime-ending uses, so we keep track of these pairs of arguments in trampolinedGuaranteedPhiArguments and emit end_borrows when function cloning is finished.
  • Create trampoline blocks for cond_br instructions to conform to ownership rules.

PullbackEmitter

  • Make pullback struct arguments have @owned ownership, for both function arguments and phi arguments.
  • Make all other phi arguments also have @owned ownership.
    • Note: Linear maps get evaluated linearly, so all values should get consumed immediately when a pullback is called. This is not the case yet since pullback functions have @guaranteed arguments. We should change this in the future so that all calls to pullbacks consume their arguments. (TF-761)
  • Emit a switch_enum even if there is only one successor. It no longer triggers any verification failure.
    • Remove the TF-585 workaround (emitting a fix_lifetime on boxed enums) since the crasher in AllocBoxToStack is no longer reproducible.
  • Handle destructure_tuple instead of tuple_extract in array literal initialization adjoint emission logic.
  • Add pullback emission visitors for destructure_tuple, load_borrow, store_borrow, copy_value, and begin_borrow.

Resolves TF-709 and TF-585. Also improves some source locations in diagnostics since the differentiation transform runs on a higher-level IR.

@rxwei rxwei added the tensorflow This is for "tensorflow" branch PRs. label Jul 16, 2019
@rxwei rxwei requested review from gottesmm and dan-zheng July 16, 2019 03:03
@rxwei rxwei force-pushed the ownership branch 4 times, most recently from 2447939 to 4ba9e7a Compare July 22, 2019 10:34
@rxwei rxwei force-pushed the ownership branch 3 times, most recently from 69d0849 to 6e948f9 Compare July 26, 2019 12:12
rxwei added a commit to rxwei/swift that referenced this pull request Jul 28, 2019
…nks.

In `ADContext::promoteToDifferentiableFunction` when we emit a subset parameters thunk, the closure produced by `emitAssociatedFunctionReference` is not being used or released. This patch fixes that by releasing the closure.

In a future design, we should determine the actual parameter indices and whether thunking is needed before emitting an unused closure. This shall be addressed in a future patch.

A value leak checking test was added for protocols because this bug was originally discovered in eaplatanios/swift-ale#1 by using the newly split `Layer` protocol. However, value leak checking tests do not test for closure (`partial_apply`) leaks. The proper setup for catching closure leaks is moving AD before ownership stripping (swiftlang#26157).
rxwei added a commit that referenced this pull request Jul 28, 2019
…nks. (#26384)

In `ADContext::promoteToDifferentiableFunction` when we emit a subset parameters thunk, the closure produced by `emitAssociatedFunctionReference` is not being used or released. This patch fixes that by releasing the closure.

In a future design, we should determine the actual parameter indices and whether thunking is needed before emitting an unused closure. This shall be addressed in a future patch.

A value leak checking test was added for protocols because this bug was originally discovered in eaplatanios/swift-ale#1 by using the newly split `Layer` protocol. However, value leak checking tests do not test for closure (`partial_apply`) leaks. The proper setup for catching closure leaks is moving AD before ownership stripping (#26157).
@gottesmm
Copy link
Contributor

gottesmm commented Aug 5, 2019

@rxwei question. How many bugs have you guys found so far through this? 3?

@rxwei
Copy link
Contributor Author

rxwei commented Aug 8, 2019

@gottesmm More than 5 and counting :)

@rxwei rxwei force-pushed the ownership branch 3 times, most recently from 023b9cb to 4ed2883 Compare August 15, 2019 01:58
@rxwei rxwei force-pushed the ownership branch 2 times, most recently from 76f43b2 to ad8144f Compare August 17, 2019 02:09
@rxwei rxwei requested a review from bartchr808 August 17, 2019 03:12
@rxwei rxwei marked this pull request as ready for review August 17, 2019 03:15
@@ -4908,14 +5016,19 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
auto activeValueAdj = getAdjointValue(origBB, activeValue);
auto concreteActiveValueAdj =
materializeAdjointDirect(activeValueAdj, loc);
trampolineArguments.push_back(concreteActiveValueAdj);

if (!adjointTrampolineBlockMap.count(concreteActiveValueAdj)) {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

This logic is hacky and is not straightforward to describe. Fortunately it will be defined away with the upcoming ad-all-indirect that @dan-zheng is working on.

@rxwei
Copy link
Contributor Author

rxwei commented Aug 17, 2019

@swift-ci please test tensorflow

1 similar comment
@rxwei
Copy link
Contributor Author

rxwei commented Aug 17, 2019

@swift-ci please test tensorflow

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.

Congrats on this big undertaking and permanently improving robustness against memory leaks!
LGTM if tests pass 🙂

for (auto i : range(tupleType->getNumElements()))
result.push_back(builder.createTupleExtract(val.getLoc(), val, i));
else
if (auto tupleType = val->getType().getAs<TupleType>()) {
Copy link
Contributor

Choose a reason for hiding this comment

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

Just FYI: another version of extractAllElements lives in SILGenPoly.cpp.
It calls SILBuilder:: emitDestructureValueOperation.

@rxwei
Copy link
Contributor Author

rxwei commented Aug 18, 2019

@swift-ci please test tensorflow

@rxwei
Copy link
Contributor Author

rxwei commented Aug 18, 2019

@swift-ci please test tensorflow

3 similar comments
@rxwei
Copy link
Contributor Author

rxwei commented Aug 18, 2019

@swift-ci please test tensorflow

@rxwei
Copy link
Contributor Author

rxwei commented Aug 18, 2019

@swift-ci please test tensorflow

@rxwei
Copy link
Contributor Author

rxwei commented Aug 18, 2019

@swift-ci please test tensorflow

@rxwei rxwei merged commit 5dd9c50 into swiftlang:tensorflow Aug 18, 2019
@rxwei rxwei deleted the ownership branch August 19, 2019 08:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
tensorflow This is for "tensorflow" branch PRs.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants