[dynamo] another 3.11 resume codegen KeyError #166176
Description
Originally reported in an internal model.
import torch
def fn(x):
torch._dynamo.graph_break()
with torch.no_grad():
with torch.no_grad():
torch._dynamo.graph_break()
return x + 1
inp = torch.ones(3)
opt_m = torch.compile(fn, backend="eager")
opt_m(inp)gives error
...
File "/data/users/williamwen/pytorch4/torch/_dynamo/symbolic_convert.py", line 1030, in handle_graph_break
+ self.create_call_resume_at(
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/users/williamwen/pytorch4/torch/_dynamo/symbolic_convert.py", line 2975, in create_call_resume_at
resume_code, resume_name = cur_tx.create_resume(
^^^^^^^^^^^^^^^^^^^^^
File "/data/users/williamwen/pytorch4/torch/_dynamo/symbolic_convert.py", line 2879, in create_resume
new_code: types.CodeType = ContinueExecutionCache.lookup(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/users/williamwen/pytorch4/torch/_dynamo/resume_execution.py", line 307, in lookup
cls.cache[code][key] = cls.generate(code, lineno, *key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/users/williamwen/pytorch4/torch/_dynamo/resume_execution.py", line 336, in generate
return cls.generate_based_on_original_code_object(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/users/williamwen/pytorch4/torch/_dynamo/resume_execution.py", line 683, in generate_based_on_original_code_object
setup_fn_target_offsets = tuple(
^^^^^^
File "/data/users/williamwen/pytorch4/torch/_dynamo/resume_execution.py", line 684, in <genexpr>
meta.block_target_offset_remap[new_offset][n]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^
torch._dynamo.exc.InternalTorchDynamoError: KeyError: 236
from user code:
File "/data/users/williamwen/pytorch4/playground.py", line 18, in torch_dynamo_resume_in_fn_at_14
y = gn(x)
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @Lucaskabela