Skip to content

[AutoDiff] Fix autodiff_function_extract operand ownership kind and memory leaks. #27199

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.

Sign up for GitHub

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
rxwei merged 1 commit into from
Sep 16, 2019
Merged

[AutoDiff] Fix autodiff_function_extract operand ownership kind and memory leaks. #27199

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion lib/SIL/ValueOwnership.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,8 @@ CONSTANT_OWNERSHIP_INST(Unowned, ValueToBridgeObject)
}
CONSTANT_OR_TRIVIAL_OWNERSHIP_INST(Guaranteed, StructExtract)
CONSTANT_OR_TRIVIAL_OWNERSHIP_INST(Guaranteed, TupleExtract)
// SWIFT_ENABLE_TENSORFLOW
CONSTANT_OR_TRIVIAL_OWNERSHIP_INST(Guaranteed, AutoDiffFunctionExtract)
// OpenExistentialValue opens the boxed value inside an existential
// CoW box. The semantics of an existential CoW box implies that we
// can only consume the projected value inside the box if the box is
Expand Down Expand Up @@ -254,7 +256,6 @@ FORWARDING_OWNERSHIP_INST(SelectEnum)
FORWARDING_OWNERSHIP_INST(Enum)
// SWIFT_ENABLE_TENSORFLOW
FORWARDING_OWNERSHIP_INST(AutoDiffFunction)
FORWARDING_OWNERSHIP_INST(AutoDiffFunctionExtract)
#undef FORWARDING_OWNERSHIP_INST

ValueOwnershipKind
Expand Down
20 changes: 14 additions & 6 deletions lib/SILOptimizer/Mandatory/Differentiation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3749,9 +3749,14 @@ class VJPEmitter final
context.getResultIndices()[autoDiffFuncInst] =
activeResultIndices.front();

vjpValue = getBuilder().createAutoDiffFunctionExtract(
auto borrowedADFunc =
builder.emitBeginBorrowOperation(loc, autoDiffFuncInst);
auto extractedVJP = getBuilder().createAutoDiffFunctionExtract(
loc, AutoDiffFunctionExtractInst::Extractee::VJP,
/*differentiationOrder*/ 1, autoDiffFuncInst);
/*differentiationOrder*/ 1, borrowedADFunc);
vjpValue = builder.emitCopyValueOperation(loc, extractedVJP);
builder.emitEndBorrowOperation(loc, borrowedADFunc);
builder.emitDestroyValueOperation(loc, autoDiffFuncInst);
}

// Record desired/actual VJP indices.
Expand Down Expand Up @@ -4850,7 +4855,6 @@ class JVPEmitter final
// on the remapped original function operand and `autodiff_function_extract`
// the JVP. The actual JVP functions will be populated in the
// `autodiff_function` during the transform main loop.
SILValue differentiableFunc;
if (!jvpValue) {
// FIXME: Handle indirect differentiation invokers. This may require some
// redesign: currently, each original function + attribute pair is mapped
Expand Down Expand Up @@ -4911,16 +4915,20 @@ class JVPEmitter final
auto *autoDiffFuncInst =
context.createAutoDiffFunction(builder, loc, indices.parameters,
/*differentiationOrder*/ 1, original);
differentiableFunc = autoDiffFuncInst;

// Record the `autodiff_function` instruction.
context.getAutoDiffFunctionInsts().push_back(autoDiffFuncInst);
context.getResultIndices()[autoDiffFuncInst] =
activeResultIndices.front();

jvpValue = builder.createAutoDiffFunctionExtract(
auto borrowedADFunc =
builder.emitBeginBorrowOperation(loc, autoDiffFuncInst);
auto extractedJVP = builder.createAutoDiffFunctionExtract(
loc, AutoDiffFunctionExtractInst::Extractee::JVP,
/*differentiationOrder*/ 1, autoDiffFuncInst);
/*differentiationOrder*/ 1, borrowedADFunc);
jvpValue = builder.emitCopyValueOperation(loc, extractedJVP);
builder.emitEndBorrowOperation(loc, borrowedADFunc);
builder.emitDestroyValueOperation(loc, autoDiffFuncInst);
}

// Call the JVP using the original parameters.
Expand Down