add an API for external backends to register custom device names #86992
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/86992
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 FailuresAs of commit 2fad0f2: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
… names"
This API adds some improvements to external backends who are building C++ backends out of tree using the `PrivateUse1` dispatch key.
The docs and linked examples go over the API in more detail, but you should be able to use it like:
```
# This should probably be in the __init__.py file of a external backend's python package
> torch.register_privateuse1_backend("foo")`
# And it will allow the user to do this:
> a = torch.ones(2, device="foo")
```
[ghstack-poisoned]
… names"
This API adds some improvements to external backends who are building C++ backends out of tree using the `PrivateUse1` dispatch key.
The docs and linked examples go over the API in more detail, but you should be able to use it like:
```
# This should probably be in the __init__.py file of a external backend's python package
> torch.register_privateuse1_backend("foo")`
# And it will allow the user to do this:
> a = torch.ones(2, device="foo")
```
[ghstack-poisoned]
torch/_torch_docs.py
Outdated
| ) | ||
|
|
||
| add_docstr( | ||
| torch.register_privateuse1_backend, |
There was a problem hiding this comment.
Not sure about the name here.
In particular, keeping in the back of my mind that we might need to extend this to support multiple backends.
I am wondering if we want to make an API that will support the multi backend case directly here. And thus naming it something like torch.register_new_backend("foo"). And then from c++ recommend they use getDispatchKeyForDeviceType(parse_type("foo")) (not sure what is the exact name for this?) to register their kernels.
The other option is to say they need to use PrivateUse1 manually and thus this API is just to rename it and much more limited. In such a case, I would suggest we move it to a smaller namespace like torch.utils.rename_privateuse1_backend("foo").
Neither of these options require changing the way things are implemented right now btw, just the python naming, the doc and the error message when multiple backends are being used.
WDYT?
There was a problem hiding this comment.
And then from c++ recommend they use
getDispatchKeyForDeviceType(parse_type("foo"))
I don't think this will easily work without some bigger API changes, since right now on the C++ side you still need to hardcode the PrivateUse1 key name directly into our registration macro: TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {...}.
If we wanted to fully go down the multi-backend route now (and abstract PrivateUse1 away from any public API's), one way would be with a new macro like TORCH_BACKEND_IMPL(aten, foo, m {...}) where we set up the link between a PrivateUse key and the backend name at library load time. That would make a python API like torch.register_new_backend() redundant though.
LMK if you agree with that line of reasoning - but I'm leaning towards your second proposal, which is to use a smaller namespace (torch.utils.rename_privateuse1_backend)
… names"
This API adds some improvements to external backends who are building C++ backends out of tree using the `PrivateUse1` dispatch key.
The docs and linked examples go over the API in more detail, but you should be able to use it like:
```
# This should probably be in the __init__.py file of a external backend's python package
> torch.register_privateuse1_backend("foo")`
# And it will allow the user to do this:
> a = torch.ones(2, device="foo")
```
[ghstack-poisoned]
… names"
This API adds some improvements to external backends who are building C++ backends out of tree using the `PrivateUse1` dispatch key.
The docs and linked examples go over the API in more detail, but you should be able to use it like:
```
# This should probably be in the __init__.py file of a external backend's python package
> torch.register_privateuse1_backend("foo")`
# And it will allow the user to do this:
> a = torch.ones(2, device="foo")
```
[ghstack-poisoned]
… names"
This API adds some improvements to external backends who are building C++ backends out of tree using the `PrivateUse1` dispatch key.
The docs and linked examples go over the API in more detail, but you should be able to use it like:
```
# This should probably be in the __init__.py file of a external backend's python package
> torch.register_privateuse1_backend("foo")`
# And it will allow the user to do this:
> a = torch.ones(2, device="foo")
```
[ghstack-poisoned]
… names"
This API adds some improvements to external backends who are building C++ backends out of tree using the `PrivateUse1` dispatch key.
The docs and linked examples go over the API in more detail, but you should be able to use it like:
```
# This should probably be in the __init__.py file of a external backend's python package
> torch.register_privateuse1_backend("foo")`
# And it will allow the user to do this:
> a = torch.ones(2, device="foo")
```
[ghstack-poisoned]
… names"
This API adds some improvements to external backends who are building C++ backends out of tree using the `PrivateUse1` dispatch key.
The docs and linked examples go over the API in more detail, but you should be able to use it like:
```
# This should probably be in the __init__.py file of a external backend's python package
> torch.register_privateuse1_backend("foo")`
# And it will allow the user to do this:
> a = torch.ones(2, device="foo")
```
[ghstack-poisoned]
|
@pytorchbot merge -f "flaky failure" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
Hey @bdhirsh. |
|
@pytorchbot revert -m "breaking internal builds - D40534212 - arstudio-windows-tests-landcastle-0" -c "ghfirst" |
|
Seems it is breaking windows builds: check toupper/tolower usage |
|
@pytorchbot successfully started a revert job. Check the current status here. |
|
@bdhirsh your PR has been successfully reverted. |
…mes (#86992)" This reverts commit fb6826b. Reverted #86992 on behalf of https://github.com/jeanschmidt due to breaking internal builds - D40534212 - arstudio-windows-tests-landcastle-0
github-actions
bot
requested review from
albanD,
anjali411,
antoniojkim,
Chillee,
ezyang,
Krovatkin,
miladm and
wconstab
|
Closing, reland at #87453 |
…mes (#86992)" (#87453) Re-land of #86992 This reverts commit a895af9. Pull Request resolved: #87453 Approved by: https://github.com/ezyang, https://github.com/albanD
…mes (pytorch#86992)" (pytorch#87453) Re-land of pytorch#86992 This reverts commit a895af9. Pull Request resolved: pytorch#87453 Approved by: https://github.com/ezyang, https://github.com/albanD
…mes (pytorch#86992)" (pytorch#87453) Re-land of pytorch#86992 This reverts commit a895af9. Pull Request resolved: pytorch#87453 Approved by: https://github.com/ezyang, https://github.com/albanD
This API adds some improvements to external backends who are building C++ backends out of tree using the
PrivateUse1dispatch key.The docs and linked examples go over the API in more detail, but you should be able to use it like:
Stack from ghstack (oldest at bottom):