Skip to content

[AutoDiff] Fix AllocStack SILDebugVariable assertion failure. #27486

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
Merged
merged 2 commits into from
Oct 2, 2019
Merged

[AutoDiff] Fix AllocStack SILDebugVariable assertion failure. #27486

merged 2 commits into from
Oct 2, 2019

Conversation

dan-zheng
Copy link
Contributor

@dan-zheng dan-zheng commented Oct 2, 2019 •
edited
Loading

#27238 tightened debug info requirements, adding assertions to
AllocStack and AllocBox constructors.

This requires differentiation transform changes:

  • When cloning alloc_stack instructions that may have SILDebugVariable info,
    propagate the info.
    • This is relevant for JVPEmitter::emitTangentForAllocStackInst.
  • When using AllocStack for "local allocations" that do not correspond to
    user-written VarDecls, use RegularLocation::getAutoGeneratedLocation() as
    the location so that assertions do not trigger.
    • This is relevant for PullbackEmitter::getAdjointBuffer.

Example assertion failure prior to this patch:

Assertion failed: ((!dyn_cast_or_null<VarDecl>(Loc.getAsASTNode<Decl>()) || Var) &&
"location is a VarDecl, but SILDebugVariable is empty"), function createAllocStack,
file /Users/danielzheng/swift-merge/swift/include/swift/SIL/SILBuilder.h, line 407.

Stack dump:
...
1.      Swift version 5.1-dev (LLVM c5340df2d1, Swift 7bc4a06c32)
2.      While running pass #56204 SILModuleTransform "Differentiation".

swift::SILBuilder::createAllocStack(...)
(anonymous namespace)::PullbackEmitter::getAdjointBuffer(...)

@dan-zheng dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Oct 2, 2019
@dan-zheng dan-zheng requested review from burmako and rxwei October 2, 2019 07:14
dan-zheng
dan-zheng commented Oct 2, 2019
View reviewed changes
lib/SILOptimizer/Mandatory/Differentiation.cpp Outdated
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Note: this code is constructing a temporary buffer inside PullbackEmitter::getAdjointBuffer - maybe it's even correct for this buffer to have a synthesized location (because it's an auxiliary variable that doesn't directly relate to user-written variables).

@dan-zheng
Copy link
Contributor Author

dan-zheng commented Oct 2, 2019 •
edited
Loading

This patch isn't yet robust, triggering tests to see if TensorFlow compilation passes. (I found locally that it did.)

@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

@dan-zheng
[WIP] [AutoDiff] Fix AllocStack SILDebugVariable assertion failure.
8cee057 
swiftlang#27238 tightened debug info requirements,
adding assertions to AllocStack and AllocBox constructors.

This caused assertion failures for the differentiation transform, which does
not provide SILDebugVariable when building AllocStack during differential/pullback
generation:

```
Assertion failed: ((!dyn_cast_or_null<VarDecl>(Loc.getAsASTNode<Decl>()) || Var) &&
"location is a VarDecl, but SILDebugVariable is empty"), function createAllocStack,
file /Users/danielzheng/swift-merge/swift/include/swift/SIL/SILBuilder.h, line 407.

Stack dump:
...
1.      Swift version 5.1-dev (LLVM c5340df2d1, Swift 7bc4a06)
2.      While running pass swiftlang#56204 SILModuleTransform "Differentiation".

swift::SILBuilder::createAllocStack(...)
(anonymous namespace)::PullbackEmitter::getAdjointBuffer(...)
```

There are a few potential fixes:
- Use `RegularLocation::getAutoGeneratedLocation()` so that the assertion
  does not trigger. This needs to be done for all `createAllocStack`
  invocations in Differentiation.cpp.
- Find a way to provide the correct SILDebugVariable. This seems hard because
  SILDebugVariable is constructed during SILGen and has a particular "ArgNo"
  field that seems hard to compute.

This is an initial minimal test to see if tensorflow/swift-apis compilation
will pass. More changes are needed for a robust final fix:
- Update all `createAllocStack` calls, including those in differential generation.
- Add apple/swift tests.
@rxwei
Copy link
Contributor

rxwei commented Oct 2, 2019

Use RegularLocation::getAutoGeneratedLocation() so that the assertion
does not trigger. This needs to be done for all createAllocStack
invocations in Differentiation.cpp.

Doing this for all createAllocStack innovations isn't quite right. It's only necessary for what we call "local allocations", not for allocations that do correspond to VarDecls, e.g. the tangent of an alloc_stack in a differential.

@dan-zheng
Fix differential generation for AllocStack.
2c6f07c 
Use the VarInfo of the original AllocStack instruction.
@dan-zheng dan-zheng force-pushed the diff-transform-fix-alloc-stack-nodebugvar branch from 9e2a28d to 2c6f07c Compare October 2, 2019 17:10
rxwei
rxwei approved these changes Oct 2, 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.

This fix LGTM.

@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

1 similar comment
@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

burmako
burmako approved these changes Oct 2, 2019
View reviewed changes
Copy link

@burmako burmako left a comment

Choose a reason for hiding this comment

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

Thank you guys for explaining how this should be fixed!

@dan-zheng dan-zheng merged commit 7569ff1 into swiftlang:tensorflow-merge Oct 2, 2019
@dan-zheng dan-zheng deleted the diff-transform-fix-alloc-stack-nodebugvar branch October 2, 2019 21:42
@dan-zheng dan-zheng changed the title [WIP] [AutoDiff] Fix AllocStack SILDebugVariable assertion failure. [AutoDiff] Fix AllocStack SILDebugVariable assertion failure. Oct 2, 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

+1 more reviewer

@burmako burmako burmako approved these changes

Reviewers whose approvals may not affect merge requirements

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.

3 participants

@dan-zheng @rxwei @burmako