-
Notifications
You must be signed in to change notification settings - Fork 10.5k
[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
Conversation
2447939
to
4ba9e7a
Compare
69d0849
to
6e948f9
Compare
…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).
…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).
@rxwei question. How many bugs have you guys found so far through this? 3? |
@gottesmm More than 5 and counting :) |
023b9cb
to
4ed2883
Compare
76f43b2
to
ad8144f
Compare
@@ -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)) { |
There was a problem hiding this comment.
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.
@swift-ci please test tensorflow |
1 similar comment
@swift-ci please test tensorflow |
There was a problem hiding this 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>()) { |
There was a problem hiding this comment.
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
.
@swift-ci please test tensorflow |
@swift-ci please test tensorflow |
3 similar comments
@swift-ci please test tensorflow |
@swift-ci please test tensorflow |
@swift-ci please test tensorflow |
Moves Differentiation before OwnershipModelEliminator. Now Differentiation happens right after DefiniteInitialization.
Changes
Activity analysis
store_borrow
instructions.mayReadFromMemory()
case. Remove special handlingmayReadFromMemory()
and propagate usefulness through all buffer operands instead.destructure_tuple
instead oftuple_extract
in array literal initialization pattern matcher.autodiff_function
canonicalizationcopyParameterArgumentsForApply
, used inreapplyFunctionConversions
and the curry thunk cloning logic inADContext::promoteToDifferentiableFunction
.VJPEmitter
@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 intrampolinedGuaranteedPhiArguments
and emitend_borrow
s when function cloning is finished.cond_br
instructions to conform to ownership rules.PullbackEmitter
@owned
ownership, for both function arguments and phi arguments.@owned
ownership.@guaranteed
arguments. We should change this in the future so that all calls to pullbacks consume their arguments. (TF-761)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.destructure_tuple
instead oftuple_extract
in array literal initialization adjoint emission logic.destructure_tuple
,load_borrow
,store_borrow
,copy_value
, andbegin_borrow
.Resolves TF-709 and TF-585. Also improves some source locations in diagnostics since the differentiation transform runs on a higher-level IR.