Skip to content

Conversation

@justinchuby
Copy link
Collaborator

@justinchuby justinchuby commented Sep 7, 2024

The ONNX custom ops registration API.

Design

  1. Create a "custom_translation_table: dict[Callable, Sequence[Callable] | Callable" parameter for specifying extra functions
  2. Use a callable as the key to support all possible call_function targets in the fx graph
  3. Allow a callable or a Sequence of callables as values.
    - When there is a single callable, it is the translation function for the op
    - When there is a Sequence of callable, the exporter's dispatcher will dispatch to these callables in order based on input dtypes.
    - The translation functions can be a plain python function that calls onnxscript ops (traced), or an onnxscript function.
    - Complex input support: We create special type annotations for annotating real representations of complex inputs, which are needed to handle complex computation in the ONNX graph, as we don't have any ops in ONNX that handle complex inputs. The dispatcher will have knowledge of these newly created type annotations and dispatch correctly. The complex functions will be in the same overload pool as the real functions.
torch.onnx.export(dynamo=True, 
	custom_translation_table = { 
	torch.ops.aten.add: [overload1, overload2], 
	torch.sym_not: sym_not_onnx, 
})

Support for functions that handles complex inputs will be in separate PRs.

fixes #138391

cc @gramalingam @xadupre @borisfom

@pytorch-bot pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Sep 7, 2024
@pytorch-bot
Copy link

pytorch-bot bot commented Sep 7, 2024

🔗 Helpful Links

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

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

✅ No Failures

As of commit 7e739df with merge base d338499 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@justinchuby justinchuby added module: onnx Related to torch.onnx and removed release notes: onnx torch.onnx related changes that should show up in the release notes labels Sep 7, 2024
@pytorch-bot pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Sep 7, 2024
@titaiwangms titaiwangms requested a review from xadupre September 7, 2024 00:22
@xadupre
Copy link
Collaborator

xadupre commented Sep 7, 2024

What about having the decompositions functions as a parameter of the exporter? Maybe it could be seen as an attributes of a class OptimizationOptions handling the rewriting patterns and other postprocessings?

@justinchuby
Copy link
Collaborator Author

justinchuby commented Sep 7, 2024

Sending in option objects tends to be a little inconvenient in my experience because constructing those objects is usually a hassle (need to type a long name, figure out what fields to fill, etc). (I used to think that’s a good idea and pushed for it, but every time i needed to supply those to dynamo export or ORT I didn’t enjoy it. ) I wonder if we can simply allow a dictionary as an argument if desired.

@ezyang ezyang added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Sep 9, 2024
@titaiwangms titaiwangms added the topic: new features topic category label Sep 9, 2024
@justinchuby justinchuby force-pushed the justinchu/onnx-registration branch from 34a8103 to 65d1c9e Compare September 16, 2024 19:02
Copy link
Collaborator

@titaiwangms titaiwangms left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the status of this?

@justinchuby

This comment was marked as resolved.

@justinchuby justinchuby changed the title [ONNX][RFC] New registration API [ONNX] New registration API Oct 23, 2024
Copy link
Collaborator

@titaiwangms titaiwangms left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should have a test or docs for demonstration.

@justinchuby
Copy link
Collaborator Author

@titaiwangms I added tests, PTAL. I will create a separate PR for better docs if that works with you

Copy link
Collaborator

@titaiwangms titaiwangms left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@justinchuby
Copy link
Collaborator Author

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 1, 2024
@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

rahulsingh-intel pushed a commit to rahulsingh-intel/pytorch that referenced this pull request Nov 5, 2024
The ONNX custom ops registration API.

## Design

1. Create a "custom_translation_table: dict[Callable, Sequence[Callable] | Callable" parameter for specifying extra functions
2. Use a callable as the key to support all possible call_function targets in the fx graph
3. Allow a callable or a Sequence of callables as values.
		- When there is a single callable, it is the translation function for the op
		- When there is a Sequence of callable, the exporter's dispatcher will dispatch to these callables in order based on input dtypes.
		- The translation functions can be a plain python function that calls onnxscript ops (traced), or an onnxscript function.
		- Complex input support: We create special type annotations for annotating real representations of complex inputs, which are needed to handle complex computation in the ONNX graph, as we don't have any ops in ONNX that handle complex inputs. The dispatcher will have knowledge of these newly created type annotations and dispatch correctly. The complex functions will be in the same overload pool as the real functions.

```py
torch.onnx.export(dynamo=True,
	custom_translation_table = {
	torch.ops.aten.add: [overload1, overload2],
	torch.sym_not: sym_not_onnx,
})
```
Support for functions that handles complex inputs will be in separate PRs.

fixes pytorch#138391

Pull Request resolved: pytorch#135403
Approved by: https://github.com/titaiwangms
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk Trigger trunk jobs on your pull request Merged module: onnx Related to torch.onnx open source release notes: onnx torch.onnx related changes that should show up in the release notes topic: new features topic category triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[ONNX] Custom op registration api

7 participants