Add const shadows for forward mode AD on RegionBranchOpInterface#2780
Open
vimarsh6739 wants to merge 1 commit intomainfrom
Open
Add const shadows for forward mode AD on RegionBranchOpInterface#2780vimarsh6739 wants to merge 1 commit intomainfrom
RegionBranchOpInterface#2780vimarsh6739 wants to merge 1 commit intomainfrom
Conversation
51333d0 to
450fff0
Compare
wsmoses
reviewed
May 5, 2026
2fcb520 to
d31bb21
Compare
wsmoses
reviewed
May 7, 2026
wsmoses
reviewed
May 7, 2026
| // forceAugmentedReturns will not shadow const operands. No need to add | ||
| // to the invertPointers map since `operand` is const, the shadow will | ||
| // be unused. | ||
| for (const RegionSuccessor &successor : entrySuccessors) { |
Member
There was a problem hiding this comment.
I don't think this should be here, but within the impl of createWithShadows, right?
Member
Author
There was a problem hiding this comment.
No, it has to be here. We are shadowing only the const operands for any augmented newOp.
Member
Author
There was a problem hiding this comment.
createWithShadows correctly adds const shadows. But newOp doesn't have them(since it was created with forceAugmentedReturns), so the takeBody() discards the args in the replacement op
We have 2 choices here - either add the const shadows to newOp or the replacement (after takeBody). The end result is the same.
736e317 to
a0a2d7b
Compare
RegionBranchOpInterface
When an operand has const activity, we may still need to create a shadow for it (particularly in scf.for, where a const iter_arg may still have to be shadowed if the result and the terminator are active)
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
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.
This primarily affects scf.for and affine.for, but should also handle a the same case in scf.parallel and affine.parallel ops (iter-arg is a constant/dead inside loop but terminators still have activity).
Aside: I think one interesting thing we can do here is to reuse the successor information to effectively eliminate the creation of some dead regions. This already kind of happens in scf.for which prunes the entry successors depending on if the loop has no iterations, or non-zero trip count(e.g. if the only entry successor is the parent, then just skip diffing the regionBranchOp body).