Sufficiently deep nesting + inlining + closure fails with KeyError #99639
Description
🐛 Describe the bug
def closure_repro():
def outer(a):
return a + 1
def indirect(x):
return direct(x)
def direct(x):
def inner(b):
return b + 2
def deep(c):
d = outer(c)
return inner(d)
return deep(x)
dynamo.export(indirect, torch.randn(3)) # changing this to export `direct` instead works!
closure_repro()Error:
[2023-04-20 09:53:27,732] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing indirect
[2023-04-20 09:53:27,738] torch._dynamo.symbolic_convert: [DEBUG] TRACE starts_line <...>scripts/avik/pt2/example.py:96
[2023-04-20 09:53:27,738] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_DEREF direct []
[2023-04-20 09:53:27,738] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_FAST x [UserFunctionVariable()]
[2023-04-20 09:53:27,739] torch._dynamo.symbolic_convert: [DEBUG] TRACE CALL_FUNCTION 1 [UserFunctionVariable(), TensorVariable()]
[2023-04-20 09:53:27,740] torch._dynamo.symbolic_convert: [DEBUG] INLINING <code object direct at 0x7f0e8d5f2b30, file "<...>scripts/avik/pt2/example.py", line 98>
[2023-04-20 09:53:27,740] torch._dynamo.symbolic_convert: [DEBUG] TRACE starts_line <...>scripts/avik/pt2/example.py:99
[2023-04-20 09:53:27,740] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_CONST <code object inner at 0x7f0e8d5f7870, file "<...>scripts/avik/pt2/example.py", line 99> []
[2023-04-20 09:53:27,740] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_CONST closure_repro.<locals>.direct.<locals>.inner [ConstantVariable(code)]
[2023-04-20 09:53:27,741] torch._dynamo.symbolic_convert: [DEBUG] TRACE MAKE_FUNCTION 0 [ConstantVariable(code), ConstantVariable(str)]
[2023-04-20 09:53:27,741] torch._dynamo.symbolic_convert: [DEBUG] TRACE STORE_DEREF inner [NestedUserFunctionVariable()]
[2023-04-20 09:53:27,741] torch._dynamo.symbolic_convert: [DEBUG] TRACE starts_line <...>scripts/avik/pt2/example.py:102
[2023-04-20 09:53:27,741] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_CLOSURE inner []
...closure_cells={'inner': NewCellVariable()}, symbolic_locals={'x': TensorVariable(), 'outer': UserFunctionVariable()}
[2023-04-20 09:53:27,741] torch._dynamo.symbolic_convert: [DEBUG] TRACE LOAD_CLOSURE outer [NewCellVariable()]
...closure_cells={'inner': NewCellVariable()}, symbolic_locals={'x': TensorVariable(), 'outer': UserFunctionVariable()}
[2023-04-20 09:53:27,741] torch._dynamo.symbolic_convert: [DEBUG] FAILED INLINING <code object direct at 0x7f0e8d5f2b30, file "<...>scripts/avik/pt2/example.py", line 98>
[2023-04-20 09:53:27,742] torch._dynamo.output_graph: [DEBUG] restore_graphstate: removed 0 nodes
Traceback (most recent call last):
File "<...>torch/_dynamo/convert_frame.py", line 401, in _compile
out_code = transform_code_object(code, transform)
File "<...>torch/_dynamo/bytecode_transformation.py", line 1000, in transform_code_object
transformations(instructions, code_options)
File "<...>torch/_dynamo/convert_frame.py", line 386, in transform
tracer.run()
File "<...>torch/_dynamo/symbolic_convert.py", line 1942, in run
super().run()
File "<...>torch/_dynamo/symbolic_convert.py", line 656, in run
and self.step()
File "<...>torch/_dynamo/symbolic_convert.py", line 616, in step
getattr(self, inst.opname)(inst)
File "<...>torch/_dynamo/symbolic_convert.py", line 351, in wrapper
return inner_fn(self, inst)
File "<...>torch/_dynamo/symbolic_convert.py", line 1065, in CALL_FUNCTION
self.call_function(fn, args, {})
File "<...>torch/_dynamo/symbolic_convert.py", line 507, in call_function
self.push(fn.call_function(self, args, kwargs))
File "<...>torch/_dynamo/variables/functions.py", line 269, in call_function
return super().call_function(tx, args, kwargs)
File "<...>torch/_dynamo/variables/functions.py", line 102, in call_function
return tx.inline_user_function_return(
File "<...>torch/_dynamo/symbolic_convert.py", line 543, in inline_user_function_return
result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
File "<...>torch/_dynamo/symbolic_convert.py", line 2047, in inline_call
return cls.inline_call_(parent, func, args, kwargs)
File "<...>torch/_dynamo/symbolic_convert.py", line 2125, in inline_call_
tracer.run()
File "<...>torch/_dynamo/symbolic_convert.py", line 656, in run
and self.step()
File "<...>torch/_dynamo/symbolic_convert.py", line 616, in step
getattr(self, inst.opname)(inst)
File "<...>torch/_dynamo/symbolic_convert.py", line 2242, in LOAD_CLOSURE
self.push(self.closure_cells[inst.argval])
KeyError: outer
from user code:
File "<...>scripts/avik/pt2/example.py", line 96, in indirect
return direct(x)
File "<...>scripts/avik/pt2/example.py", line 102, in direct
def deep(c):
Versions
v2.0
cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire