xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/mkldnn/Gelu.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
2 #include <ATen/core/Tensor.h>
3 #include <ATen/Config.h>
4 #include <ATen/native/Activation.h>
5 
6 #ifndef AT_PER_OPERATOR_HEADERS
7 #include <ATen/NativeFunctions.h>
8 #else
9 #include <ATen/ops/gelu_native.h>
10 #include <ATen/ops/gelu_backward_native.h>
11 #endif
12 
13 #if !AT_MKLDNN_ENABLED()
14 
15 namespace at { namespace native {
16 
mkldnn_gelu(const Tensor & input,c10::string_view approximate)17 Tensor mkldnn_gelu(const Tensor& input, c10::string_view approximate) {
18   TORCH_CHECK(false, "mkldnn_gelu: ATen not compiled with MKLDNN support");
19 }
20 
mkldnn_gelu_backward(const Tensor & grad_output,const Tensor & input,c10::string_view approximate)21 Tensor mkldnn_gelu_backward(const Tensor& grad_output, const Tensor& input, c10::string_view approximate) {
22   TORCH_CHECK(false, "mkldnn_gelu_backward: ATen not compiled with MKLDNN support");
23 }
24 
25 }}
26 
27 #else // AT_MKLDNN_ENABLED
28 
29 #include <ATen/native/mkldnn/MKLDNNCommon.h>
30 #include <ATen/native/mkldnn/Utils.h>
31 
32 namespace at { namespace native {
33 
mkldnn_gelu(const Tensor & input,c10::string_view approximate)34 Tensor mkldnn_gelu(const Tensor& input, c10::string_view approximate) {
35   if (input.scalar_type() == ScalarType::BFloat16) {
36     TORCH_CHECK(mkldnn_bf16_device_check(),
37         "mkldnn_gelu: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq");
38   }
39   TORCH_CHECK(get_gelutype_enum(approximate) == GeluType::None,
40                   "mkldnn_gelu: fast, approximate gelu is not supported");
41   const ideep::tensor& x = itensor_from_tensor(input);
42   ideep::tensor y;
43   ideep::eltwise_forward::compute(
44       x, y, ideep::algorithm::eltwise_gelu_erf, ideep::prop_kind::forward_training, /*alpha*/ 0.0);
45   return new_with_itensor_mkldnn(std::move(y), optTypeMetaToScalarType(input.options().dtype_opt()),
46                                  input.options().device_opt());
47 }
48 
mkldnn_gelu_backward(const Tensor & grad_output,const Tensor & input,c10::string_view approximate)49 Tensor mkldnn_gelu_backward(const Tensor& grad_output, const Tensor& input, c10::string_view approximate) {
50   TORCH_CHECK(get_gelutype_enum(approximate) == GeluType::None,
51                   "mkldnn_gelu_backward: fast, approximate gelu is not supported");
52   const ideep::tensor& x = itensor_from_tensor(input);
53   ideep::tensor grady = itensor_from_tensor(grad_output);
54   ideep::tensor gradx;
55   ideep::eltwise_backward::compute(x, grady, gradx,
56       ideep::algorithm::eltwise_gelu_erf, /*alpha*/ 0.0);
57   return new_with_itensor_mkldnn(std::move(gradx),
58                                  optTypeMetaToScalarType(grad_output.options().dtype_opt()),
59                                  grad_output.options().device_opt());
60 }
61 
62 }}
63 
64 #endif // AT_MKLDNN_ENABLED
65