[AutoDiff] Support differentiation of loops. #25558

dan-zheng · 2019-06-18T06:55:52Z

Change predecessor enum generation to support loops.
- Generate indirect predecessor enums for BBs in loops.
- Handle boxed payloads of indirect enums.
Traverse basic blocks in post-order post-dominance order.
- This is necessary for computational correctness.
Add loop tests (for-in, while, repeat-while, break/continue, nested).

Examples:

func for_loop(_ x: Float) -> Float {
  var result = x
  for _ in 1..<3 {
    result = result * x
  }
  return result
}
expectEqual((8, 12), valueWithGradient(at: 2, in: for_loop))
expectEqual((27, 27), valueWithGradient(at: 3, in: for_loop))

func nested_loop(_ x: Float, count: Int) -> Float {
  var outer = x
  for _ in 1..<count {
    outer = outer * x

    var inner = outer
    var i = 1
    while i < count {
      inner = inner + x
      i += 1
    }
    outer = inner
  }
  return outer
}
expectEqual((20, 22), valueWithGradient(at: 2, in: { x in nested_loop(x, count: 3) }))
expectEqual((104, 66), valueWithGradient(at: 4, in: { x in nested_loop(x, count: 3) }))

Verified loop differentiation correctness against PyTorch and Tangent.

Exposed TF-584: derivative computation bug for repeat-while loops.
Exposed TF-585: AllocBoxToStack crash with -O for nested loop AD.

- Change predecessor enum generation to support loops. - Generate indirect predecessor enums for BBs in loops. - Handle boxed payloads of indirect enums. - Traverse basic blocks in post-order post-dominance order. - This is necessary for computational correctness. - Add loop tests (`for-in`, `while`, nested).

dan-zheng · 2019-06-18T07:11:50Z

lib/SILOptimizer/Mandatory/Differentiation.cpp

+    if (!boxType)
+      return builder.createEnum(loc, pbStructVal, enumEltDecl, enumLoweredTy);
+    // Otherwise, box the pullback struct value and create an enum.
+    auto *allocBox = builder.createAllocBox(loc, boxType);


This alloc_box is not being deallocated and may cause leaks - to be verified.
Loop AD leak checking tests are added in test/AutoDiff/leakchecking.swift.

Resolved - see this comment. Deallocation is handled automatically during release_value for pullback structs.

dan-zheng · 2019-06-18T07:14:01Z

lib/SILOptimizer/Mandatory/Differentiation.cpp

+  /// Mapping from original basic blocks to local adjoint values to be cleaned
+  /// up. This is populated when adjoint emission is run on one basic block and
+  /// cleaned before processing another basic block.
+  DenseMap<SILBasicBlock *, SmallVector<AdjointValue, 8>>


Note: changed blockLocalAdjointValues to a DenseMap for more principled/predictable behavior.
Previously, using SmallVector caused bugs when the original entry BB isn't the last visited BB during adjoint generation.

rxwei

~~I couldn't find any emission of dealloc_box and this is apparently leaking a lot more memory (all pullback structs in loops are kept alive). This needs to be fixed before merging.~~
Actually, this should be handled automatically when you release_value the pullback struct.

rxwei · 2019-06-18T07:14:56Z

test/AutoDiff/control_flow.swift

