Skip to content

[TORCH] Modified fx_importer to support hop_while_loop #4338

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

[TORCH] Modified fx_importer to support hop_while_loop #4338

wants to merge 7 commits into from
+314 −14

Conversation

@keshavvinayak01
Copy link

@keshavvinayak01 keshavvinayak01 commented Oct 10, 2025 •
edited
Loading

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_program to 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_graph is marked as "deprecated".

@keshavvinayak01 keshavvinayak01 changed the title [WIP] Modified fx_imoprter to support hop_while_loop [WIP] Modified fx_importer to support hop_while_loop Oct 10, 2025
@keshavvinayak01 keshavvinayak01 marked this pull request as ready for review October 13, 2025 09:33
@keshavvinayak01 keshavvinayak01 marked this pull request as draft October 13, 2025 19:36
@keshavvinayak01 keshavvinayak01 marked this pull request as ready for review October 15, 2025 15:14
@keshavvinayak01 keshavvinayak01 changed the title [WIP] Modified fx_importer to support hop_while_loop [TORCH] Modified fx_importer to support hop_while_loop Oct 15, 2025
@keshavvinayak01 keshavvinayak01 mentioned this pull request Oct 15, 2025
Open
@zjgarvey
Copy link
Collaborator

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.

@keshavvinayak01
Copy link
Author

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.

@zjgarvey
Copy link
Collaborator

@keshavvinayak01 can you sync with main?

@keshavvinayak01
Modified fx_importer to support hop_while_loop
c8c711c 
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
@keshavvinayak01 keshavvinayak01 force-pushed the keshavvinayak01/torch-hop-while branch from 3eb338f to c8c711c Compare October 22, 2025 16:42
@keshavvinayak01
Copy link
Author

Should be synced. Had to force push.

zjgarvey
zjgarvey requested changes Oct 22, 2025
View reviewed changes
Copy link
Collaborator

@zjgarvey zjgarvey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@keshavvinayak01
Addressed Comments | Simplified unique child_func_name creation
b250583 
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
zjgarvey
zjgarvey requested changes Oct 23, 2025
View reviewed changes
Copy link
Collaborator

@zjgarvey zjgarvey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

@keshavvinayak01
Addressed comments
db1e7e9 
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
@keshavvinayak01
Formatting
d9646c6 
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
@keshavvinayak01
Copy link
Author

I've addressed your comments, but specifically about the generated while_loop IR, you're right; prim will not support function calls in body. I'm going to modify the implementation to inline the function within the body.

@keshavvinayak01
Added children module imports to import_frozen_program flow
cc03291 
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
@keshavvinayak01
Formatting and reordered CHECKs
6a70e1c 
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
@keshavvinayak01
Changes done to TorchToScf:
85e3acd 
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 <keshavvinayakjha@gmail.com>
@keshavvinayak01
Copy link
Author

Had to modify TorchToScf conversion to support these changes. In general, the current state of TorchToScf would error out for any tensor values. @zjgarvey Please take a look and let me know. I also saw that PR https://github.com/llvm/torch-mlir/pull/3040/files solved this issue, but they were also removed?

@zjgarvey
Copy link
Collaborator

Had to modify TorchToScf conversion to support these changes. In general, the current state of TorchToScf would error out for any tensor values. @zjgarvey Please take a look and let me know. I also saw that PR https://github.com/llvm/torch-mlir/pull/3040/files solved this issue, but they were also removed?

That PR added tensor arg support for "for-like" loop conversion, but not "while-like" loop conversion.

zjgarvey
zjgarvey reviewed Oct 24, 2025
View reviewed changes
Copy link
Collaborator

@zjgarvey zjgarvey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@zjgarvey zjgarvey zjgarvey requested changes

@penguin-wwy penguin-wwy Awaiting requested review from penguin-wwy

@stellaraccident stellaraccident Awaiting requested review from stellaraccident

@rsuderman rsuderman Awaiting requested review from rsuderman

@Groverkss Groverkss Awaiting requested review from Groverkss

Requested changes must be addressed to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@keshavvinayak01 @zjgarvey