Add a miniature backend select implementation for prims #82311
Closed
Conversation
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
It turns out that for factory function prims (prims with no Tensor arguments), we were always going to the ATen implementation of the operator. Prior to the next PR in this stack, the change is a bit hard to test, but you can indirectly observe the impact by running arange with trace dispatching on (well, you need #82277 patched in too.) ``` $ TORCH_SHOW_DISPATCH_TRACE=1 python -c "import torch._refs; torch._refs.arange(4, device='meta')" [callBoxed] op=[prims::arange], key=[BackendSelect] [call] op=[aten::empty_strided], key=[BackendSelect] [redispatch] op=[aten::empty_strided], key=[Meta] ``` Previously, the prims::arange call was dispatching to Undefined. For maximum fidelity, technically we're supposed to redispatch to a specific dispatch key, but the Python bindings to do this don't exist and it was easy to route to the implementations which we already intended to go to. We would have to fix this if we wanted external backends to register custom implementations to OTHER dispatch keys via Python op registration. Signed-off-by: Edward Z. Yang <ezyang@fb.com> [ghstack-poisoned]
🔗 Helpful links
✅ No Failures (0 Pending)As of commit 59d859c (more details on the Dr. CI page): Expand to see more💚 💚 Looks good so far! There are no failures yet. 💚 💚 This comment was automatically generated by Dr. CI (expand for details).Please report bugs/suggestions to the (internal) Dr. CI Users group. |
It turns out that for factory function prims (prims with no Tensor arguments), we were always going to the ATen implementation of the operator. Prior to the next PR in this stack, the change is a bit hard to test, but you can indirectly observe the impact by running arange with trace dispatching on (well, you need #82277 patched in too.) ``` $ TORCH_SHOW_DISPATCH_TRACE=1 python -c "import torch._refs; torch._refs.arange(4, device='meta')" [callBoxed] op=[prims::arange], key=[BackendSelect] [call] op=[aten::empty_strided], key=[BackendSelect] [redispatch] op=[aten::empty_strided], key=[Meta] ``` Previously, the prims::arange call was dispatching to Undefined. For maximum fidelity, technically we're supposed to redispatch to a specific dispatch key, but the Python bindings to do this don't exist and it was easy to route to the implementations which we already intended to go to. We would have to fix this if we wanted external backends to register custom implementations to OTHER dispatch keys via Python op registration. Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]
This was referenced Jul 27, 2022
ngimel
approved these changes
Jul 27, 2022
bdhirsh
reviewed
Jul 27, 2022
|
|
||
| from torch._subclasses.fake_tensor import contains_tensor_types | ||
|
|
||
| if not any(contains_tensor_types(a.type) for a in _prim._schema.arguments): |
There was a problem hiding this comment.
Should this map 1-to-1 with the logic that codegen uses to determine when to generate backend select kernels?
Line 859 in 1231194
There was a problem hiding this comment.
Hypothetically yes, but prims are simpler than native_functions.yaml so I guessed a simplified version would work.
bdhirsh
approved these changes
Jul 27, 2022
|
@pytorchbot merge |
|
@pytorchbot successfully started a merge job. Check the current status here |
|
Hey @ezyang. |
facebook-github-bot
pushed a commit
that referenced
this pull request
Jul 28, 2022
) Summary: It turns out that for factory function prims (prims with no Tensor arguments), we were always going to the ATen implementation of the operator. Prior to the next PR in this stack, the change is a bit hard to test, but you can indirectly observe the impact by running arange with trace dispatching on (well, you need #82277 patched in too.) ``` $ TORCH_SHOW_DISPATCH_TRACE=1 python -c "import torch._refs; torch._refs.arange(4, device='meta')" [callBoxed] op=[prims::arange], key=[BackendSelect] [call] op=[aten::empty_strided], key=[BackendSelect] [redispatch] op=[aten::empty_strided], key=[Meta] ``` Previously, the prims::arange call was dispatching to Undefined. For maximum fidelity, technically we're supposed to redispatch to a specific dispatch key, but the Python bindings to do this don't exist and it was easy to route to the implementations which we already intended to go to. We would have to fix this if we wanted external backends to register custom implementations to OTHER dispatch keys via Python op registration. Signed-off-by: Edward Z. Yang <ezyang@fb.com> Pull Request resolved: #82311 Approved by: https://github.com/ngimel, https://github.com/bdhirsh Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/3b6b27e9d7ad3951b23364f31c19b1d3ebbeaf7c Reviewed By: osalpekar Differential Revision: D38234285 Pulled By: ezyang fbshipit-source-id: e73b770529d3e221dd85f8c7b1607da975d1f211
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
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.
Stack from ghstack (oldest at bottom):
It turns out that for factory function prims (prims with no Tensor
arguments), we were always going to the ATen implementation of
the operator.
Prior to the next PR in this stack, the change is a bit hard to
test, but you can indirectly observe the impact by running arange
with trace dispatching on (well, you need
#82277 patched in too.)
Previously, the prims::arange call was dispatching to Undefined.
For maximum fidelity, technically we're supposed to redispatch to a
specific dispatch key, but the Python bindings to do this don't exist
and it was easy to route to the implementations which we already
intended to go to. We would have to fix this if we wanted external
backends to register custom implementations to OTHER dispatch keys
via Python op registration.
Signed-off-by: Edward Z. Yang ezyang@fb.com