Keep output specs in sync in ExportedProgramPassBase (#22864)

[mcremon-meta]

Summary:

A pass that replaces the graph's output node leaves output_specs naming the
node that is no longer there. Nothing catches it at the time; it surfaces later
as a signature that disagrees with the graph, and every pass author has to
remember to fix it up by hand.

ExportedProgramPassBase.__call__ now realigns the specs with the graph between
the pass and the postcondition check:

self.requires(exported_program)
res = self.call(exported_program)
signature_modified = _sync_output_specs(res.exported_program)
self.ensures(res.exported_program)

Rewriting a spec to name the current output node is always safe. Changing an
output between a node and a literal is not, so _sync_output_specs raises
rather than guessing, and a length mismatch between outputs and specs raises
too -- a pass that adds or removes an output is expected to maintain its own
signature.

Two consequences worth calling out:

• A pass that only renames an output now reports modified=True even if it
  returned False, because the signature did change. That is what callers
  driving a pass to fixpoint need to see.
• ensures() now receives the pass's result rather than the program that went
  in. The previous behaviour looks like an oversight -- a postcondition check
  that inspects the pre-pass program cannot check much -- and no pass in the
  tree relies on it: of the subclasses of ExportedProgramPassBase, exactly one
  overrides ensures(), and it calls exported_program.validate(), which wants
  the result.

This was previously implemented as a fused_quant-local subclass. It is not
specific to that package, so it moves here; the subclass goes away in the diff
above this one.

Reviewed By: mergennachin, DrJessop

Differential Revision: D120259092

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• linux.dgx.b200 runner queuing exceeds 4 hours and dgxb200-05 "CUDA not available"

✅ No Failures

As of commit d174667 with merge base 1cf2927 ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D120259092.

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

[mergennachin]

get_replace_hook() also rewrites input_specs, so I think we need an output-only hook here. I reproduced this at fc2a112 with two equal buffers where one is returned directly: after old.replace_all_uses_with(new), delete_constant_placeholder(ep, old) fails because the hook has already renamed the old input spec. The same pass works on the base commit. This also looks like the cause of the T5 failures in FuseEqualPlaceholdersPass in the current ARM CI run.

There’s another issue when call() works on a deep copy. GraphModule.__deepcopy__ copies the hook, but its closure still points at the original signature. Rewriting the copy changes the original signature, and the returned graph keeps the hook after the context exits. I reproduced this by copying an x + x program and replacing its output with x * x: the copy validates, but the original no longer does. Both validate on the base commit.

Could we add tests for both cases and keep the hook scoped to its own graph and signature? All 21 pass-infra tests passed locally, so these cases aren’t caught by the current suite.