Aoti minifier flatten #141156
Aoti minifier flatten #141156
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/141156
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 1751740 with merge base 16ea0dd ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
yushangdi
force-pushed
the
aoti_minifier_flatten_2
branch
2 times, most recently
from
d2a7339 to
aaebc35
Compare
yushangdi
force-pushed
the
aoti_minifier_flatten_2
branch
2 times, most recently
from
137b8c1 to
6a4d3be
Compare
torch/_inductor/compile_fx.py
Outdated
| warn_and_skip(node.get_device()) | ||
|
|
||
|
|
||
| def _flatten_inputs( |
There was a problem hiding this comment.
I understand this is moved from _inductor/init.py, but I feel flattening inputs should be a functionality provided by Export. @angelayi , does this already exist in Export?
There was a problem hiding this comment.
@desertfire (Angela pls correct me if I'm wrong) My understanding is that export does flatten the inputs, but that's the inputs for the "flattened graph" (exported_program.graph), but this graph might be different from the input graph module to AOTI (exported_program.module()) .
The main difference is that, exported_program.module() may contain submodules, but exported_program.graph is fully flatten.
So the two flattening are a little different. Here in _flatten_inputs, we are only flattening the inputs, but not other parts of the graph (e.g. the submodules).
There was a problem hiding this comment.
Yeah this code is pretty similar to what we have here, but I'm not sure if it's easy to refactor things to reuse this.
There was a problem hiding this comment.
@desertfire The _flatten_inputs function here does 3 things:
- update the
gm's codegen to take in the flattened inputs - get the flattened inputs (this is similar to the code that Angela linked above, but that's basically 2 lines)
- update the inductor_config's
serialized_in_specandserialized_out_specfields
Maybe I could rename this function to e.g. _update_gm_input_spec_and_flatten to be more clear that this function does more than just getting the flattened inputs?
There was a problem hiding this comment.
Or you rename to something more generic but with aoti in it.
There was a problem hiding this comment.
Or you rename to something more generic but with
aotiin it.
sure, renamed to _aoti_flatten_inputs now.
|
|
||
| options = ( | ||
| { | ||
| "aot_inductor.serialized_in_spec": serialized_in_spec, |
There was a problem hiding this comment.
This is a legacy issue, but we really shouldn't use inductor config for passing these. It will be nice to have some refactoring someday.
torch/_inductor/compile_fx.py
Outdated
| warn_and_skip(node.get_device()) | ||
|
|
||
|
|
||
| def _flatten_inputs( |
There was a problem hiding this comment.
Or you rename to something more generic but with aoti in it.
yushangdi
force-pushed
the
aoti_minifier_flatten_2
branch
4 times, most recently
from
2d9c61a to
b2f1579
Compare
yushangdi
force-pushed
the
aoti_minifier_flatten_2
branch
from
b2f1579 to
1751740
Compare
|
@pytorchbot merge -i |
Merge startedYour change will be merged while ignoring the following 1 checks: pull / linux-focal-py3_9-clang9-xla / test (xla, 1, 1, lf.linux.12xlarge) Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Flatten the inputs to minifier so AOTI Minifier can handle unflattened inputs and kwargs. - flatten the inputs in minifier - changed the "load_and_run" part of the minifier verification to run on the flattened inputs. - refactored code to keep `torch._inductor.__init__.py` clean - update doc `python test/inductor/test_minifier.py` Pull Request resolved: pytorch#141156 Approved by: https://github.com/desertfire
4 participants
Flatten the inputs to minifier so AOTI Minifier can handle unflattened inputs and kwargs.
torch._inductor.__init__.pycleanpython test/inductor/test_minifier.pycc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov