Skip to content
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

Jump to bottom

[AutoDiff] Support differentiation of conditionals. #25057

Merged
merged 22 commits into from
Jun 4, 2019
Merged

[AutoDiff] Support differentiation of conditionals. #25057

merged 22 commits into from
Jun 4, 2019

Conversation

dan-zheng
Copy link
Contributor

@dan-zheng dan-zheng commented May 25, 2019 •
edited
Loading

  • Support control flow in adjoint generation.
    • Make adjoint value/buffer mappings be per basic block.
    • Change AdjointValue to not be move-only. Original values from
      different basic blocks may share the same AdjointValue.
    • Propagate adjoint values from active bb arguments to predecessor
      terminator operands.
    • Propagate adjoint values/buffers of dominated active values/buffers
      to predecessor blocks.
      • For active values: propagate adjoint values as adjoint bb arguments.
      • For active buffers: propagate adjoint buffers via copy_addr.
  • Revamp AdjointEmitter handling of begin_access and end_access.
    • getAdjointValue of begin_access now returns the adjoint base buffer.
      Previously, it returned a begin_access of the adjoint base buffer
      without generating a corresponding end_access.
    • AdjointEmitter::visitBeginAccessInst now generates no code.
    • AdjointEmitter::visitEndAccessInst now does nothing.
  • Add various control flow differentiation tests.
    • Test differentiation of conditionals (nested), recursion,
      var allocations (tuples, structs).
    • Add negative leakchecking tests.

Todos:

  • Fix adjoint value/buffer propagation memory leaks.
  • Add more tests (adjoint SIL, leakchecking).
  • Support differentiation of enum-related instructions and loops.

@dan-zheng dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label May 25, 2019
@dan-zheng dan-zheng force-pushed the autodiff-control-flow branch from 8dd1735 to 04b5e15 Compare May 25, 2019 00:59
@dan-zheng
Copy link
Contributor Author

dan-zheng commented May 25, 2019 •
edited
Loading

This is a work in-progress; there are at least two bugs.

  • Generic adjoint functions with control flow fail verification due to "use of un-dominated operand". Reproducer and log here.
    • I think the crash is that adjoint buffers are created in the wrong buffer. For example, in the Gist, %40 = alloc_stack $τ_0_0.TangentVector should be allocated in adjoint bb0, not bb3.
  • Gradient values are sometimes incorrectly 0.
    • grep FIXME: Debug incorrect "0" gradient values wrt x in tests.

@dan-zheng dan-zheng requested a review from rxwei May 25, 2019 01:00
@dan-zheng
Support differentiation of conditionals.
fe1fbd8 
- Support control flow in adjoint generation.
  - Make adjoint value mapping be per basic block.
  - Change `AdjointValue` to not be move-only. Original values from
    different basic blocks can share the same `AdjointValue`.
  - Propagate adjoint values from active basic block arguments to
    predecessor terminator operands.
  - Propagate adjoint values of active values to adjoint successor blocks.
- Add runtime control flow differentiation tests.
@dan-zheng dan-zheng force-pushed the autodiff-control-flow branch from 04b5e15 to fe1fbd8 Compare May 25, 2019 01:01
dan-zheng added 2 commits May 24, 2019 18:25
@dan-zheng
Fix type remapping crash, expose "use of un-dominated operand" crash.
020f24e 
Expose new crash in adjoint verification:
```
SIL verification failed: instruction isn't dominated by its operand:
properlyDominates(valueI, I)
```
@dan-zheng
Rename "forwarding basic block" to "trampoline basic block".
e52d9fb
rxwei
rxwei reviewed May 26, 2019
View reviewed changes
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
@rxwei
Copy link
Contributor

rxwei commented May 28, 2019

FYI: Since @dan-zheng is OOO, I'm taking over this today to make progress.

