[AutoDiff] Fix use after free when pullback is used multiple times #64647
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@swift-ci please test |
|
Tagging @BradLarson @fibrechannelscsi |
rxwei
approved these changes
Mar 27, 2023
There was a problem hiding this comment.
Looks good to me.
Aside from this change, I think that we should also try to migrate from the underlying bump pointer allocator to StackAllocator that's currently used for Concurrency. StackAllocator allows for the pullback to deallocate things along the way.
Oh, great, I was not aware of it. Definitely will explore. |
etcwilde
pushed a commit
to etcwilde/swift
that referenced
this pull request
Apr 19, 2023
…4647) Linear maps are captured in vjp routine via callee-guaranteed partial apply and are passed as @owned references to the enclosing pullback that finally consumes them. Necessary retains are inserted by a partial apply forwarder. However, this is not the case when the function being differentiated contains loops as heap-allocated context is used and bare pointer is captured by the pullback partial apply. As a result, partial apply forwarder does not retain the linear maps that are owned by a heap-allocated context, however, they are still treated as @owned references and therefore are released in the pullback after the first call. As a result, subsequent pullback calls release linear maps and we'd end with possible use-after-free. Ensure we retain values when we load values from the context. Reproducible only when: * Loops (so, heap-allocated context) * Pullbacks of thick functions (so context is non-zero) * Multiple pullback calls * Some cleanup while there Fixes swiftlang#64257
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Linear maps are captured in vjp routine via callee-guaranteed partial apply and are passed as
@ownedreferences to the enclosing pullback that finally consumes them. Necessary retains are inserted by a partial apply forwarder.However, this is not the case when the function being differentiated contains loops as heap-allocated context is used and bare pointer is captured by the pullback partial apply. As a result, partial apply forwarder does not retain the linear maps that are owned by a heap-allocated context, however, they are still treated as
@ownedreferences and therefore are released in the pullback after the first call. As a result, subsequent pullback calls release linear maps and we'd end with possible use-after-free.Ensure we retain values when we load values from the context.
Reproducible only when:
Some cleanup while there
Fixes #64257