Add structure for use of TOSA CUSTOM ops

[robell]

TOSA custom ops allow for wrapped custom operators of their own dialect. This adds the tosa op and recognition in partitioning to enable custom torch.library operators to be passed into the arm backend and mapped to backend provided implementations.

The broader implementation using these mechanisms will follow.

• Enables operators to be registered in the partitioner
• add tosa.CUSTOM fake op registration in the dialect
• for backend passes (in tree or registed) to create+use tosa.custom only within partition
• register a TOSA CUSTOM node visitor for serialization
• as custom wraps an operator, adds support for register_fake_tosa to register the shape-only operator in tosa

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson

[pytorch-bot]

🔗 Helpful Links

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

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

• Rolling out OSDC (ARC) runners on pull workflow for PyTorch trunk commits

❌ 4 New Failures, 1 Cancelled Job, 3 Unrelated Failures

As of commit a3b7ee1 with merge base 2eaa16c ():

NEW FAILURES - The following jobs have failed:

• Cadence Build & Test / cpu-test / test-aot / test-aot (gh)
  backends/cadence/aot/tests/test_replace_ops_passes.py::TestReplaceOpsPasses::test_replace_transposed_conv_with_linear_1
• trunk / test-torchao-huggingface-checkpoints (lfm2_5_1_2b, linux.arm64.2xlarge, executorch-ubuntu-22.04-g... / linux-job (gh)
  RuntimeError: Command docker exec -t 67a121197cb4abc82ad42761b0635196b1cd629b15b221d8b707bb853ff8e747 /exec failed with exit code 1
• trunk / test-torchao-huggingface-checkpoints (phi_4_mini, linux.arm64.2xlarge, executorch-ubuntu-22.04-gc... / linux-job (gh)
  RuntimeError: Command docker exec -t a37770dd1b7cd648101284a377d66ad0b98d865f3dec03ceb29359055a0ec9e0 /exec failed with exit code 1
• trunk / test-torchao-huggingface-checkpoints (qwen3_4b, linux.arm64.2xlarge, executorch-ubuntu-22.04-gcc1... / linux-job (gh)
  RuntimeError: Command docker exec -t e896010047c2057026886f51ab2808cf1810dc342b9ffae491c37e3674619f1c /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• trunk / test-models-macos-coreml (mv3) / macos-job (gh)
  ##[error]The operation was canceled.

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

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

• pull / unittest / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• trunk / unittest-release / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[Copilot]

Pull request overview

Adds initial infrastructure for representing and plumbing backend-defined custom operators through the Arm TOSA pipeline, including partitioning support, fake-op propagation, and serialization.

Changes:

• Extend the TOSA partitioner to treat explicitly registered custom ops as supported and to prevent their decomposition.
• Add a generic tosa.CUSTOM fake op plus a registry mechanism to plug in per-wrapped-op fake implementations.
• Add a TOSA CUSTOM node visitor to serialize tosa.CUSTOM into the TOSA serializer graph, and update metadata extraction to better handle list-valued meta["val"].

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
backends/arm/tosa/partitioner.py	Adds custom-op registration and integrates custom-op support into partitioning and decomposition policy.
backends/arm/tosa/mapping.py	Extends extract_tensor_meta to handle meta["val"] as a non-empty list (multi-output cases).
backends/arm/tosa/dialect/ops/custom.py	Introduces tosa.CUSTOM fake op plus registry/dispatch for wrapped custom-op fake implementations.
backends/arm/tosa/dialect/init.py	Ensures the new custom dialect ops module is imported/registered.
backends/arm/operators/op_tosa_custom.py	Adds serializer lowering for tosa.CUSTOM.default via a new node visitor.
backends/arm/operators/init.py	Registers the new custom node visitor module for import side effects.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (2)

backends/arm/ethosu/partitioner.py:37

• torch is referenced in the type annotation for _custom_partition_ops, but this module does not import torch. Because annotations are evaluated at import time (no from __future__ import annotations here), importing this module will raise NameError: name 'torch' is not defined. Add import torch (or quote the annotation) to fix the import-time failure.

        )
        self.additional_checks = additional_checks
        self.tosa_spec = compile_spec.tosa_spec
        self._custom_partition_ops: set[torch._ops.OpOverload] = set()

backends/arm/vgf/partitioner.py:37

• torch is referenced in the type annotation for _custom_partition_ops, but this module does not import torch. Because annotations are evaluated at import time (no from __future__ import annotations here), importing this module will raise NameError: name 'torch' is not defined. Add import torch (or quote the annotation) to fix the import-time failure.

        )
        self.additional_checks = additional_checks
        self.tosa_spec = compile_spec.tosa_spec
        self._custom_partition_ops: set[torch._ops.OpOverload] = set()

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.