rxwei and others added 6 commits May 30, 2019 03:25
@rxwei
Address comments.
5160703
@rxwei
Merge branch 'tensorflow' of github.com:apple/swift into autodiff-con…
e2cc993 
…trol-flow
@rxwei
Merge branch 'tensorflow' of github.com:apple/swift into autodiff-con…
1f77d30 
…trol-flow
@rxwei
Merge branch 'tensorflow' of github.com:apple/swift into autodiff-con…
bdeb342 
…trol-flow
@rxwei
Expose incorrect active buffer handling.
00f1776 
Also fix `getRemappedTangentType` to handle addresses.
@dan-zheng
Fix active adjoint buffer handling.
feaa1b8 
- Make adjoint buffer mapping per basic block.
- For active buffers in original blocks, allocate and zero initialize adjoint
  buffers in adjoint entry.
- Propagate adjoint buffers via `copy_addr` before branching in adjoint blocks.
- Re-enable fixed tests.
@dan-zheng
Change adjoint function local allocation insertion order.
91c36fe 
Insert local allocations before previous local allocations to ensure
that allocation and zero initialization instructions are grouped together.

Fixes test/AutoDiff/generics.swift.
rxwei
rxwei reviewed Jun 1, 2019
View reviewed changes
lib/SILOptimizer/Mandatory/Differentiation.cpp
@@ -4216,11 +4216,15 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
}

