Add partitioner pre-decomposition transform hook

[YufengShi-dudu]

Some backend-specific transformations must run before EXIR's initial run_decompositions({}) call because decomposition can remove the ATen patterns they need to inspect or rewrite. Previously, backends needed frontend wrappers or call-site-specific pass scheduling, exposing lowering order details outside the backend.

This PR adds a partitioner hook for this preprocessing. The relevant lowering order becomes:

to_edge_transform_and_lower()
  -> partitioner.transform_for_pre_decomposition()
  -> program.run_decompositions({})
  -> Edge transformation and partitioning

Pre-decomposition transforms from multiple partitioners run in the supplied order. The default implementation is a no-op, so existing partitioners and call sites remain unchanged.

Assisted by Codex.

Change-Id: I9e67b896c22fd2043c3b4f5439016016acc05362

cc @JacobSzwejbka @angelayi

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

As of commit 1110df8 with merge base 88a5f60 ():

NEW FAILURE - The following job has failed:

• trunk / test-custom-ops-macos (cmake) / macos-job (gh)
  torchaudio

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

[digantdesai]

IIRC do_not_decomp will complain if a partitioner asked for something to be not decomposed and then it didn't partition, is this true for this as well? I can imagine that as a good thing if we are planning to do some transforms here in this pass.

[YufengShi-dudu]

Thanks for your comment. No, we don't record the same info to check whether the generated ops remain undelegated after the following decomposition and partitioning. Tracking them may require a new metedata field.
transform_for_annotation_pipeline() also applies transforms without tracking generated nodes. Do you think we should introduce such tracking for this hook?

[digantdesai]

yeah quantizer being backend specific yet no direct enforcement for its actions is something I am not very happy about.

The reason I am a bit nervous this is, now we are modifying the graph in the paritioner and without any consequences if this hook misbehaves.

That said, with your hook, since its a pass, writing a '_sanity_check_graph_for_non_decomp_ops` like fn can get tricky if multiple passes or even multiple partitioners have worked on the graph.

[digantdesai]

ok let's hope this doesn't get out of control, I will stamp this as is.

[YufengShi-dudu]

thank you, we need to be more cautious when adding new passes to this hook.

[digantdesai]

Thanks. Also please list your motivation for this hook in the PR summary before merging.