Resolve externalized call sites by op target, not by recorded node name

[metascroy]

Problem

add_exported_program(_externalized_exported_programs=...) is documented to emit
composite graphs "for the patched call sites in exported_program", but it
registers each lowering under _ExternalizedExportedProgram.source_nodes — FX
node names captured back when _subexport_and_restore ran.

The two-step API introduced in #53 exists precisely so a caller can do work
between the phases; _patch_model_for_externalization's own docstring shows
ep = my_export_or_quantize_pipeline(model). Any pass in that window that
rebuilds or renames nodes leaves the recorded names matching nothing. The
lowerings are then registered under dead keys, and conversion fails later with
an opaque error that names neither the submodule nor the cause:

ValueError: unable to handle call function op:
  target: norm_rmsnorm_impl_7b5f5a2a.default, namespace: coreai_torch_ext

The call sites are still there — only their names changed. The custom op target
survives any such transform.

Fix

TorchConverter._resolve_source_nodes locates each call site by op target and
pairs it with its _ExternalizedExportedProgram, falling back to the recorded
names when it cannot. _perform_externalization uses the resolved names.

• One pass per program, bucketing nodes by op, so cost stays linear in total
  graph size rather than one graph walk per op name.
• Every program is searched, not just the whole-model one: a nested submodule's
  call site lives in its parent's program.
• Multiple call sites of one op are paired in graph order, which is preparation
  order. This assumes a transform preserves the relative order of an op's call
  sites, as renaming and the usual lowering passes do.
• If an op's call-site count no longer matches its prepared submodules, the
  graph changed shape rather than just its names, so pairing by position would
  be meaningless. That op keeps its recorded names and a UserWarning explains
  why, instead of silently degrading to the error above.

_utils._externalized_op_name is the inverse of the existing
_find_custom_op_node / _find_all_custom_op_nodes lookups and shares
_EXTERNALIZE_NAMESPACE with them.

Tests

• test_call_sites_resolved_after_a_renaming_transform — renames every
  call_function node between _subexport_and_restore and
  add_exported_program, then asserts both call sites still lower. Two call
  sites, so it covers ordered pairing as well as resolution. Fails without the
  change.
• test_mismatched_call_site_count_warns_and_falls_back — asserts the
  UserWarning on a count mismatch.

Notes

No behaviour change for add_pytorch_module: it converts the same program the
submodules were prepared from, so resolution returns the recorded names.

[metascroy]

cc @cymbalrush

[gokulkrishna98]

Hi @metascroy, thanks for the PR.

Any pass in that window that rebuilds or renames nodes leaves the recorded names matching nothing.

We would like to understand the use-case or need for renaming the custom ops introduced by coreai-torch (we do so in separate custom op namespace).

These APIs were introduced recently and are experimental, we would like to understand how these are getting used.

[metascroy]

Hi @metascroy, thanks for the PR.

Any pass in that window that rebuilds or renames nodes leaves the recorded names matching nothing.

We would like to understand the use-case or need for renaming the custom ops introduced by coreai-torch (we do so in separate custom op namespace).

These APIs were introduced recently and are experimental, we would like to understand how these are getting used.

The concrete consumer is the ExecuTorch Core AI backend (see sister PR pytorch/executorch#21750). A minimal version of the flow is:

ep, externalized = externalize_modules(
    model,
    [ExternalizeSpec(target_class=RMSNormImpl)],
    export_fn=lambda m: torch.export.export(m, inputs).run_decompositions(
        get_decomp_table()
    ),
)

to_edge_transform_and_lower(
    ep,
    partitioner=[CoreAIPartitioner(externalized_modules=externalized)],
)

ExecuTorch’s standard ATen-to-edge conversion rebuilds the FX graph. For example, I observed:

before: node.name = n0_rmsnorm_impl_<suffix>
after:  node.name = coreai_torch_ext_n0_rmsnorm_impl_<suffix>_default

The target remains coreai_torch_ext::n0_rmsnorm_impl_.default; ExecuTorch is not renaming the custom op or changing its namespace. The problem is only that the prepared externalization record refers to the old FX node name.

Maybe this isn't the best way to integrate externalization? Any suggestions you have would be appreciated

[gokulkrishna98]

Hi @metascroy, this PR looks good. I will refactor some of the logic for externalization, so we can remove the need for maintaining this mapping. I will be using this test as a reference (please feel free to update it or add more, so I do not break your use case).

[gokulkrishna98]

For merging the PR, you must have all the commits to have verified signatures.

reference: https://docs.github.com/en/authentication/managing-commit-signature-verification/about-commit-signature-verification