[dynamo, nested graph breaks] fix nested step graph break related issues (#162737)

Turns out codegen'ing a nested step graph break is significantly more complicated than first thought. The optimized function should actually do:
- call graph/load values/do side effects etc.
- call into the leaf's resume function, but skipped (this essentially step graph break function for just the leaf function)
- call into all the other resume functions, traced.

This PR also adds `torch._dynamo.step_unsupported()`, which can be used for internal testing purposes to better test step graph break handling.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162737
Approved by: https://github.com/Lucaskabela
ghstack dependencies: #160601
This commit is contained in:
William Wen
2025-10-08 10:50:05 -07:00
committed by PyTorch MergeBot
parent 486b4d2414
commit af4c29fea8
11 changed files with 542 additions and 242 deletions

View File

@ -263,6 +263,11 @@ class RecompileLimitExceeded(Unsupported):
pass
# debug exception thrown when tracing torch._dynamo.step_unsupported()
class StepUnsupported(TorchDynamoException):
pass
class UnsafeScriptObjectError(TorchDynamoException):
pass