Cache backend FakeTensor dispatches during ExportPass replay (#21700)

[apullin]

Summary:

Split from D97528110 so the FakeTensor cache extension can be reviewed independently from cold-node fast-copy.

FakeTensorMode normally caches only aten, prim, and prims operations. ExportPass replay repeatedly dispatches deterministic backend operators from quantized_decomposed, tosa, and cortex_m, but PyTorch stores non-builtin negative cache entries for them. This change temporarily extends torch._library.utils.is_builtin while ExportPass runs, evicts only matching non-builtin bypass entries from the class-global and active ShapeEnv caches, and restores the exact original predicate in finally. Positive entries remain available across successive passes, which is the source of the speedup.

This intentionally remains a narrow process-global monkeypatch. Current FakeTensorMode has no per-instance cache-policy hook, and its non-symbolic cache is class-global. The reentrant lock serializes participating ExportPass contexts, but unrelated threads can observe the extended predicate while the context is active. Fully isolated behavior requires an upstream policy hook plus cache-policy identity or per-mode cache storage; keeping this in a separate diff makes that tradeoff explicit.

A/B benchmark against the fast-copy-only parent on the same host, each side run twice:

• CombinedControl U55 lowering: 124.343s / 126.491s with this change vs 148.465s / 149.726s in the parent. Warm speedup: 15.5%; two-run mean speedup: 15.9%.
• Synthetic U55 suite: 80.915s / 80.905s vs 81.238s / 82.371s. The 1.1% mean difference is within run-to-run noise; CombinedControl is the representative workload.

Differential Revision: D115374097

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21700

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 7 New Failures

As of commit 173c7bd with merge base fb5eedc ():

NEW FAILURES - The following jobs have failed:

• pull / test-arm-backend-no-driver (test_pytest_ops_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t c1b826a63b795c08da7428a69b48bdab83b3b8b94bb03b868e034dfb4ad39df1 /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  exir/emit/test/test_emit.py::TestEmit::test_emit_sym_int
• pull / unittest-editable / linux / linux-job (gh)
  exir/emit/test/test_emit.py::TestEmit::test_final_graph_module_update_debug_handle
• pull / unittest-editable / macos / macos-job (gh)
  exir/tests/test_memory_planning.py::TestMap::test_multi_map
• trunk / test-arm-backend-ethos-u (test_pytest_ops_ethos_u85) / linux-job (gh)
  RuntimeError: Command docker exec -t 9c2cbc64d378694caa7eef3d7e68c7687d9b0e95b34c1499a927d0220cb00e61 /exec failed with exit code 1
• trunk / unittest-release / linux / linux-job (gh)
  exir/emit/test/test_emit.py::TestEmit::test_emit_sym_int
• trunk / unittest-release / macos / macos-job (gh)
  exir/tests/test_memory_planning.py::TestMap::test_multi_state_plan

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@apullin has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115374097.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[apullin]

Hold on any reviewing here. There iis a way to do this without any monkeypatching, but it requires an extension to torch itself. That diff is written, will be submitted RFC to torch.