[BC-BREAKING] Replace Type dispatch with ATenDispatch #21320
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
pytorchbot
added
module: autograd
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
li-roy
changed the title
added 3 commits
June 4, 2019 01:14
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Contributor
|
Do you want initial review now or wait until it's not WIP? |
li-roy
changed the title
Contributor
Author
|
@ezyang sorry this is ready for review, i forgot to change the title. |
Contributor
|
cc @ailzhang What's the plan for coordinating this diff with XLA? |
Contributor
|
UBSAN seems to have caught an error: |
ezyang
reviewed
Jun 5, 2019
ezyang
reviewed
Jun 5, 2019
ezyang
reviewed
Jun 5, 2019
| type_registry[static_cast<int>(b)] = std::move(t); | ||
| detail::getVariableHooks().registerVariableTypeFor(this, b); | ||
| } | ||
| private: |
Contributor
Author
There was a problem hiding this comment.
For now, I'm calling this directly from Functions.h for all factory functions.
ezyang
reviewed
Jun 5, 2019
ezyang
reviewed
Jun 5, 2019
| } | ||
| } | ||
|
|
||
| // example |
Contributor
There was a problem hiding this comment.
What's blocking removing Type.h entirely?
Contributor
Author
There was a problem hiding this comment.
A few functions like unsafeTensorFromTH are still on Type and haven't found a new home yet.
ezyang
reviewed
Jun 5, 2019
| } | ||
|
|
||
| Tensor & VariableType::detach_(Tensor & self) const { | ||
| Tensor & VariableType::detach_(Tensor & (*_op)(Tensor &), Tensor & self) { |
Contributor
There was a problem hiding this comment.
This is pretty weird, especially since you never actually call _op here. Are you sure this is the pattern you want for "Variable-only" operations?
ezyang
reviewed
Jun 5, 2019
ezyang
reviewed
Jun 5, 2019
| template<class FuncType> | ||
| ATenDispatch& registerOp(Backend backend, const char* schema, FuncType* fn) { | ||
| auto id = getSchemaId(schema); | ||
| function_table[static_cast<int64_t>(backend)][id] = reinterpret_cast<void*>(fn); |
Contributor
There was a problem hiding this comment.
You don't need a reinterpret_cast to cast a pointer into a void pointer. (It's only needed the other way.)
Contributor
There was a problem hiding this comment.
So... no error checking, I guess?
ezyang
reviewed
Jun 5, 2019
ezyang
reviewed
Jun 5, 2019
|
|
||
| if (is_variable) { | ||
| if (wrapper_table[id] == nullptr) { | ||
| AT_ERROR("No autograd wrapper is registered for ", name, ". Please report a bug to PyTorch."); |
Contributor
There was a problem hiding this comment.
If this is really a bug you should use TORCH_INTERNAL_ASSERT
ezyang
reviewed
Jun 5, 2019
| int64_t getSchemaId(std::string schema) { | ||
| static std::unordered_map<std::string, int64_t> schema_to_id = { | ||
| ${schema_to_id_pairs} | ||
| }; |
Contributor
There was a problem hiding this comment.
@zdevito Do you know if a static variable like this takes a long time to compile? Also even though this isn't perf sensitive code, we should really avoid using std::unordered_map
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
ezyang
reviewed
Jun 14, 2019
aten/src/ATen/core/ATenDispatch.h
Outdated
| class CAFFE2_API ATenOpTable { | ||
| public: | ||
| ATenOpTable(std::string schema) | ||
| : schema_(schema) {} |
ezyang
reviewed
Jun 14, 2019
| } | ||
| private: | ||
| void registerOp(Backend backend, void* fn) { | ||
| TORCH_CHECK(function_table_[static_cast<int64_t>(backend)] == nullptr, |
Contributor
There was a problem hiding this comment.
Internal asserts here? (I guess if you're going to expose directly to extensions an internal assert is not appropriate, but for internal use, this isn't a public API right)
Contributor
There was a problem hiding this comment.
There's no type checking in the registration API right now; as long as we don't have any type checking I have a pretty strong preference of NOT allowing external parties to poke this API. I don't know what the situation is with XLA and this diff at the moment. (The reason for this preference is that if XLA uses this directly, and we start changing the types of functions, they'll start segfaulting. Ick!)
ezyang
reviewed
Jun 14, 2019
| void* getBaseOp(Backend backend) const { | ||
| if (function_table_[static_cast<int64_t>(backend)] == nullptr) { | ||
| TORCH_CHECK(function_table_[static_cast<int64_t>(Backend::Undefined)] != nullptr, | ||
| "No function is registered for schema ", schema_, " on backend ", toString(backend)); |
Contributor
There was a problem hiding this comment.
Re error message: not only is no function not registered, but there is no undefined implementation registered either.
Actually, why are we dispatching to Undefined at all, at this point? Can't we just error out straight up at this point?
ezyang
reviewed
Jun 14, 2019
|
|
||
| namespace at { | ||
|
|
||
| ATenDispatch & globalATenDispatch() { |
Contributor
There was a problem hiding this comment.
Registering to the global table is not thread safe, right? Gotta be careful: library loads can happen in different threads.
ezyang
approved these changes
Jun 14, 2019
Contributor
ezyang
left a comment
ezyang
left a comment
There was a problem hiding this comment.
I'm not going to block landing this on account of lack of type checking, but if XLA is going to use this interface I would be veeeeery careful. But I don't know what the XLA plan is at the moment.
Contributor
|
The current path to XLA is that there is an include file generated, which provides a matching between C++ signature, and the magic-string pytorch generates for table based registrations. |
added 4 commits
June 18, 2019 09:52
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
This was referenced Jun 19, 2019
Closed
added 3 commits
June 19, 2019 08:11
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
Replace Type dispatch with ATenDispatch gh-metadata: pytorch pytorch 21320 gh/li-roy/25/head
zdevito
pushed a commit
to zdevito/ATen
that referenced
this pull request
Jun 20, 2019
Summary: Pull Request resolved: pytorch/pytorch#21320 ghimport-source-id: cc18f746a1c74df858cb0f6d8b7d4de4315683c7 Test Plan: Imported from OSS Differential Revision: D15637222 Pulled By: li-roy fbshipit-source-id: fcfaea0b5480ab966175341cce92e3aa0be7e3cb
Contributor
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.
Stack from ghstack:
Replace Type object based dispatch with a table of void*. This is bc breaking because it will break cpp extensions that extend Type.
This gives us a registration system that has a similar API to c10 dispatch. Backend extensions can be folded into use this API directly rather than having its own dispatch system. This also saves us a dispatch because Type dispatch was doing a dynamic dispatch and a vtable dispatch, while this table only does a dynamic dispatch.
As a follow up to this change, we can clean up a lot of code involving Type.
Outline of this change:
Benchmarks are a bit noisy because we're timing such a small op, but I ran benchmarks on a null op, with and without variable unwrapping.
Before, without variable unwrapping: ~365ns
After, without variable unwrapping: ~330ns
Before, with variable unwrapping: ~450ns
After, with variable unwrapping: ~420ns
Differential Revision: D15637222