[dynamo, nested graph breaks] move cell codegen before side effects codegen (#160601)

This is needed because if we codegen cells for nested frames AFTER side effects, then reconstruction could get messed up. From below:

>The added test case demonstrates the reconstruction failure if we kept cell codegen at the original place (only happens with nested graph breaks since we reconstruct nested frame cells from VariableTracker rather than directly using LOAD_CLOSURE).

>At a high level, what happened before this change was that side_effects was pruning the cells (I don't recall exactly why this happens), and because cells were codegen'd after the side effects were applied, we were unable to properly reconstruct the cell. The error I was seeing was a list/tuple IndexError.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/160601
Approved by: https://github.com/mlazos
This commit is contained in:
William Wen
2025-10-08 10:50:04 -07:00
committed by PyTorch MergeBot
parent 8f83b3e71c
commit 486b4d2414
8 changed files with 181 additions and 74 deletions

View File

@ -449,7 +449,7 @@ class ZipVariable(IteratorVariable):
codegen.create_load_const("strict"),
codegen.create_load_const(self.strict),
create_instruction("BUILD_MAP", arg=1),
*create_call_function_ex(True),
*create_call_function_ex(True, False),
]
)
@ -487,7 +487,7 @@ class MapVariable(ZipVariable):
codegen.extend_output(
[
create_build_tuple(len(self.iterables) + 1),
*create_call_function_ex(False),
*create_call_function_ex(False, False),
]
)