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] Refactor convert_frame._compile_inner to return compiled bytecode + output graph. [3/n] (#160855)

Browse Source 
We are refactoring dynamo code for convert frame so that we can have modularized pieces sharable between different compiler frontends (e.g. torch.compile, precompile and torch.export).

This PR adds a new helper function compile_frame() which takes a bytecode and a transform function and return compiled bytecode + output graph as DynamoOutput type.

Differential Revision: [D80430802](https://our.internmc.facebook.com/intern/diff/D80430802/)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/160855
Approved by: https://github.com/tugsbayasgalan
ghstack dependencies: #160814, #160815
This commit is contained in:
zhxchen17
2025-08-18 12:52:15 -07:00
committed by PyTorch MergeBot
parent b3e215b864
commit 9e050b6339
5 changed files with 81 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
torch/_dynamo/bytecode_transformation.py
View File

@ -1259,7 +1259,7 @@ def debug_bytes(*args: bytes) -> str:
def debug_checks(code: types.CodeType) -> None:
"""Make sure our assembler produces same bytes as we start with"""
dode = transform_code_object(code, lambda x, y: None, safe=True)
dode, _ = transform_code_object(code, lambda x, y: None, safe=True)
assert code.co_code == dode.co_code, debug_bytes(code.co_code, dode.co_code)
assert code.co_lnotab == dode.co_lnotab, debug_bytes(code.co_lnotab, dode.co_lnotab)
@ -1456,7 +1456,7 @@ def transform_code_object(
[list[Instruction], dict[str, Any]], Optional["DynamoTracerOutput"]
],
safe: bool = False,
) -> types.CodeType:
) -> tuple[types.CodeType, Optional["DynamoTracerOutput"]]:
keys = get_code_keys()
code_options = {k: getattr(code, k) for k in keys}
assert len(code_options["co_varnames"]) == code_options["co_nlocals"]
@ -1465,8 +1465,9 @@ def transform_code_object(
# propagate line nums again for added instructions
propagate_line_nums(instructions)
transformations(instructions, code_options)
return clean_and_assemble_instructions(instructions, keys, code_options)[1]
tracer_output = transformations(instructions, code_options)
_, bytecode = clean_and_assemble_instructions(instructions, keys, code_options)
return bytecode, tracer_output
def clean_and_assemble_instructions(