[CoreML] Cache the quantizer's pattern graphs instead of rebuilding them per model

[vijay-kapse]

Fixes the root cause behind #21623.

Problem

The CoreML quantizer builds its subgraph-match patterns by exporting small pattern modules,
and it rebuilds them for every model it annotates. Instrumenting _get_aten_graph_module
for a single prepare_pt2e on a one-conv module:

prepare_pt2e total           :   28.51s
pattern graph builds         : 405
  time in them               :   28.24s   (99.1% of prepare_pt2e)
  mean per build             :    0.070s

The patterns do not depend on the model being annotated. Instrumenting the getters with
their actual arguments across three different models gives 1215 builds with only 405 distinct
argument signatures, and the key set for a linear model is identical to the key set for a
conv model. Every model after the first rebuilds patterns that were already built.

Nothing amortizes it. Five consecutive prepare_pt2e calls in one process, each on a fresh
module:

28.94s, 32.24s, 34.23s, 35.40s, 32.16s

For context on where that lands, the quantize stage is 95% of the per-case cost of the
coreml_static_int8 flow, and prepare_pt2e is 99.4% of the quantize stage:

sub-step	time	share
torch.export(strict=True)	0.15 s	0.5%
CoreMLQuantizer(...) ctor	0.00 s	0.0%
prepare_pt2e()	28.29 s	99.4%
calibrate	0.02 s	0.1%
convert_pt2e()	0.01 s	0.0%

Change

Memoizes the five pattern getters. Each builds its example_inputs internally at fixed
shapes and takes only hashable arguments (conv_dim: int, act_fn, act_in_place: bool),
so each is a pure function of its arguments.

The cached graph cannot be handed out directly, which is worth knowing because the
obvious lru_cache is incorrect here. SubgraphMatcherWithNameNodeMap.__init__ calls
_split_to_graph_and_name_node_map(pattern_gm), which splits name_node_map out of the
pattern's return tuple in place; a second consumer of the same object then fails with:

AssertionError: Expecting the pattern graph to return a tuple

So callers get a copy. That is still overwhelmingly worth it: building one pattern costs
62.0 ms, deepcopy of it costs 0.7 ms, roughly 93x cheaper.

Measured

prepare_pt2e + convert_pt2e on an M4, alternating conv and linear modules so the cache is
exercised across model types:

case
1	conv2d	29.92 s (cold, builds the 405 patterns once)
2	conv2d	0.53 s
3	linear	0.54 s
4	conv2d	0.55 s
5	linear	0.54 s

55x on warm calls.

Output is unchanged

The fully converted graph for a conv module and for a linear module is byte identical before
and after this change, both in node targets and in print_readable() output — produced by
running the same script against the unpatched and patched file and diffing.

Effect on #21623

Per-case quantize in that suite drops from roughly 29 s to 0.5 s, with one ~29 s warmup per
xdist worker process. I could not run all 467 operator tests to completion locally, so the
suite-level figure is arithmetic rather than a measurement: total quantize work goes from
about 467 x 29 s to (workers x 29 s) + (467 x 0.5 s), which should bring the flow near the
coreml flow's 11 min instead of hitting the 180 min cap. A nightly run would confirm it.

This also speeds up every other user of the CoreML quantizer, not only CI.

Environment: macOS 26.5 on an M4, torch 2.14.1, coremltools 9.0 (matching
COREMLTOOLS_VERSION="9.0" and the torch==2.14.0 pin).

cc @kimishpatel @YifanShenSZ @cymbalrush @metascroy @nil-is-all

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 2 Unrelated Failures

As of commit a212d83 with merge base d036d50 ():

NEW FAILURE - The following job has failed:

• trunk / test-qnn-optimum-model (fp32, dit) / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Test CoreML Backend / test-coreml / test-backend-macos (coreml, models) / macos-job (gh) (trunk failure)
  ##[error]The operation was canceled.

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

• trunk / test-samsung-models-linux / linux-job (gh) (#23507)
  test_yolo26_fp16

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

[meta-cla]

Hi @vijay-kapse!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: vijay-kapse / name: Vijay Suryakant Kapse (b172a36)

[vijay-kapse]

@pytorchbot label "release notes: apple"

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[vijay-kapse]

Correcting a number I got wrong above, before anyone spends time on it.

I originally wrote that a single prepare_pt2e performs "405 pattern builds of which only 9
are distinct, 396 duplicates (97.8%)". That characterization was wrong. It came from
keying my instrumentation on type(pattern).__qualname__, and every pattern class is defined
inside its builder function, so conv1d/conv2d/conv_transpose and the activation variants all
report the same qualname and collapsed together in my count.

Re-instrumented with the getters' actual arguments:

per prepare_pt2e	405 builds, all 405 distinct within the call
across 3 different models	1215 builds, 405 distinct overall
key set, linear model vs conv model	identical

So there is no duplication within a single annotate. The waste is entirely across models:
the pattern set is completely independent of the model being annotated, and every model after
the first rebuilds all 405 from scratch. For the 467-case operators suite that is roughly
189,000 builds where 405 would do.

Everything else in the writeup is unaffected and was measured directly rather than inferred:
the 405 builds being 99.1% of prepare_pt2e, the 95%/99.4% stage split, the five
non-amortizing consecutive calls, the 55x warm speedup, and the byte-identical quantized
output. The fix is the same, and if anything this is a cleaner justification for it — the
patterns are model-independent, which is exactly why memoizing them is safe.

The commit message and PR description are updated.

[metascroy]

Testing CI on trunk

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[vijay-kapse]

Thanks for triggering CI. Three jobs went red and I do not think any of them is this PR — here is what I checked.

test-phi-3-mini-runner-linux and test-samsung-models-linux both fail in the runner setup, not in a test:

Executing the custom container implementation failed. Please contact your self hosted runner administrator.
Credentials could not be loaded, please check your action inputs: Could not load credentials from any providers

configure-aws-credentials cannot load secrets on a pull request from a fork, so the container never comes up. test-phi-3-mini-runner-linux passes on the ten most recent open PRs, all of which are branches in this repo rather than forks.

test-coreml / test-backend-macos (coreml, models) hit the runner time limit rather than failing an assertion:

The job has exceeded the maximum execution time of 3h0m0s

Run script ran 23:07:54 to 02:06:53, so 2h59m, and was cancelled at the cap. The same job does the same thing on main. Run 38007909820 (push to main, f2575be5ac) ran test-backend-macos (coreml, models) from 00:21:41 to 03:22:11 and was cancelled with the identical annotation, with no commit from this PR in it. The scheduled run that started at 02:23:40 is still going at the time of writing.

For contrast, in the scheduled trunk run test-backend-macos (coreml_static_int8, models) finished in 20 minutes (02:23:40 to 02:43:53), and (coreml, operators) finishes in minutes on both trunk and here. So it looks specific to the (coreml, models) variant rather than to anything here.

This PR only touches backends/apple/coreml/quantizer/_coreml_quantizer_utils.py, and it caches pattern graphs that were previously rebuilt on every call, so it should move that path in the opposite direction. Happy to rebase or to rerun whichever jobs you want if that is useful.

[metascroy]

I'm testing whether this change brings the coreml_static_int8 operators suite within the existing 180-minute CI timeout. The temporary workflow commit enables that suite for this measurement.

Fresh CI run: https://github.com/pytorch/executorch/actions/runs/38280085324

I'll check the operators job's completion time separately from the unrelated coreml models job.

[metascroy]

Looks good to land once CI is green. Thanks for the contribution!