[TORCH] Modified fx_importer to support hop_while_loop #4338
Conversation
keshavvinayak01
changed the title
keshavvinayak01
changed the title
Groverkss
requested review from
Groverkss,
penguin-wwy,
rsuderman,
stellaraccident and
zjgarvey
and removed request for
stellaraccident
|
I'm working on resolving ci issues right now. Once that lands, I'll ping here to have you sync your branch with main and review. |
Sounds good, thanks. |
|
@keshavvinayak01 can you sync with main? |
Signed-off-by: Keshav Vinayak Jha <[email protected]>
keshavvinayak01
force-pushed
the
keshavvinayak01/torch-hop-while
branch
from
3eb338f to
c8c711c
Compare
|
Should be synced. Had to force push. |
zjgarvey
left a comment
zjgarvey
left a comment
There was a problem hiding this comment.
There are some name inconsistencies, and even after resolving these, the generated IR does not lower out of torch dialect (some issue with the conversion to scf.while). I'd double-check the correctness of the generated prim loop, and possibly consider adding an e2e test
Signed-off-by: Keshav Vinayak Jha <[email protected]>
zjgarvey
left a comment
zjgarvey
left a comment
There was a problem hiding this comment.
I added some more comments for now, thanks for addressing the earlier comments.
Please add an e2e test for this op and debug, since I wasn't able to lower the test output IR to linalg/tensor/scf (even with consistent naming).
Signed-off-by: Keshav Vinayak Jha <[email protected]>
Signed-off-by: Keshav Vinayak Jha <[email protected]>
|
I've addressed your comments, but specifically about the generated |
Signed-off-by: Keshav Vinayak Jha <[email protected]>
Signed-off-by: Keshav Vinayak Jha <[email protected]>
Change 1: Converts builtin tensors → Torch tensors when entering the loop body Change 2: Ensures Torch tensors → builtin tensors when yielding back to the loop condition Without these fixes, the conversion would fail when while loops carry tensor values Also modified basic_test.py FILECHECK statements. Signed-off-by: Keshav Vinayak Jha <[email protected]>
|
Had to modify |
That PR added tensor arg support for "for-like" loop conversion, but not "while-like" loop conversion. |
zjgarvey
left a comment
zjgarvey
left a comment
There was a problem hiding this comment.
This looks good to me, but please add at least one e2e test in projects/pt1/python/torch_mlir_e2e_test/test_suite/control_flow.py. Let me know if you want some pointers on adding one of these.
Signed-off-by: Keshav Vinayak Jha <[email protected]>
Signed-off-by: Keshav Vinayak Jha <[email protected]>
Signed-off-by: Keshav Vinayak Jha <[email protected]>
Signed-off-by: Keshav Vinayak Jha <[email protected]>
Signed-off-by: Keshav Vinayak Jha <[email protected]>
zjgarvey
left a comment
zjgarvey
left a comment
There was a problem hiding this comment.
We shouldn't merge this unless the e2e test passes the fx_importer config.
Let me know if you have trouble debugging this.
Signed-off-by: Keshav Vinayak Jha <[email protected]>
zjgarvey
left a comment
zjgarvey
left a comment
There was a problem hiding this comment.
Just two small nits after running through the changes again. Otherwise LGTM!
Signed-off-by: Keshav Vinayak Jha <[email protected]>
2 participants
This PR adds support for emitting graphs for Pytorch HOPs, beginning with
torch._higher_order_ops.while_loop.The proposed change is to modify the
import_programto call function_import_all_child_modules, which recursively imports the stateless graph for all the children modules.Since HOP operator graphs are stateless graphs with no mutation, it is correct to import them as stateless graphs, although the method
import_stateless_graphis marked as "deprecated".