[Dynamo] make bytecode of resume function resemble natural bytecode

[youkaichao]

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/126630

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (3 Unrelated Failures)

As of commit b7ad5f3 with merge base a8195f2 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_timm, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  Error response from daemon: a prune operation is already running

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#126884)
  eca_halonext26ts

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[youkaichao]

When we deal with natural bytecode (Python bytecode that is generated by Python compiler, from source code), we can observe that co_freevars seem to be sorted, no matter how they are declared:

def helper_outer_function():
    data_ptr = None
    last_dim_size = None
    last_two_dims_size = None
    shape = None
    stride = None
    qkv_format = None
    def dummy1():
        nonlocal data_ptr, last_dim_size, last_two_dims_size, shape, stride, qkv_format
        # Access variables in the order they were declared
        _ = data_ptr
        _ = last_dim_size
        _ = last_two_dims_size
        _ = shape
        _ = stride
        _ = qkv_format
        print(data_ptr)
        print(last_dim_size)
        print(last_two_dims_size)
        print(shape)
        print(stride)
        print(qkv_format)
    dummy1()

print(helper_outer_function.__code__.co_consts[1].co_freevars)

The output is:

('data_ptr',
 'last_dim_size',
 'last_two_dims_size',
 'qkv_format',
 'shape',
 'stride')

Note that our declaration order is data_ptr, last_dim_size, last_two_dims_size, shape, stride, qkv_format , but Python sorts them.

If Dynamo generated bytecode does not obey this rule, it means we cannot generate source code that can compile the same bytecode as Dynamo, which makes it impossible to understand Dynamo bytecode by decompiling it into source code.

[youkaichao]

Note: we require the new bytecode has exactly the same co_freevar as the old one, in order to have faster new frame construction. It is proposed by me in #115062 .

[youkaichao]

reference:

https://github.com/python/cpython/blob/caf6064a1bc15ac344afd78b780188e60b9c628e/Python/compile.c#L530-L534

/* Sort the keys so that we have a deterministic order on the indexes
   saved in the returned dictionary.  These indexes are used as indexes
   into the free and cell var storage.  Therefore if they aren't
   deterministic, then the generated bytecode is not deterministic.
*/

indexes of free and cell var storage are sorted.

[youkaichao]

TODO: how to add test for this. We need to generate a resume function with freevars.

[ezyang]

Trying @williamwen42 as reviewer, but this looks pretty harmless, shout if it gets lost

[williamwen42]

LGTM - just add a test

[youkaichao]

@williamwen42 thanks for the reply. Can you give me some guidance on how to add tests? In particular, I don't know how to manually instruct Dynamo to generate a resume function as I want.

[williamwen42]

You can compile a function with a graph break and then search globals/locals for the resume function:

import torch

def fn(x):
    x = x + 1
    torch._dynamo.graph_break()
    x = torch.sin(x)
    return x

opt_fn = torch.compile(fn, backend="eager")
opt_fn(torch.randn(10))

for k, v in list(globals().items()):
    if k.startswith("__resume_at"):
        print(k)
        print(v)

[youkaichao]

This kind of resume function does not have freevars:

__resume_at_16_3.__code__.co_freevars == ()

[williamwen42]

You can use a different function that has freevars - compile a function with a closure?

[youkaichao]

Still no freevars:

import torch

def fn(x):
    x = x + 1
    @torch.compile(backend="eager")
    def inner(x):
        x = x + 1
        torch._dynamo.graph_break()
        x = x * 2
        return x
    y = inner(torch.sin(x))
    return y

fn(torch.randn(10))

for k, v in list(globals().items()):
    if k.startswith("__resume_at"):
        print(k)
        print(v)

[youkaichao]

When will resume function contain free vars?

[williamwen42]

Looks like resume functions with freevars are not exposed to the global scope, so we'll need to add something like

            # expose code object for debugging purposes
            self.output.install_global_unsafe(name, new_code)

before the cg.make_function_with_closure(name, new_code, True, stack_len) line in def create_call_resume_at (symbolic_convert.py).

Then a function like this should work:

import torch

def create():
    cl = 1
    def fn(x):
        x = x + 1
        torch._dynamo.graph_break()
        x = x + cl
        return x
    return fn

fn = create()
opt_fn = torch.compile(fn, backend="eager")
print(opt_fn(torch.randn(10)))

breakpoint()
for k, v in list(globals().items()):
    if k.startswith("__resume_at"):
        print(k)
        print(v)
        print(v.co_freevars)
        print(v.co_cellvars)

[youkaichao]

@williamwen42 thanks for the guidance! do you know if the test failures in the commit are related with the changes in this PR?

[williamwen42]

They don't look related, but we'll see what CI shows.

[youkaichao]

@williamwen42 can you take a look at whether ci test failures are related?

[williamwen42]

They don't look related - we're having some issues with CI atm.

[youkaichao]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 15 jobs have failed, first few of them are: pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, linux.2xlarge), inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu), inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu), inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu), inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 1, linux.g5.4xlarge.nvidia.gpu, unstable)

Details for Dev Infra team

Raised by workflow job

[youkaichao]

@williamwen42 then how can we merge this? do we need to wait until the ci team fixes the issue?

[huydhn]

@pytorchbot merge -r

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased youkaichao-patch-1 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout youkaichao-patch-1 && git pull --rebase)

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here