Add structure for use of TOSA CUSTOM ops#18837
Open
robell wants to merge 7 commits intopytorch:mainfrom
Open
Conversation
- Enables operators to be registered in the partitioner - it's expected that custom dialects are registered in this way - These should be paird with an ArmBackend pass Signed-off-by: Rob Elliott <Robert.Elliott@arm.com> Change-Id: If48fd1a0031b91ef0b0d442939557c23df4bf00a
- add tosa.CUSTOM fake op registration in the dialect - for backend passes to create+use tosa.custom only within partition - register a TOSA CUSTOM node visitor for serialization - needing shape for the wrapped op, adding a decorator for tosa shape Signed-off-by: Rob Elliott <Robert.Elliott@arm.com> Change-Id: I085ffe8656ffc1edf22d70c92bfc80aa2e602694
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/18837
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 4 New Failures, 1 Cancelled Job, 3 Unrelated FailuresAs of commit a3b7ee1 with merge base 2eaa16c ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
robell
added
the
partner: arm
robell
added
release notes: arm
Contributor
There was a problem hiding this comment.
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.CUSTOMfake op plus a registry mechanism to plug in per-wrapped-op fake implementations. - Add a TOSA CUSTOM node visitor to serialize
tosa.CUSTOMinto the TOSA serializer graph, and update metadata extraction to better handle list-valuedmeta["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.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
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.
Comment on lines
+135
to
+159
| @register_fake_tosa_op( | ||
| "CUSTOM(Tensor[] inputs, str operator_name, str domain_name, int[] implementation_attrs) -> Tensor[]", | ||
| TosaSpecification.all_versions_and_profiles(), | ||
| ) | ||
| def CUSTOM( | ||
| inputs: list[torch.Tensor], | ||
| operator_name: str, | ||
| domain_name: str, | ||
| implementation_attrs: list[int], | ||
| ) -> list[torch.Tensor]: | ||
| """Fake implementation for TOSA CUSTOM op. | ||
|
|
||
| The CUSTOM op is backend-defined. The fake implementation dispatches to a | ||
| registered compiler-side fake implementation for the specific custom op. | ||
|
|
||
| """ | ||
| _ = get_context_spec() # ensure a spec context exists | ||
| if not inputs: | ||
| raise RuntimeError("tosa.CUSTOM requires at least one input tensor") | ||
| return run_registered_fake_tosa_impl( | ||
| inputs, | ||
| operator_name, | ||
| domain_name, | ||
| implementation_attrs, | ||
| ) |
There was a problem hiding this comment.
Consider adding targeted tests for the new tosa.CUSTOM fake op dispatch/registration path (e.g., successful dispatch to a registered fake impl, and the failure mode when no impl is registered). The repo already has dialect fake-op tests under backends/arm/test/misc/ (e.g. conv2d/shape_ops), so this new behavior should be covered similarly.
Signed-off-by: Rob Elliott <Robert.Elliott@arm.com> Change-Id: I49be3aaeb79c2a0b0124dc1267626e7e9098f8a5
zingo
approved these changes
Apr 13, 2026
Signed-off-by: Rob Elliott <Robert.Elliott@arm.com> Change-Id: I8f12a57d6781266c6a3cbd43b1d5310d7c4ba4e2
Contributor
There was a problem hiding this comment.
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
torchis referenced in the type annotation for_custom_partition_ops, but this module does not importtorch. Because annotations are evaluated at import time (nofrom __future__ import annotationshere), importing this module will raiseNameError: name 'torch' is not defined. Addimport 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
torchis referenced in the type annotation for_custom_partition_ops, but this module does not importtorch. Because annotations are evaluated at import time (nofrom __future__ import annotationshere), importing this module will raiseNameError: name 'torch' is not defined. Addimport 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson