Skip to content

Add const shadows for forward mode AD on RegionBranchOpInterface#2780

Open
vimarsh6739 wants to merge 1 commit intomainfrom
scf.for-fix2
Open

Add const shadows for forward mode AD on RegionBranchOpInterface#2780
vimarsh6739 wants to merge 1 commit intomainfrom
scf.for-fix2

Conversation

@vimarsh6739
Copy link
Copy Markdown
Member

@vimarsh6739 vimarsh6739 commented Apr 8, 2026

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).

Comment thread enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp Outdated
@vimarsh6739 vimarsh6739 force-pushed the scf.for-fix2 branch 2 times, most recently from 2fcb520 to d31bb21 Compare May 6, 2026 20:21
Comment thread enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp Outdated
// 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) {
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I don't think this should be here, but within the impl of createWithShadows, right?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

No, it has to be here. We are shadowing only the const operands for any augmented newOp.

Copy link
Copy Markdown
Member Author

@vimarsh6739 vimarsh6739 May 7, 2026

Choose a reason for hiding this comment

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

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.

@vimarsh6739 vimarsh6739 force-pushed the scf.for-fix2 branch 6 times, most recently from 736e317 to a0a2d7b Compare May 7, 2026 01:37
@vimarsh6739 vimarsh6739 changed the title Fix forward mode AD for for-like ops Add const shadows for forward mode AD on RegionBranchOpInterface May 7, 2026
@vimarsh6739 vimarsh6739 requested a review from wsmoses May 7, 2026 01:49
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants