Skip to content
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.

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

Fix markNonContextParamsAsSideEffectFree. #6054

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

csyonghe
Copy link
Collaborator

markNonContextParamsAsSideEffectFree is supposed to mark all parameters except the intermediate context parameter as IgnoreSideEffect. However, the function is not implemented correctly. It checks if a parameter is an intermediate context by checking if its type is IRIntermediateContextType. However at the time this function is called, the intermediate context type is already lowered into a struct, so the current implementation will always mark all parameters as IngoreSideEffect. which result in the compiler to aggressively remove/ignore many primal function calls during autodiff pass when it shouldn't.

This PR fixes this bug by explicitly marking the intermediate context parameter with a PrimalContext decoration, and instead of checking parameter types in markNonContextParamsAsSideEffectFree, we simply look for this decoration to test if a parameter is an intermediate context.

Closes #6039.
Closes #6041.
Closes #6048.

@csyonghe csyonghe requested a review from a team as a code owner January 10, 2025 02:55
@csyonghe csyonghe added the pr: non-breaking PRs without breaking changes label Jan 10, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pr: non-breaking PRs without breaking changes
Projects
None yet
1 participant