1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
2 #include <ATen/core/Tensor.h>
3 #include <ATen/Config.h>
4
5 #ifndef AT_PER_OPERATOR_HEADERS
6 #include <ATen/NativeFunctions.h>
7 #else
8 #include <ATen/ops/relu_native.h> // for mkldnn_relu, mkldnn_...
9 #include <ATen/ops/threshold_backward_native.h> // for mkldnn_relu_backward
10 #endif
11
12 #if !AT_MKLDNN_ENABLED()
13
14 namespace at { namespace native {
15
mkldnn_relu(const Tensor & input)16 Tensor mkldnn_relu(const Tensor& input) {
17 TORCH_CHECK(false, "mkldnn_relu: ATen not compiled with MKLDNN support");
18 }
19
mkldnn_relu_(Tensor & input)20 Tensor& mkldnn_relu_(Tensor& input) {
21 TORCH_CHECK(false, "mkldnn_relu_: ATen not compiled with MKLDNN support");
22 }
23
mkldnn_relu_backward(const Tensor & grad_output,const Tensor & input,const Scalar & threshold)24 Tensor mkldnn_relu_backward(const Tensor& grad_output, const Tensor& input, const Scalar& threshold) {
25 TORCH_CHECK(false, "mkldnn_relu_backward: ATen not compiled with MKLDNN support");
26 }
27
28 }}
29
30 #else // AT_MKLDNN_ENABLED
31
32 #include <ATen/native/mkldnn/MKLDNNCommon.h>
33 #include <ATen/native/mkldnn/Utils.h>
34
35 namespace at { namespace native {
36
mkldnn_relu(const Tensor & input)37 Tensor mkldnn_relu(const Tensor& input) {
38 if (input.scalar_type() == ScalarType::BFloat16) {
39 TORCH_CHECK(mkldnn_bf16_device_check(),
40 "mkldnn_relu: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq");
41 }
42
43 const ideep::tensor& x = itensor_from_mkldnn(input);
44 ideep::tensor y;
45 ideep::eltwise_forward::compute(
46 x, y, ideep::algorithm::eltwise_relu, ideep::prop_kind::forward_training, /*alpha*/ 0.0);
47 return new_with_itensor_mkldnn(std::move(y), optTypeMetaToScalarType(input.options().dtype_opt()),
48 input.options().device_opt());
49 }
50
mkldnn_relu_(Tensor & input)51 Tensor& mkldnn_relu_(Tensor& input) {
52 if (input.scalar_type() == ScalarType::BFloat16) {
53 TORCH_CHECK(mkldnn_bf16_device_check(),
54 "mkldnn_relu_: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq");
55 }
56
57 ideep::tensor& x = itensor_from_mkldnn(input);
58 ideep::eltwise_forward::compute(
59 x, x, ideep::algorithm::eltwise_relu, ideep::prop_kind::forward_training, /*alpha*/ 0.0);
60 return input;
61 }
62
mkldnn_relu_backward(const Tensor & grad_output,const Tensor & input,const Scalar & threshold)63 Tensor mkldnn_relu_backward(const Tensor& grad_output, const Tensor& input, const Scalar& threshold) {
64 ideep::tensor& x = itensor_from_mkldnn(input);
65 ideep::tensor grady = itensor_from_mkldnn(grad_output);
66 ideep::tensor gradx;
67 ideep::eltwise_backward::compute(x, grady, gradx,
68 ideep::algorithm::eltwise_relu, /*alpha*/ 0.0);
69 return new_with_itensor_mkldnn(std::move(gradx),
70 optTypeMetaToScalarType(grad_output.options().dtype_opt()),
71 grad_output.options().device_opt());
72 }
73
74 }}
75
76 #endif // AT_MKLDNN_ENABLED
77