Disable argument preservation optimization#3017
Open
Pangoraw wants to merge 3 commits into
Open
Conversation
The blunt `XLA_BYPASS_OPTIMIZATION` flag broke the clamp! test because `maximum(x_ra)` started donating x_ra's buffer: with all block-arg results forced into nresults, preserved_args_idx was always empty, making every arg donated. XLA then rejected the already-donated clamp output with "Donation requested for invalid buffer". Replace the global bypass with two targeted checks: 1. Preserve a block-arg result only when `linear_results[i] === linear_args[arg_num+1]` (the result is the identical traced object as the arg at that MLIR position). Cross-arg aliasing — e.g. state.u⁰ capturing state.u's block arg — fails this identity check and is returned through XLA explicitly instead. 2. Collect `block_arg_outputs` (arg indices whose block arg appears directly in nresults) and skip donation for those args, preventing XLA from modifying the buffer in-place and corrupting the value being returned as the cross-arg output. Also removes the now-unnecessary XLA_BYPASS_OPTIMIZATION Ref and reverts the kernelabstractions filecheck pattern, which was loosened only because unchanged args were no longer preserved. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
wsmoses
reviewed
Jul 1, 2026
wsmoses
left a comment
wsmoses
left a comment
Member
There was a problem hiding this comment.
can we have this as an optional optimization, rather than full disable rn?
Collaborator
Author
|
hitting test failures with an unconditional deactivation of the optimization due to #3023. |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
we can see if there are potential unwanted side effects of disabling, and remove the optimization altogether if there are none.
But this optimization has caused a few issues in the past