forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy.cpp
43 lines (33 loc) · 1.05 KB
/
Copy.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
#include <ATen/ATen.h>
#include <ATen/Config.h>
#include <ATen/NativeFunctions.h>
#if !AT_MKLDNN_ENABLED()
namespace at {
namespace native {
Tensor& copy_mkldnn_(Tensor& self, const Tensor& src, bool non_blocking) {
TORCH_CHECK(false, "copy_mkldnn_: ATen not compiled with MKLDNN support");
}
} // namespace native
} // namespace at
#else // AT_MKLDNN_ENABLED
#include <ATen/native/mkldnn/MKLDNNCommon.h>
namespace at {
namespace native {
Tensor& copy_mkldnn_(Tensor& self, const Tensor& src, bool non_blocking) {
TORCH_CHECK(
self.sizes() == src.sizes(),
"copy_mkldnn_: only support same size tensor.");
TORCH_CHECK(
self.is_mkldnn() && src.is_mkldnn(),
"copy_mkldnn_: between mkldnn layout and dense Tensors is not implemented! Found self type = ",
self.toString(),
" and src type = ",
src.toString());
ideep::tensor& x = itensor_from_mkldnn(src);
ideep::tensor& y = itensor_from_mkldnn(self);
ideep::direct_copy::compute(x, y);
return self;
}
} // namespace native
} // namespace at
#endif // AT_MKLDNN_ENABLED