+  expectEqual((8, 12), valueWithGradient(at: 2, in: for_loop))
+  expectEqual((27, 27), valueWithGradient(at: 3, in: for_loop))
+
+  func while_loop(_ x: Float) -> Float {


Please also test a repeat-while loop.

Added tests for repeat-while loops and continue/break statements in 32f6629.
Exposed TF-584: derivative computation bug for repeat-while loops. Thanks for pointing this out!

rxwei

The patch is pretty clean. LGTM!

lib/SILOptimizer/Mandatory/Differentiation.cpp

Add tests for repeat-while loops and break/continue statements. Expose TF-584: incorrect derivative computation for repeat-while loops.

dan-zheng · 2019-06-18T22:10:42Z

@swift-ci Please test tensorflow

dan-zheng · 2019-06-18T23:56:15Z

Two tests fail only with optimizations, which I hadn't tested locally:

Failing Tests (2):
    Swift(linux-x86_64) :: AutoDiff/leakchecking.swift
    Swift(linux-x86_64) :: AutoDiff/control_flow.swift

bool swift::StackNesting::solve(): Assertion `!(BI.Block->getTerminator()->isFunctionExiting() && Bits.any()) && "stack location is missing dealloc"' failed.

Perhaps there are some missing deallocations in generated code. Will investigate later.

EDIT: Filed TF-585 with details about the AllocBoxToStack crash.

dan-zheng · 2019-06-19T01:35:29Z

@swift-ci Please test tensorflow

lib/SILOptimizer/Mandatory/Differentiation.cpp

Generate `fix_lifetime` instruction on `alloc_box` to avoid crash.

dan-zheng · 2019-06-19T01:53:01Z

@swift-ci Please test tensorflow

dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Jun 18, 2019

dan-zheng requested a review from rxwei June 18, 2019 06:55

dan-zheng force-pushed the autodiff-loops branch 3 times, most recently from 8ebed4a to eab1067 Compare June 18, 2019 07:06

dan-zheng force-pushed the autodiff-loops branch from eab1067 to 16dc9a2 Compare June 18, 2019 07:10

dan-zheng commented Jun 18, 2019

View reviewed changes

rxwei reviewed Jun 18, 2019

View reviewed changes

rxwei approved these changes Jun 18, 2019

View reviewed changes

lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved

Add tests.

32f6629

Add tests for repeat-while loops and break/continue statements. Expose TF-584: incorrect derivative computation for repeat-while loops.

dan-zheng force-pushed the autodiff-loops branch from a92cdb7 to 32f6629 Compare June 18, 2019 20:38

Minor fixes.

3eb19ba

rxwei reviewed Jun 19, 2019

View reviewed changes

lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated Show resolved Hide resolved

dan-zheng force-pushed the autodiff-loops branch from dfd9e3d to ea63e6f Compare June 19, 2019 01:51

Avoid AllocBoxToStack crash (TF-585).

46caa45

Generate `fix_lifetime` instruction on `alloc_box` to avoid crash.

dan-zheng force-pushed the autodiff-loops branch from ea63e6f to 46caa45 Compare June 19, 2019 01:52

dan-zheng merged commit 766a1de into swiftlang:tensorflow Jun 19, 2019

dan-zheng deleted the autodiff-loops branch June 19, 2019 03:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[AutoDiff] Support differentiation of loops. #25558

[AutoDiff] Support differentiation of loops. #25558

Uh oh!

dan-zheng commented Jun 18, 2019 •

edited

Loading

Uh oh!

dan-zheng Jun 18, 2019

Uh oh!

dan-zheng Jun 18, 2019

Uh oh!

dan-zheng Jun 18, 2019

Uh oh!

rxwei left a comment •

edited

Loading

Uh oh!

rxwei Jun 18, 2019

Uh oh!

dan-zheng Jun 18, 2019

Uh oh!

rxwei left a comment

Uh oh!

Uh oh!

dan-zheng commented Jun 18, 2019

Uh oh!

dan-zheng commented Jun 18, 2019 •

edited

Loading

Uh oh!

dan-zheng commented Jun 19, 2019

Uh oh!

Uh oh!

dan-zheng commented Jun 19, 2019

Uh oh!

Uh oh!

[AutoDiff] Support differentiation of loops. #25558

[AutoDiff] Support differentiation of loops. #25558

Uh oh!

Conversation

dan-zheng commented Jun 18, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dan-zheng Jun 18, 2019

Choose a reason for hiding this comment

Uh oh!

dan-zheng Jun 18, 2019

Choose a reason for hiding this comment

Uh oh!

dan-zheng Jun 18, 2019

Choose a reason for hiding this comment

Uh oh!

rxwei left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rxwei Jun 18, 2019

Choose a reason for hiding this comment

Uh oh!

dan-zheng Jun 18, 2019

Choose a reason for hiding this comment

Uh oh!

rxwei left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dan-zheng commented Jun 18, 2019

Uh oh!

dan-zheng commented Jun 18, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dan-zheng commented Jun 19, 2019

Uh oh!

Uh oh!

dan-zheng commented Jun 19, 2019

Uh oh!

Uh oh!

dan-zheng commented Jun 18, 2019 •

edited

Loading

rxwei left a comment •

edited

Loading

dan-zheng commented Jun 18, 2019 •

edited

Loading