Modest code motion in compile_fx #141574
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. |
github-actions
bot
requested review from
SherlockNoMad,
albanD,
antoniojkim,
bdhirsh and
miladm
| 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.
New docblock (and a few comments sprinkled here and there later)
| model_, | ||
| example_inputs_, | ||
| recursive_compile_fx, | ||
| ) |
There was a problem hiding this comment.
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.
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.
I inlined _fw_compiler_base into fw_compiler_base at the cost of one level of indentation (see autoformat fix below)
|
|
||
| # 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.
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
4 participants
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