Skip to content

[AutoDiff] Partially fix AD control flow memory leaks #25294

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
merged 5 commits into from
Jun 7, 2019
Merged

[AutoDiff] Partially fix AD control flow memory leaks #25294

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a0a1bed
Make pullback structs the last BB argument.
rxwei Jun 6, 2019
87aeb91
Reduce memory leaks by applying cleanups in each adjoint block.
rxwei Jun 6, 2019
2d1d637
Merge branch 'tensorflow' of github.com:apple/swift into memory-leak-fix
rxwei Jun 6, 2019
1577947
Fix tests.
rxwei Jun 7, 2019
60e2c01
Add cleanup tests.
rxwei Jun 7, 2019
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
28 changes: 18 additions & 10 deletions lib/SILOptimizer/Mandatory/Differentiation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4106,6 +4106,10 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
}

public:
//--------------------------------------------------------------------------//
// Entry point
//--------------------------------------------------------------------------//

/// Performs adjoint synthesis on the empty adjoint function. Returns true if
/// any error occurs.
bool run() {
Expand Down Expand Up @@ -4187,11 +4191,6 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
continue;
}

// Otherwise, we create a phi argument for the corresponding pullback
// struct, and handle dominated active values/buffers.
auto *pbStructArg = adjointBB->createPhiArgument(
pbStructLoweredType, ValueOwnershipKind::Guaranteed);
adjointPullbackStructArguments[origBB] = pbStructArg;
// Get all active values in the original block.
// If the original block has no active values, continue.
auto &bbActiveValues = activeValues[origBB];
Expand All @@ -4217,6 +4216,10 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
activeValueAdjointBBArgumentMap[{origBB, activeValue}] = adjointArg;
}
}
// Add a pullback struct argument.
auto *pbStructArg = adjointBB->createPhiArgument(
pbStructLoweredType, ValueOwnershipKind::Guaranteed);
adjointPullbackStructArguments[origBB] = pbStructArg;
// - Create adjoint trampoline blocks for each successor block of the
// original block. Adjoint trampoline blocks only have a pullback
// struct argument, and branch from the adjoint successor block to the
Expand Down Expand Up @@ -4364,8 +4367,6 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
assert(adjointSuccBB && adjointSuccBB->getNumArguments() == 1);
SILBuilder adjointTrampolineBBBuilder(adjointSuccBB);
SmallVector<SILValue, 8> trampolineArguments;
// Propagate pullback struct argument.
trampolineArguments.push_back(adjointSuccBB->getArguments().front());
// Propagate adjoint values/buffers of active values/buffers to
// predecessor blocks.
auto &predBBActiveValues = activeValues[predBB];
Expand All @@ -4374,6 +4375,11 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
auto activeValueAdj = getAdjointValue(bb, activeValue);
auto concreteActiveValueAdj =
materializeAdjointDirect(activeValueAdj, adjLoc);
// Emit cleanups for children.
if (auto *cleanup = concreteActiveValueAdj.getCleanup()) {
cleanup->disable();
cleanup->applyRecursively(builder, activeValue.getLoc());
}
trampolineArguments.push_back(concreteActiveValueAdj);
// If the adjoint block does not yet have a registered adjoint
// value for the active value, set the adjoint value to the
Expand All @@ -4383,9 +4389,9 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
if (!hasAdjointValue(predBB, activeValue)) {
auto *adjointBBArg =
getActiveValueAdjointBlockArgument(predBB, activeValue);
// FIXME: Propagate cleanups to fix memory leaks.
auto forwardedArgAdj =
makeConcreteAdjointValue(ValueWithCleanup(adjointBBArg));
auto forwardedArgAdj = makeConcreteAdjointValue(
ValueWithCleanup(adjointBBArg,
makeCleanup(adjointBBArg, emitCleanup)));
initializeAdjointValue(predBB, activeValue, forwardedArgAdj);
}
} else {
Expand All @@ -4399,6 +4405,8 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
adjLoc, adjBuf, predAdjBuf, IsNotTake, IsNotInitialization);
}
}
// Propagate pullback struct argument.
trampolineArguments.push_back(adjointSuccBB->getArguments().front());
// Branch from adjoint trampoline block to adjoint block.
adjointTrampolineBBBuilder.createBranch(
adjLoc, adjointBB, trampolineArguments);
Expand Down
39 changes: 25 additions & 14 deletions test/AutoDiff/control_flow_sil.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,32 +72,43 @@ func cond(_ x: Float) -> Float {
// CHECK-SIL: switch_enum [[BB3_PRED]] : $_AD__cond_bb3__Pred__src_0_wrt_0, case #_AD__cond_bb3__Pred__src_0_wrt_0.bb2!enumelt.1: bb3, case #_AD__cond_bb3__Pred__src_0_wrt_0.bb1!enumelt.1: bb1

// CHECK-SIL: bb1([[BB3_PRED1_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_bb1__PB__src_0_wrt_0):
// CHECK-SIL: br bb2([[BB3_PRED1_TRAMP_PB_STRUCT]] : $_AD__cond_bb1__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}}: $Float)
// CHECK-SIL: br bb2({{%.*}} : $Float, {{%.*}}: $Float, [[BB3_PRED1_TRAMP_PB_STRUCT]] : $_AD__cond_bb1__PB__src_0_wrt_0)

// CHECK-SIL: bb2([[BB1_PB_STRUCT:%.*]] : $_AD__cond_bb1__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}} : $Float):
// CHECK-SIL: bb2({{%.*}} : $Float, {{%.*}} : $Float, [[BB1_PB_STRUCT:%.*]] : $_AD__cond_bb1__PB__src_0_wrt_0):
// CHECK-SIL: [[BB1_PB:%.*]] = struct_extract [[BB1_PB_STRUCT]]
// CHECK-SIL: [[BB1_ADJVALS:%.*]] = apply [[BB1_PB]]([[SEED]]) : $@callee_guaranteed (Float) -> (Float, Float)
// CHECK-SIL: [[BB1_PRED:%.*]] = struct_extract [[BB1_PB_STRUCT]]
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: [[BB1_PB_STRUCT_DATA:%.*]] = unchecked_enum_data [[BB1_PRED]]
// CHECK-SIL: br bb5([[BB1_PB_STRUCT_DATA]] : $_AD__cond_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb3([[BB3_PRED2_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_bb2__PB__src_0_wrt_0):
// CHECK-SIL: br bb4([[BB3_PRED2_TRAMP_PB_STRUCT]] : $_AD__cond_bb2__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}}: $Float)
// CHECK-SIL: br bb4({{%.*}} : $Float, {{%.*}}: $Float, [[BB3_PRED2_TRAMP_PB_STRUCT]] : $_AD__cond_bb2__PB__src_0_wrt_0)

// CHECK-SIL: bb4([[BB2_PB_STRUCT:%.*]] : $_AD__cond_bb2__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}} : $Float):
// CHECK-SIL: bb4({{%.*}} : $Float, {{%.*}} : $Float, [[BB2_PB_STRUCT:%.*]] : $_AD__cond_bb2__PB__src_0_wrt_0):
// CHECK-SIL: [[BB2_PB:%.*]] = struct_extract [[BB2_PB_STRUCT]]
// CHECK-SIL: [[BB2_ADJVALS:%.*]] = apply [[BB2_PB]]([[SEED]]) : $@callee_guaranteed (Float) -> (Float, Float)
// CHECK-SIL: [[BB2_PRED:%.*]] = struct_extract [[BB2_PB_STRUCT]]
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: [[BB2_PB_STRUCT_DATA:%.*]] = unchecked_enum_data [[BB2_PRED]]
// CHECK-SIL: br bb6([[BB2_PB_STRUCT_DATA]] : $_AD__cond_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb5([[BB1_PRED0_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_bb0__PB__src_0_wrt_0):
// CHECK-SIL: br bb7([[BB1_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_bb0__PB__src_0_wrt_0, {{%.*}} : $Float)
// CHECK-SIL: br bb7({{%.*}} : $Float, [[BB1_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb6([[BB2_PRED0_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_bb0__PB__src_0_wrt_0):
// CHECK-SIL: br bb7([[BB2_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_bb0__PB__src_0_wrt_0, {{%.*}} : $Float)
// CHECK-SIL: br bb7({{%.*}} : $Float, [[BB2_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb7([[BB0_PB_STRUCT:%.*]] : $_AD__cond_bb0__PB__src_0_wrt_0, {{%.*}} : $Float):
// CHECK-SIL: bb7({{%.*}} : $Float, [[BB0_PB_STRUCT:%.*]] : $_AD__cond_bb0__PB__src_0_wrt_0):
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: release_value {{%.*}} : $Float
// CHECK-SIL: return {{%.*}} : $Float

@differentiable
Expand Down Expand Up @@ -147,30 +158,30 @@ func cond_tuple_var(_ x: Float) -> Float {
// CHECK-SIL: switch_enum [[BB3_PRED]] : $_AD__cond_tuple_var_bb3__Pred__src_0_wrt_0, case #_AD__cond_tuple_var_bb3__Pred__src_0_wrt_0.bb2!enumelt.1: bb3, case #_AD__cond_tuple_var_bb3__Pred__src_0_wrt_0.bb1!enumelt.1: bb1

// CHECK-SIL: bb1([[BB3_PRED1_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb1__PB__src_0_wrt_0):
// CHECK-SIL: br bb2([[BB3_PRED1_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb1__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}}: $Float)
// CHECK-SIL: br bb2({{%.*}} : $Float, {{%.*}} : $Float, [[BB3_PRED1_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb1__PB__src_0_wrt_0)

// CHECK-SIL: bb2([[BB1_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb1__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}} : $Float):
// CHECK-SIL: bb2({{%.*}} : $Float, {{%.*}} : $Float, [[BB1_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb1__PB__src_0_wrt_0):
// CHECK-SIL: [[BB1_PRED:%.*]] = struct_extract [[BB1_PB_STRUCT]]
// CHECK-SIL: copy_addr {{%.*}} to {{%.*}} : $*(Float, Float)
// CHECK-SIL-NOT: copy_addr {{%.*}} to {{%.*}} : $*Float
// CHECK-SIL: [[BB1_PB_STRUCT_DATA:%.*]] = unchecked_enum_data [[BB1_PRED]]
// CHECK-SIL: br bb5([[BB1_PB_STRUCT_DATA]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb3([[BB3_PRED2_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb2__PB__src_0_wrt_0):
// CHECK-SIL: br bb4([[BB3_PRED2_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb2__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}}: $Float)
// CHECK-SIL: br bb4({{%.*}} : $Float, {{%.*}} : $Float, [[BB3_PRED2_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb2__PB__src_0_wrt_0)

// CHECK-SIL: bb4([[BB2_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb2__PB__src_0_wrt_0, {{%.*}} : $Float, {{%.*}} : $Float):
// CHECK-SIL: bb4({{%.*}} : $Float, {{%.*}} : $Float, [[BB2_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb2__PB__src_0_wrt_0):
// CHECK-SIL: [[BB2_PRED:%.*]] = struct_extract [[BB2_PB_STRUCT]]
// CHECK-SIL: copy_addr {{%.*}} to {{%.*}} : $*(Float, Float)
// CHECK-SIL-NOT: copy_addr {{%.*}} to {{%.*}} : $*Float
// CHECK-SIL: [[BB2_PB_STRUCT_DATA:%.*]] = unchecked_enum_data [[BB2_PRED]]
// CHECK-SIL: br bb6([[BB2_PB_STRUCT_DATA]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb5([[BB1_PRED0_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0):
// CHECK-SIL: br bb7([[BB1_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0, {{%.*}} : $Float)
// CHECK-SIL: br bb7({{%.*}} : $Float, [[BB1_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb6([[BB2_PRED0_TRAMP_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0):
// CHECK-SIL: br bb7([[BB2_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0, {{%.*}} : $Float)
// CHECK-SIL: br bb7({{%.*}} : $Float, [[BB2_PRED0_TRAMP_PB_STRUCT]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0)

// CHECK-SIL: bb7([[BB0_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0, {{%.*}} : $Float):
// CHECK-SIL: bb7({{%.*}} : $Float, [[BB0_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb0__PB__src_0_wrt_0):
// CHECK-SIL: return {{%.*}} : $Float
8 changes: 4 additions & 4 deletions test/AutoDiff/leakchecking.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ LeakCheckingTests.test("ControlFlow") {
// FIXME: Fix control flow AD memory leaks.
// See related FIXME comments in adjoint value/buffer propagation in
// lib/SILOptimizer/Mandatory/Differentiation.cpp.
testWithLeakChecking(expectedLeakCount: 105) {
testWithLeakChecking(expectedLeakCount: 74) {
func cond_nestedtuple_var(_ x: Tracked<Float>) -> Tracked<Float> {
// Convoluted function returning `x + x`.
var y = (x + x, x - x)
Expand All @@ -116,7 +116,7 @@ LeakCheckingTests.test("ControlFlow") {
// FIXME: Fix control flow AD memory leaks.
// See related FIXME comments in adjoint value/buffer propagation in
// lib/SILOptimizer/Mandatory/Differentiation.cpp.
testWithLeakChecking(expectedLeakCount: 379) {
testWithLeakChecking(expectedLeakCount: 300) {
func cond_nestedstruct_var(_ x: Tracked<Float>) -> Tracked<Float> {
// Convoluted function returning `x + x`.
var y = FloatPair(x + x, x - x)
Expand All @@ -140,7 +140,7 @@ LeakCheckingTests.test("ControlFlow") {
// FIXME: Fix control flow AD memory leaks.
// See related FIXME comments in adjoint value/buffer propagation in
// lib/SILOptimizer/Mandatory/Differentiation.cpp.
testWithLeakChecking(expectedLeakCount: 9) {
testWithLeakChecking(expectedLeakCount: 3) {
var model = ExampleLeakModel()
let x: Tracked<Float> = 1.0
_ = model.gradient(at: x) { m, x in
Expand All @@ -157,7 +157,7 @@ LeakCheckingTests.test("ControlFlow") {
// FIXME: Fix control flow AD memory leaks.
// See related FIXME comments in adjoint value/buffer propagation in
// lib/SILOptimizer/Mandatory/Differentiation.cpp.
testWithLeakChecking(expectedLeakCount: 14) {
testWithLeakChecking(expectedLeakCount: 6) {
var model = ExampleLeakModel()
let x: Tracked<Float> = 1.0
_ = model.gradient(at: x) { m, x in
Expand Down