Optimize AOTInductor: Caching, Reduced Decompositions, and Improved JSON Handling #148616
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/148616
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 7890c7b with merge base 4b35139 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "topic: not user facing" |
desertfire
left a comment
There was a problem hiding this comment.
I will let @EikanWang to comment the aoti_eager part.
| @@ -0,0 +1,259 @@ | |||
| torch/_higher_order_ops/utils.py: operator: OperatorBase, delayed_error: bool, *args: Any, **kwargs: Any | |||
There was a problem hiding this comment.
Please revert this third-party change.
| self.module.to(self.src_device_obj) | ||
|
|
||
| def forward(self, *inputs: Any, **kwargs: Any) -> Any: | ||
| def forward(self, *inputs: P.args, **kwargs: P.kwargs) -> R: |
There was a problem hiding this comment.
Is this relevant to this PR?
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
PR Description:
This PR improves AOTInductor's compilation performance by introducing caching, limiting unnecessary decompositions, and optimizing JSON handling.
Changes:
Added persistent caching to avoid redundant recompilation.
Restricted decompositions to only necessary operators (aten::add, aten::mul).
Optimized JSON metadata updates to prevent unnecessary file writes.
Impact:
Reduces compilation time for repeated runs.
Improves efficiency by only updating metadata when needed.
Helps prevent excessive decompositions, leading to better overall performance.
Testing:
Ran pytest test/inductor to check for regressions.
Verified that AOT compilation is significantly faster on repeated runs.
cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov