frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[dynamo, nested graph breaks] support very simple nested graph breaks (#159329)

Browse Source 
e.g. this graph breaks once now:
```python
import torch

torch._dynamo.config.nested_graph_breaks = True

def inner(x):
    x = x + 1
    torch._dynamo.graph_break()
    return x + 2

@torch.compile(backend="eager")
def outer(x):
    return inner(x)

print(outer(torch.ones(3)))
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/159329
Approved by: https://github.com/anijain2305
ghstack dependencies: #157971, #159281, #144516
This commit is contained in:
William Wen
2025-08-25 13:27:42 -07:00
committed by PyTorch MergeBot
parent 9a756c2d71
commit 8dab6d4c41
5 changed files with 226 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
torch/_dynamo/bytecode_transformation.py
View File

@ -212,6 +212,10 @@ def create_jump_absolute(target: Instruction) -> Instruction:
return create_instruction(inst, target=target)
def is_jump_absolute(target: Instruction) -> bool:
return target.opname in ("JUMP_FORWARD", "JUMP_ABSOLUTE")
def create_load_const(val: Any, checked: bool = True) -> Instruction:
"""
In general we should only create `LOAD_CONST` for immutable objects, but
@ -504,15 +508,6 @@ def create_binary_slice(
]
def create_reverse(n: int) -> list[Instruction]:
# Reverse the top n values on the stack
# UNPACK_SEQUENCE reverses the sequence
return [
create_instruction("BUILD_TUPLE", arg=n),
create_instruction("UNPACK_SEQUENCE", arg=n),
]
def lnotab_writer(
lineno: int, byteno: int = 0
) -> tuple[list[int], Callable[[int, int], None]]: