[reland] Fix coreml to edge transform and lower

[metascroy]

Re-land of: #12564

Previous attempt had conflict with #12306 that caused CI failure.

---

The current design of using EDGE_DO_NOT_DECOMP to prevent decomposition has long standing issues, and often fails lowering when certain ops are requested for preservation. This shows up most notably in the CoreML backend, where most ops are requested for preservation.

As a band-aid, we introduced _remove_invalid_ops_for_not_decompose to cover certain kinds of ops. But when an op is encountered that we do not have an exception for, lowering still fails.

We also recently found another bug that shows up for SDPA related to contiguous views (https://fb.workplace.com/groups/pytorch.edge.users/permalink/1796069037930048/) that we still do not fully understand the root cause of.

EDGE_DO_NOT_DECOMP is actually only used to support the "check_op_support" argument in the partitioner; ops_to_not_decompose only modifies the default composition table.

In CoreML's case, "check_op_support" is not used, and the issues with EDGE_DO_NOT_DECOMP's design causes lots of lowering issues that are hard to keep up with. This PR enables a new path that bypasses EDGE_DO_NOT_DECOMP's when possible (_can_skip_using_EDGE_DO_NOT_DECOMP).

Long term, we need to address the buggy design of EDGE_DO_NOT_DECOMP. There are some ideas here: https://fb.workplace.com/groups/pytorch.edge2.team/permalink/1241898747065975/

cc @kimishpatel @YifanShenSZ @cymbalrush

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 3f67eef with merge base 413dee4 ():

NEW FAILURES - The following jobs have failed:

• pull / test-eval_llama-mmlu-linux / linux-job (gh)
  RuntimeError: Command docker exec -t e8687c19154158471f2297c9cf67624670b7c43312227ab6e4d7f2f74e049ebd /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  backends/apple/coreml/test/test_coreml_partitioner.py::TestCoreMLPartitioner::test_ops_to_not_decompose
• pull / unittest / macos / macos-job (gh)
  backends/apple/coreml/test/test_coreml_partitioner.py::TestCoreMLPartitioner::test_ops_to_not_decompose
• pull / unittest-editable / linux / linux-job (gh)
  backends/apple/coreml/test/test_coreml_partitioner.py::TestCoreMLPartitioner::test_ops_to_not_decompose
• pull / unittest-editable / macos / macos-job (gh)
  backends/apple/coreml/test/test_coreml_partitioner.py::TestCoreMLPartitioner::test_ops_to_not_decompose
• trunk / unittest-release / linux / linux-job (gh)
  backends/apple/coreml/test/test_coreml_partitioner.py::TestCoreMLPartitioner::test_ops_to_not_decompose
• trunk / unittest-release / macos / macos-job (gh)
  backends/apple/coreml/test/test_coreml_partitioner.py::TestCoreMLPartitioner::test_ops_to_not_decompose

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

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

[lucylq]

Do backends treat these the same at to_backend?

[metascroy]

They do not necessarily (CoreML does for many of them), but it's already applied in _generate_edge_program after the verifier, but before to_backend is called, so this doesn't change the behavior that exists today.

Some other options are:

1. Always filter out view ops from preservation
2. Query curr_partitioner.ops_to_not_decompose for view op's copy variant support, and only preserve the view op if the copy variant is present.

[metascroy]

I can also remove this line. It shouldn't be needed after you land D78535519

[billmguo]

can we test the model export and make sure last Friday issue not appear again with this diff?

[derekxu]

So this will fix the numerical lowering issues?

[metascroy]

So this will fix the numerical lowering issues?

When combined with #12665, I think it should (this switches to using torchao APIs for quantization)

[metascroy]

can we test the model export and make sure last Friday issue not appear again with this diff?

I will test to_edge_transform_and_lower on the ANE static model in the examples/apple/coreml/llama folder.