-
Notifications
You must be signed in to change notification settings - Fork 26.3k
[ONNX] New registration API #135403
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[ONNX] New registration API #135403
Conversation
🔗 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 FailuresAs of commit 7e739df with merge base d338499 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
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? |
|
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. |
34a8103 to
65d1c9e
Compare
titaiwangms
left a comment
There was a problem hiding this 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?
This comment was marked as resolved.
This comment was marked as resolved.
titaiwangms
left a comment
There was a problem hiding this 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.
|
@titaiwangms I added tests, PTAL. I will create a separate PR for better docs if that works with you |
titaiwangms
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks!
|
@pytorchbot merge |
Merge startedYour 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 |
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
The ONNX custom ops registration API.
Design
- 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.
Support for functions that handles complex inputs will be in separate PRs.
fixes #138391
cc @gramalingam @xadupre @borisfom