SIGN IN SIGN UP

Fix Dynamo tracing of compile support checks (#185528)

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.

Test Plan:

python test/dynamo/test_trace_rules.py -k test_is_compile_supported_constant

python test/dynamo/test_trace_rules.py

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_after.csv

git diff --check

lintrunner -a

Fixes #154259

Generated by my agent

Pull Request resolved: https://github.com/pytorch/pytorch/pull/185528
Approved by: https://github.com/bobrenjc93
J
Jason Ansel committed
e26e72702afee0c6c6a1b9091bc364b02a434b2d
Parent: 3b0b5dc
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/24/2026, 4:53:18 AM