CARVIEW |
Navigation Menu
-
Notifications
You must be signed in to change notification settings - Fork 24.7k
Modest code motion in compile_fx #141574
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
Modest code motion in compile_fx #141574
Conversation
đź”— Helpful Linksđź§Ş See artifacts and rendered test results at hud.pytorch.org/pr/141574
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit df08081 with merge base 4c1f50a ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
||
NB: This function TAKES OWNERSHIP of the input ``model_`` and can potentially | ||
mutate it! Make a copy if you need to preserve the original GraphModule. | ||
""" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
New docblock (and a few comments sprinkled here and there later)
model_, | ||
example_inputs_, | ||
recursive_compile_fx, | ||
) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
All of these early short circuits (above, not shown due to whitespace changed ignored) have been hoisted out of the top level context managers here. The reasoning is that they immediately call back to compile_fx
, and so therefore we don't have to set the context manager multiple times.
|
||
if isinstance(model_, GraphModule) and isinstance( | ||
model_.graph._codegen, _PyTreeCodeGen | ||
): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This short circuit was nestled with recursive pre-grad passes conditional, I have hoisted it out and duplicated the condition
model: GraphModule, | ||
example_inputs: List[InputType], | ||
is_inference: bool, | ||
) -> CompiledFxGraph: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I inlined _fw_compiler_base
into fw_compiler_base
at the cost of one level of indentation (see autoformat fix below)
@@ -1688,10 +1722,6 @@ def bw_compiler( | |||
boxed_forward_device_index=forward_device, | |||
) | |||
|
|||
# TODO: can add logging before/after the call to create_aot_dispatcher_function | |||
# in torch._functorch/aot_autograd.py::aot_module_simplified::aot_function_simplified::new_func | |||
# once torchdynamo is merged into pytorch |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
torchdynamo is merged into pytorch so I assume this is obsolete
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 jobs have failed, first few of them are: inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu) Details for Dev Infra teamRaised by workflow job |
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #141654 Approved by: https://github.com/aorenste, https://github.com/jansel ghstack dependencies: #141491, #141492, #141574
Do code review with whitespace changes off. Check comments for what I changed. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch#141574 Approved by: https://github.com/bobrenjc93, https://github.com/jansel ghstack dependencies: pytorch#141491, pytorch#141492
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch#141654 Approved by: https://github.com/aorenste, https://github.com/jansel ghstack dependencies: pytorch#141491, pytorch#141492, pytorch#141574
Stack from ghstack (oldest at bottom):
Do code review with whitespace changes off. Check comments for what I changed.
Signed-off-by: Edward Z. Yang ezyang@meta.com
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov