Skip to content

Conversation

@dan-zheng
Copy link
Contributor

-enable-strip-ownership-after-serialization is now true by default,
skipping ownership model eliminator after differentiation.

This is problematic for control flow differentiation. VJPEmitter creates
trampoline bbs, which create predecessor enum values and immediately branch to
a destination bb. Trampoline bbs share the argument of their destination bb,
which may have @guaranteed ownership kind. @guaranteed values need a single
lifetime-ending use, e.g. an end_borrow.

Previously, VJPEmitter::trampolinedGuaranteedPhiArguments tracked all
@guaranteed trampoline/destination bb argument pairs and emitted an
end_borrow for the trampoline bb argument after the cloned end_borrow for
the destination bb argument. However, this does not work when ownership model
eliminator is skipped: mandatory inlining invokes mergeBasicBlocks on VJPs,
which causes the end_borrow instructions to be redundant.

// Generated by differentiation:
bb1(%6 : @guaranteed $Tracked<Float>):            // Preds: bb0
  %7 = enum $_AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF_bb1__Pred__src_0_wrt_1, #_AD
__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF_bb1__Pred__src_0_wrt_1.bb0!enumelt.1, %4 : $
_AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF_bb0__PB__src_0_wrt_1 // user: %8
  br bb2(%6 : $Tracked<Float>, %7 : $_AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF_bb1_
_Pred__src_0_wrt_1) // id: %8
// %9                                             // user: %11
// %10                                            // user: %13
bb2(%9 : @guaranteed $Tracked<Float>, %10 : $_AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_A
FtF_bb1__Pred__src_0_wrt_1): // Preds: bb1
  end_borrow %9 : $Tracked<Float>                 // id: %11
  end_borrow %6 : $Tracked<Float>                 // id: %12
  ...
// After mandatory inlining, which calls `mergeBasicBlocks`:
bb1(%6 : @guaranteed $Tracked<Float>):            // Preds: bb0
  %7 = enum $_AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF_bb1__Pred__src_0_wrt_1, #_AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF_bb1__Pred__src_0_wrt_1.bb0!enumelt.1, %4 : $_AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF_bb0__PB__src_0_wrt_1 // user: %10
  end_borrow %6 : $Tracked<Float>                 // id: %8
  end_borrow %6 : $Tracked<Float>                 // id: %9
  ...
Function: 'AD__$s1e15enum_notactive2y23DifferentiationUnittest7TrackedVySfGAA4EnumO_AFtF__vjp_src_0_wrt_1'
Found over consume?!
Value: %13 = argument of bb2 : $(Tracked<Float>, Tracked<Float>) // users: %15, %17, %18
User:   end_borrow %13 : $(Tracked<Float>, Tracked<Float>) // id: %17
Block: bb2

Now, VJPEmitter preemptively calls mergeBasicBlocks, which merges trampoline
bbs with their destinations. @guaranteed bb arguments are merged and no longer
duplicated, so there is no need to specially emit end_borrow for them.

Consequences:

  • VJP functions now have simplified CFGs.
  • All VJP trampoline bbs are merged with their destinations. Consider rewriting
    the transform to generate the merged destination bbs in the first place.

Resolves TF-990.

`-enable-strip-ownership-after-serialization` is now true by default,
skipping ownership model eliminator after differentiation.

This is problematic for control flow differentiation. `VJPEmitter` creates
trampoline bbs, which create predecessor enum values and immediately branch to
a destination bb. Trampoline bbs share the argument of their destination bb,
which may have `@guaranteed` ownership kind. `@guaranteed` values need a single
lifetime-ending use, e.g. an `end_borrow`.

Previously, `VJPEmitter::trampolinedGuaranteedPhiArguments` tracked all
`@guaranteed` trampoline/destination bb argument pairs and emitted an
`end_borrow` for the trampoline bb argument after the cloned `end_borrow` for
the destination bb argument. However, this does not work when ownership model
eliminator is skipped: mandatory inlining invokes `mergeBasicBlocks` on VJPs,
which causes the `end_borrow` instructions to be redundant.

Now, `VJPEmitter` preemptively calls `mergeBasicBlocks`, which merges trampoline
bbs with their destinations. `@guaranteed` bb arguments are merged and no longer
duplicated, so there is no need to specially emit `end_borrow` for them.

Consequences:
- VJP functions now have simplified CFGs.
- All VJP trampoline bbs are merged with their destinations. Consider rewriting
  the transform to generate the merged destination bbs in the first place.

Resolves TF-990.
@dan-zheng dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Nov 23, 2019
@dan-zheng dan-zheng requested review from marcrasi and rxwei November 23, 2019 22:53
@dan-zheng
Copy link
Contributor Author

@swift-ci Please clean test tensorflow

@dan-zheng
Copy link
Contributor Author

I created a PR to document this VJP generation change.
Happy to discuss in detail later, merging now to unblock progress.

@dan-zheng dan-zheng merged commit f9956c1 into swiftlang:tensorflow-merge Nov 23, 2019
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

tensorflow This is for "tensorflow" branch PRs.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant