[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. |
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
| for child_name, child_module in prog.graph.owning_module.named_children(): | ||
| if isinstance(child_module, GraphModule) and hasattr(child_module, 'graph'): | ||
| # Generate function name: parent_childname | ||
| child_func_name = f"{parent_name}_{child_name}_{id(child_module)}" |
There was a problem hiding this comment.
This is not consistent with the _hop_while_loop callee names (hence the ci failure). Since you can't easily pass the parent name to this func, it would make sense to simply put the child name + an additional uniqueifier (I'm not a fan of id, since it will not be reproducible between runs).
You might be better off defining a mapping between graph modules and mlir func names as an attribute of the FxImporter, and handling name collisions as necessary there.
Signed-off-by: Keshav Vinayak Jha <[email protected]>
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. |
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".