Add MKLDNN support for Half GELU (pytorch#145339)

CaoE · pytorchmergebot · commit 0dbff7e4be4d · 2025-01-22T15:14:51.000Z
Add MKLDNN support for Half GELU to align with BFloat16. Pull Request resolved: pytorch#145339 Approved by: https://github.com/yanbing-j, https://github.com/leslie-fang-intel, https://github.com/Skylion007
diff --git a/aten/src/ATen/native/Activation.cpp b/aten/src/ATen/native/Activation.cpp
@@ -382,7 +382,8 @@ static bool use_mkldnn(const Tensor& input) {
   return (input.is_mkldnn()) || // input is mkldnn Tensor
     (input.device().is_cpu() &&
     (((input.scalar_type() == kBFloat16) && mkldnn_bf16_device_check()) ||
-    (input.scalar_type() == kFloat))); // input is dense layout and bfloat16/float32
+    ((input.scalar_type() == kHalf) && mkldnn_fp16_device_check()) ||
+    (input.scalar_type() == kFloat))); // input is dense layout and bfloat16/float16/float32
 }
 #endif
 

Original file line number	Diff line number	Diff line change
`@@ -382,7 +382,8 @@ static bool use_mkldnn(const Tensor& input) {`
`382`	`382`	`return (input.is_mkldnn()) \|\| // input is mkldnn Tensor`
`383`	`383`	`(input.device().is_cpu() &&`
`384`	`384`	`(((input.scalar_type() == kBFloat16) && mkldnn_bf16_device_check()) \|\|`
`385`		`- (input.scalar_type() == kFloat))); // input is dense layout and bfloat16/float32`
	`385`	`+ ((input.scalar_type() == kHalf) && mkldnn_fp16_device_check()) \|\|`
	`386`	`+ (input.scalar_type() == kFloat))); // input is dense layout and bfloat16/float16/float32`
`386`	`387`	`}`
`387`	`388`	`#endif`
`388`	`389`