module.compile() installs torch.compile(self._call_impl). For built-in leaf modules such as Conv2d, both Module._call_impl and the module forward live in skip-rule files. The torch.compile(module) path already handles this by wrapping modules whose forward has a skip rule in wrap_inline, but module.compile() excluded _call_impl from top-level wrapping. This left no capturable frame, so a custom backend was never called.
This teaches Dynamo to make the same exception for bound nn.Module _call_impl methods when their forward has a skip rule. _wrapped_call_impl and _lazy_forward remain excluded; Module.compile() uses _call_impl directly, and broadening those other private paths is unnecessary for this bug.
The alternative would be setting wrap_top_frame for module.compile globally or changing Module.compile to pass the module object through torch.compile. Both are broader behavior changes; this keeps the fix targeted and reuses the existing wrap_inline mechanism.
Fixes #150915
Generated by my agent
Benchmark Results:
- Microbenchmark command: 30 iterations of ToyModel().compile(backend=CompileCounter()); model(x) and torch.nn.Conv2d(3, 3, 3).compile(backend=CompileCounter()); model(x), timing compile()+first call with torch._dynamo.reset() each iteration.
- Unchanged ToyModel path: main median 23.810 ms, patched median 23.017 ms, frame_count [1] in both runs.
- Conv2d issue path: main median 14.657 ms with frame_count [0] and op_count [0]; patched median 40.624 ms with frame_count [1] and op_count [1]. The first call is slower because it now invokes the backend and compiles the intended graph instead of skipping compilation.
Test Plan:
- python test/dynamo/test_compile.py InPlaceCompilationTests.test_compilation_builtin_leaf_module
- python test/dynamo/test_compile.py
- python test/dynamo/test_hooks.py HooksTests.test_wrap_top_frame_with_hooks
- lintrunner -a
- git diff --cached --check
Pull Request resolved: https://github.com/pytorch/pytorch/pull/185722
Approved by: https://github.com/williamwen42