Fix Dynamo tracing of compile support checks (#185528)
Summary:
Newer torchvision roi_align calls torch._dynamo.utils.is_compile_supported(input.device.type) to choose between the Python ROIAlign path and the custom op path. Since torch._dynamo is skiplisted, Dynamo treated this helper as skipped and inserted graph breaks when vision_maskrcnn reached roi_align. That contributed to the torchvision pin graph-break regression.
Teach Dynamo's trace rules to represent is_compile_supported as a user function and mark the helper as a constant result. The helper depends on a constant device type and runtime compiler support, so Dynamo can evaluate it during tracing and guard the result instead of tracing into skiplisted internals. This is targeted to the one helper rather than broadening torch._dynamo inlining.
Benchmark Results:
Before: PYTHONPATH=/data/users/jansel/torchvision:$PYTHONPATH python benchmarks/dynamo/torchbench.py --ci --accuracy --timing --explain --print-compilation-time --disable-cudagraphs --inductor --inference --devices cuda --only vision_maskrcnn --log-graph-breaks --output /home/jansel/local/pytorch-issue-fixer/state/154259/vision_maskrcnn_before.csv produced 37 graphs, 1753 ops, 31 graph breaks, 4 unique graph breaks, and 112.507s compilation latency.
After: the same command writing vision_maskrcnn_after.csv produced 33 graphs, 1753 ops, 27 graph breaks, 3 unique graph breaks, and 107.582s compilation latency. The is_compile_supported skipped-function graph breaks disappeared; remaining local graph breaks are other existing classes.
X-link: https://github.com/pytorch/pytorch/pull/185528
Approved by: https://github.com/bobrenjc93
Reviewed By: jeanschmidt
Differential Revision: D117157052
fbshipit-source-id: 581bf5171e05b9fdc8035ac4a13611f9da3b56e0