forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathBinaryOps.cpp
70 lines (51 loc) · 1.7 KB
/
BinaryOps.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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
#include <ATen/ATen.h>
#include <ATen/Config.h>
#include <ATen/NativeFunctions.h>
#if !AT_MKLDNN_ENABLED()
namespace at {
namespace native {
Tensor& mkldnn_add_out(
Tensor& result,
const Tensor& self,
const Tensor& other,
Scalar alpha) {
AT_ERROR("mkldnn_add_out: ATen not compiled with MKLDNN support");
}
Tensor mkldnn_add(const Tensor& self, const Tensor& other, Scalar alpha) {
AT_ERROR("mkldnn_add: ATen not compiled with MKLDNN support");
}
Tensor& mkldnn_add_(Tensor& self, const Tensor& other, Scalar alpha) {
AT_ERROR("mkldnn_add_: ATen not compiled with MKLDNN support");
}
} // namespace native
} // namespace at
#else // AT_MKLDNN_EBABLED
#include <ATen/native/mkldnn/MKLDNNCommon.h>
namespace at {
namespace native {
Tensor& mkldnn_add_out(
Tensor& result,
const Tensor& self,
const Tensor& other,
Scalar alpha) {
ideep::tensor& x = itensor_from_mkldnn(self);
ideep::tensor& y = itensor_from_mkldnn(other);
ideep::tensor& z = itensor_from_mkldnn(result);
const std::vector<float> scales{1.0, alpha.to<float>()};
ideep::sum::compute<AllocForMKLDNN>(scales, {x, y}, z);
return result;
}
Tensor mkldnn_add(const Tensor& self, const Tensor& other, Scalar alpha) {
ideep::tensor& x = itensor_from_mkldnn(self);
ideep::tensor& y = itensor_from_mkldnn(other);
ideep::tensor z;
const std::vector<float> scales{1.0, alpha.to<float>()};
ideep::sum::compute<AllocForMKLDNN>(scales, {x, y}, z);
return new_with_itensor_mkldnn(std::move(z), self.options());
}
Tensor& mkldnn_add_(Tensor& self, const Tensor& other, Scalar alpha) {
return native::mkldnn_add_out(self, self, other, alpha);
}
} // namespace native
} // namespace at
#endif // AT_MKLDNN_EBABLED