Choosing ops_to_preserve by delegating to pattern

[RahulC7]

Summary:

Context

The trace function in compiler.py returns the static graph representation of the model. Because certain hardwares are very optimized for certain operations, we don't want to decompose those operations in our graph.

Thus, currently, the function passes in ops_to_keep to trace_fn

https://www.internalfb.com/code/fbsource/[ada71b37b46a898065851182cd48b48b321a5d12]/fbcode/executorch/backends/cadence/aot/compiler.py?lines=59-86

which then removes these operations before decomposing the graph:
https://www.internalfb.com/code/fbsource/[ada71b37b46a898065851182cd48b48b321a5d12]/fbcode/executorch/backends/cadence/aot/compiler_funcs.py?lines=33-36

Issue

Right now, the ops_to_keep is hardcoded, and there's no easy way to customize which ops to keep. View this comment thread for more details.

The tl;dr is that there should be stricter separation of concerns between the compiler and the hardware -- the compiler shouldn't need to know what ops to keep or not.

Possible Solutions

| Solution | Pros | Cons | Implementation |
| -- |
| Delegate ops_to_keep to QuantizationPattern [chosen solution] | Clear separation of concerns, semantically correct, works well with current composition structure of classes, easily extensible | Changes existing behavior. i.e, CadenceFusedConvReluQuantizer doesn't have all patterns, but current behavior is that compiler is actually keeping all default ops. Direct inheritance doesn't work due to current behavior. | D84524714 |
| Delegate ops_to_keep to CadenceQuantizer | Keeps existing behavior, simple for users(just add additional ops to keep) | Doesn't work well with more complex compositions(none exist), making inheritance additive is somewhat complicated, not sure if it makes sense semantically, harder to customize | D84461403 [not used] |

This Diff

For now, we delegate ops_to_keep to QuantizationPattern(first solution). We

1. Create a recursive method to get all the operations to preserve, with the base being the pattern defined for a CadenceAtenQuantizer.
2. Use our method to find the ops to keep in compiler.py
3. Create a unit test

Now, the ops_to_keep is defined using the patterns themselves!

Differential Revision: D84524714

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Unrelated Failure

As of commit 85c5ac7 with merge base 9413da0 ():

NEW FAILURES - The following jobs have failed:

• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 81380bb4d56e04a730d29838197026244d4eb0b84b666f85b43ae0cbdd372a5a /exec failed with exit code 1
• Test CUDA Builds / export-voxtral-cuda-artifact / linux-job (gh)
  RuntimeError: Command docker exec -t d75d33c65023c9a6b68af5422b0dff71d1b3036d02bdfd10ad633c9193e6f565 /exec failed with exit code 2

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

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

• pull / unittest-arm-backend-with-no-fvp (test_pytest_models) / linux-job (gh) (trunk failure)

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

[meta-codesync]

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

[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.