[Bug]: remove_some_temp() merges temporary variables too aggressively, causing `KeyError` in compiled code

[Qiaoza666]

Your current environment

PyTorch version: 2.4.1+cu121

Python version: 3.8.17 (default, Jul 5 2023, 21:04:15) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.250-9-velinux1u1-amd64-x86_64-with-glibc2.17

Versions of relevant libraries:
[pip3] depyf==0.19.0
[pip3] torch==2.4.1
[pip3] torchaudio==2.4.0+cu121
[pip3] torchvision==0.19.0+cu121
[conda] torch 2.4.1 pypi_0 pypi
[conda] torchaudio 2.4.0+cu121 pypi_0 pypi
[conda] torchvision 0.19.0+cu121 pypi_0 pypi

🐛 Describe the bug

When using depyf.prepare_debug together with torch.compile, the decompiled code may inline a temp dict across a side-effecting call (preds.clear()), turning a safe pattern into an unsafe one and causing KeyError.

Describe the bug

Using depyf.prepare_debug with torch.compile, the decompiled/transformed code sometimes becomes invalid.
Specifically, the temporary variable introduced to preserve evaluation order is merged away by remove_some_temp(), causing preds.clear() to run before reading preds['key1'], which then raises a KeyError.

Minimal repro

import torch
import depyf
@torch.compile
def toy_example(a, b):
    preds = {"key1": 1, "key2": 2}
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    preds.update({"b": b})
    return x * b
def main():
    for _ in range(10):
        toy_example(torch.randn(10), torch.randn(10))
if __name__ == "__main__":
    with depyf.prepare_debug("./save_dir"):
        main()

Redacted traceback

Traceback (most recent call last):
  File "main.py", line 38, in <module>
    main()
  File "main.py", line 25, in main
    toy_example(torch.randn(10), torch.randn(10))
  File "<site-packages>/torch/_dynamo/eval_frame.py", line 433, in _fn
    return fn(*args, **kwargs)
  File "<dump>/__transformed_code_0_for_toy_example.py", line 6, in __transformed_code_0_for_toy_example
    return __resume_at_40_2(b, preds, x)
  File "<dump>/__transformed_code_0_for_torch_dynamo_resume_in_toy_example_at_11.py", line 6, in __transformed_code_0_for_torch_dynamo_resume_in_toy_example_at_11
    preds.update({'key1': preds['key1'], 'key2': preds['key2'], 'b': graph_out_0[1]})
KeyError: 'key1'

The generated transformed code shows the problematic ordering:

def __transformed_code_0_for_torch_dynamo_resume_in_toy_example_at_11(b, preds, x):
    graph_out_0 = __compiled_fn_5(b, x)
    preds.clear()
    preds.update({'key1': preds['key1'], 'key2': preds['key2'], 'b': graph_out_0[1]})
    return graph_out_0[0]

Here preds.clear() executes before reading preds['key1'], leading to KeyError.

Analysis

In Python bytecode, the arguments for preds.update(...) are loaded before preds.clear() is executed. Depyf initially introduces a temporary variable to preserve this order, e.g.:

__temp_6 = {'key1': preds['key1'], 'key2': preds['key2'], 'b': graph_out_0[1]}
preds.clear()
preds.update(__temp_6)

However, the post-processing step remove_some_temp() merges away the temp variable (because it appears exactly twice and doesn’t cross an indentation boundary), changing the code into:

preds.clear()
preds.update({'key1': preds['key1'], 'key2': preds['key2'], 'b': graph_out_0[1]})

This is incorrect in the presence of side effects between temp assignment and usage (here, preds.clear()).

In short, remove_some_temp() is currently too aggressive: it only checks occurrence counts and indentation scope, but does not account for side-effecting calls that invalidate inlining.

Expected behavior

remove_some_temp() should retain the temporary variable when there is any potential side effect between the temp assignment and its use, to preserve evaluation order and avoid KeyError.

Environment

• Python: 3.8
• PyTorch: 2.4.1
• Depyf: latest GitHub
• OS: Ubuntu 20.04

Additional context

A possible direction is to extend remove_some_temp() with a side-effect / alias analysis between the temp assignment site and its use (e.g., detecting mutating calls like .clear() on potentially aliased objects).
I have a rough prototype that attempts such detection, but I’m not confident enough in its completeness to propose it here. I’m happy to help validate any fix or test patches on my setup.

[youkaichao]

That's a good point, thanks for the detailed analysis.

However, doing side-effect analysis would be too difficult in Python, if you only have static code. A seemingly harmless statement like a + 1 can have side effect, if a overrides the + operation.

Maybe we can add a flag to disable the remove_some_temp stuff?

[Qiaoza666]

That's a good point, thanks for the detailed analysis.

However, doing side-effect analysis would be too difficult in Python, if you only have static code. A seemingly harmless statement like can have side effect, if overrides the operation.a + 1``a``+

Maybe we can add a flag to disable the stuff?remove_some_temp

Agreed. Full side-effect analysis isn’t practical; a disable flag would be great.

[youkaichao]

Coming back to this again, I met this issue too. The problem is Dynamo uses COPY to duplicate some elements in the stack, and then use some bytecode like DELETE_SUBSCR to mutate the last one.

Something like:

a = [1, 2, 3]
b = a[1] # dup
del a[1]
fun(b)

this is problematic, because right now COPY will directly duplicate objects without creating another reference.

I have a hacky solution to eliminate the side-effect, i.e. removing del a[1] silently if a[-1] exists in the stack.

[youkaichao]

see 7d7a9bc for details.

[youkaichao]

closed by #87

[youkaichao]

@Qiaoza666 feel free to directly contact me via email, or continue in this thread if you have further questions.