SILBasicBlock::iterator getNextFunctionLocalAllocationInsertionPoint() {
Copy link
Contributor

Choose a reason for hiding this comment

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

Might be better to turn this into a function that allocates a local buffer directly: ‘createFunctionLocalBuffer(SILType)’.

dan-zheng added a commit to dan-zheng/swift that referenced this pull request Jun 1, 2019
@dan-zheng
Debug adjoint value propagation.
9e47823 
See: swiftlang#25057 (comment)
@dan-zheng
Fix/simplify adjoint value propagation logic.
bd8d319 
- Record *dominated* active values per basic block, not just
  active values used in basic blocks.
- Remove `visitAllocStack` and `visitDeallocStack` from `AdjointEmitter`.
  - These fall out from active buffer handling.
- Add nested conditional and recursion tests.
- Remove "-differentiable-enable-control-flow" flag, now that
  limited control flow support seems robust.
@dan-zheng dan-zheng force-pushed the autodiff-control-flow branch 3 times, most recently from 5041a69 to e74e43e Compare June 3, 2019 22:45
@dan-zheng dan-zheng changed the title [WIP] [AutoDiff] Support differentiation of conditionals. [AutoDiff] Support differentiation of conditionals. Jun 3, 2019
@dan-zheng dan-zheng force-pushed the autodiff-control-flow branch from e74e43e to 399b797 Compare June 3, 2019 22:57
@dan-zheng
Copy link
Contributor Author

@swift-ci Please clean test tensorflow

@dan-zheng dan-zheng force-pushed the autodiff-control-flow branch from 5b54584 to 0aaff91 Compare June 4, 2019 03:46
@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

rxwei
rxwei reviewed Jun 4, 2019
View reviewed changes
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
rxwei
rxwei reviewed Jun 4, 2019
View reviewed changes
Copy link
Contributor

@rxwei rxwei left a comment

Choose a reason for hiding this comment

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

I couldn't find any tests that are using vars. I think mutation should be tested in order for this PR to be ready.

lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved
test/AutoDiff/control_flow.swift Show resolved Hide resolved
@dan-zheng
Add control flow var tests.
806b464 
`var` tests currently crashing due to `AdjointEmitter` handling of
`begin_access`/`end_access`.
@dan-zheng dan-zheng force-pushed the autodiff-control-flow branch from 0dc5580 to 806b464 Compare June 4, 2019 08:23
dan-zheng added 2 commits June 4, 2019 09:17
@dan-zheng
Fix var + control flow differentiation.
49937d9 
- Change `AdjointEmitter` handling of `begin_access` and `end_access`.
  - Change `AdjointEmitter::getAdjointValue` of `begin_access` instruction
    to return adjoint buffer of the base buffer. Previously, it returned a
    `begin_access` of the adjoint base buffer, without generating a
    corresponding `end_access` instruction, leading to a hole.
  - Change `AdjointEmitter::visitBeginAccessInst` to generate no code,
    only performing verification and propagating cleanups.
  - Change `AdjointEmitter::visitEndAccessInst` to do nothing.
- Fix and add `var` tests.
- Todo: fix control-flow related memory leaks (FIXMEs) related to adjoint
  value/buffer propagation.
@dan-zheng
Address review feedback.
bcd2ce5
@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

@dan-zheng
Copy link
Contributor Author

To step back, there are two problems:

  1. How do you guarantee activeValue has never been an argument to getAdjointBuffer(predBB, activeValue) before? That is, will there be a case where getAdjointBuffer(predBB, activeValue) is called on the same activeValue twice? If that does happen, this code will simply rewrite the the existing adjoint value, which may give you a miscalculation error. Otherwise, it's worth thinking about how the code can be structured in a way that the reader understands this is the first time such a buffer is created for this particular adjoint value.
  2. Back to cleanups, the way you should deal with cleanups is directly related to (1). If there are multiple calls to getAdjointBuffer(predBB, activeValue) for the same activeValue, then you should probably do an AdditiveArithmetic.+=(&succAdjBuf, adjBuf) and transfer no cleanup. Otherwise, do succAdjBuf.setCleanup(adjBuf.getCleanup()) and set adjBuf's cleanup to nullptr.

In any case, I think (1) should be resolved first.

We peer programmed and discussed the two problems in depth.

We found that getAdjointBuffer(predBB, activeValue) is called multiple times for the same activeValue. But performing AdditiveArithmetic.+=(&predAdjBuf, adjBuf) instead of copy_addr(adjBuf, predAdjBuf) lead to computation errors.

Progress in 49937d9: var + control flow tests are fixed, but cleanups during adjoint value/buffer propagation are not handled, leading to memory leaks.

Negative leak checking tests added to test/AutoDiff/leakchecking.swift. @rxwei offered to help fix the cleanup logic in a follow-up PR.

rxwei
rxwei approved these changes Jun 4, 2019
View reviewed changes
Copy link
Contributor

@rxwei rxwei left a comment

Choose a reason for hiding this comment

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

LGTM for merging as-is after existing issues are addressed. I’ll work on a fix for memory leaks.

@dan-zheng
Fix tests.
9af6170 
- Make sure `gradient` calls test the right function.
- Add struct `let` and `var` tests.

Note: tuple/struct tests are designed to compute `x + x`
(in convoluted ways) for easy verification.

Corresponding tuple/struct tests have essentially the same body
(e.g. `cond_nestedtuple_var` ~= `cond_nestedstruct_var`).
dan-zheng
dan-zheng commented Jun 4, 2019
View reviewed changes
test/AutoDiff/leakchecking.swift Outdated
}

LeakCheckingTests.test("ControlFlow") {
testWithLeakChecking(expectedLeakCount: 0) {
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 cond_tuple_var test is not effective because it operates on Float instead of Tracked<Float>.
Fixing it requires significant changes to the Tracked type. I'll fix and re-add it (and more tests) in a follow-up.

The control flow leakchecking tests below should be sufficient for testing memory leaks for now.

@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

@dan-zheng dan-zheng force-pushed the autodiff-control-flow branch from a467a11 to cf8796e Compare June 4, 2019 20:33
@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

@rxwei
Copy link
Contributor

rxwei commented Jun 4, 2019

@swift-ci please test tensorflow

@rxwei rxwei merged commit 045e192 into swiftlang:tensorflow Jun 4, 2019
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@rxwei rxwei rxwei approved these changes

Assignees
No one assigned
Labels
tensorflow This is for "tensorflow" branch PRs.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@dan-zheng @rxwei