[ExecuTorch][Core AI][1/x] Graph pass: int64/float64 dtype narrowing

[metascroy]

Adds NarrowToCoreAIDtypesPass, which casts int64/float64 graph inputs to 32-bit at the boundary and widens 64-bit outputs back. This lets index style inputs such as token ids be delegated, since Core AI supports only up to 32-bit dtypes. The model's external I/O dtypes are preserved. Includes unit tests.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ 3 Pending, 1 Unrelated Failure

As of commit a73ddca with merge base d827cf7 ():

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

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

• Apple / build-demo-ios / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 65

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.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 49s —— View job

---

Code Review — NarrowToCoreAIDtypesPass

Read the pass, both __init__.py files, and the tests. The core mechanism — narrow 64-bit placeholders right after the input, re-run FakeTensorProp, then widen the values feeding originally-64-bit outputs — is sound, and the tricky bits are handled correctly. A few observations below, mostly non-blocking.

Correctness — looks right ✅

• The replace_all_uses_with(cast) → cast.args = (placeholder,) dance (narrow_dtypes.py:74-75) is the correct idiom for avoiding the self-reference cycle, and the comment explains it well.
• Capturing orig_output_dtypes before narrowing and only widening when the current dtype actually differs (:83-92) is the right ordering. Non-Node outputs (constants) are correctly skipped.
• The aten-vs-edge _to_copy selection is validated by test_interior_is_narrowed_to_32bit, which confirms FakeTensorProp propagates through edge-dialect ops.

Points worth considering

1. Only placeholder 64-bit tensors are narrowed. Interior ops that intrinsically emit int64 (aten.arange, argmax, nonzero, comparison chains) or int64 get_attr constants aren't touched, so CoreAIPartitioner will still reject those nodes. This is likely intentional for a "1/x" PR, but worth a one-line note in the docstring so the scope is explicit.

2. Identity / pass-through inputs get a redundant narrow→widen pair. For def forward(x): return x with an int64 input, the input is cast to i32 and then immediately widened back to i64 — functionally correct but adds two dead-ish _to_copy nodes around the boundary. Probably fine in practice; flagging in case it matters for the CoreAI partitioner's node counting.

3. cast.meta["val"] at :76 is redundant — the very next FakeTensorProp(...).propagate(...) overwrites it. Harmless, but you could drop it to keep the narrowing loop leaner.

4. FakeTensorProp argument ordering assumption. propagate(*[p.meta["val"] for p in placeholders]) relies on placeholders being in graph order and matching every lifted input (params/buffers included). That holds for exported/edge programs, so it's correct — just noting it's an implicit contract. The current tests only exercise param-free modules; a test with a module that has lifted parameters (e.g. nn.Linear) plus an int64 index input would harden this path.

5. Not yet wired into the backend. Grepping the coreai tree shows the pass is only referenced by its own files — it isn't registered in a partitioner or get_default_passes yet. Assuming that's the subject of a later PR in the series; just confirming it's currently exercised only by unit tests.

Style / minor

• The pass is a plain callable returning PassResult rather than subclassing ExportPass/PassBase. That's a valid pass form and keeps it simple, which fits the repo guidance — no change needed, just confirming it's intentional.
• Tests are clean and cover the meaningful cases (int64 narrow+widen, interior narrowing, float64, and the no-op path). Nice.

Overall this is a well-scoped, well-tested change. Nothing blocking from my side; items 1 and 4 are the two I'd most consider addressing.
· branch executorch-coreai-passes-1/x

[metascroy]

@digantdesai @JacobSzwejbka can I get a stamp here.

It's already been reviewed by @cymbalrush