-
Notifications
You must be signed in to change notification settings - Fork 26.3k
Add aten mkldnn transpose #21943
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
Add aten mkldnn transpose #21943
Conversation
|
@bddppq, please help review, thanks! |
bddppq
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 for working on this!
I would like to add the inplace version aten::mkldnn_transpose_ and throw an error to indicate it's not supported
| ideep::direct_copy::compute<AllocForMKLDNN>(x, y); | ||
| y.reshape({inferred_size.cbegin(), inferred_size.cend()}); | ||
| ideep::tensor y{x}; | ||
| y.reshape<AllocForMKLDNN>({inferred_size.cbegin(), inferred_size.cend()}); |
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.
Hmm if new shapes are as same as the old shapes, ideep::tensor::reshape will not create a copy and thus y will share the same memory as x. This could be messy since then there will be two aten opaque tensors sharing the same underlying ideep tensor.
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.
@bddppq We wanted to keep the semantics of reshape as much as possible here, i.e. do shallow copy on plain-format MKL-DNN tensor (layout of which is contiguous) and do deep copy otherwise (see https://github.com/intel/ideep/blob/d7304e0345c3f0647cb020a71682d680b9f0424a/include/ideep/tensor.hpp#L1100). For the shallow copy case, ideep::tensor y{x} would make y copy the metadata of x while share its underlying storage. This is the same semantics as the shallow copy of a native CPU tensor. So this should work fine here.
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.
@jgong5 It's fine to follow the semantics of cpu tensor to return the input tensor as output if the shapes are the same, but it's not fine to let two different aten opaque tensors to hold the same underlying ideep tensor.
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.
@bddppq But y and x are two different ideep::tensor object, right? They just share the same underlying storage while hold their own metadata. It is similar to how two CPU tensors share the same underlying storage while have their own metatdata. Make sense?
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.
@bddppq @XiaobingSuper OK. Sounds like we should call y.set_descriptor(x.get_descriptor()) explicitly after ideep::tensor y{x} in order to guarantee y uses a new MKL-DNN descriptor while sharing the underlying buffer of x. We should have wrapped an explicit shallow copy API in ideep to do such a thing actually.
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.
@jgong5 @XiaobingSuper I would say just add the following lines at the beginning of this function :-)
if (self.sizes() == inferred_size) {
return self;
}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.
Yes, I will change it. Thanks!
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.
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.
@bddppq y.reshape would create a new descriptor by calling set_descriptor() so additional y.set_descriptor(x.get_descriptor()) is not necessary here. With the recent change of @XiaobingSuper , the behavior of mkldnn_reshape looks like below, mimicking the semantics of reshape of native CPU tensor:
- If the requested shape is the same as
self, returnselfdirectly. - If the format of MKL-DNN tensor is public (or plain), the storage is contiguous, the returned MKL-DNN tensor shares the storage of
selfand copies all metadata ofselfwith the new shape. - For other situations, return a deep-copied tensor of
selfwith the new shape.
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.
@jgong5 @XiaobingSuper This sounds good to me.
bddppq
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.
LGTM. Thanks!
facebook-github-bot
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.
@bddppq has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
Summary: This PR is about: 1. Make mkldnn reshape can share same memory fro plain format tensor. 2. Add mkldnn transpose operator. Pull Request resolved: pytorch/pytorch#21943 Differential Revision: D15916063 Pulled By: bddppq fbshipit-source-id: d1971c67341f277c1e80c1fa34e213b6c27f4062
| variants: function, method | ||
| device_guard: False | ||
|
|
||
| - func: mkldnn_transpose(Tensor self, int dim0, int dim1) -> Tensor |
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.
Why is this part of the public API? If you want to have helper functions prepend an underscore to their name!
|
@XiaobingSuper @Jianhui-Li @jgong5 @uyongw @gujinghui @mingfeima Can you make sure this method is made private by prefixing an In general, the cc: @bddppq for future review. |
Summary: This PR is to make mkldnn reshape and transpose not exposes as Tensor API, please see the comments in pytorch/pytorch#21943. Pull Request resolved: pytorch/pytorch#22193 Differential Revision: D15983434 Pulled By: bddppq fbshipit-source-id: ad3514dfd8a3b0d89442eef752864e5d3f3d04f0
|
@soumith fixed in 22193. We don't intend to make mkldnn_* function exposed as public. |
This PR is about:
Make mkldnn reshape can share same memory fro plain format tensor.
Add mkldnn transpose operator.