[pytorch] bfloat16 support in erfinv (#111257)

jspark1105 · facebook-github-bot · commit 29ea5cdab75e · 2023-10-24T10:59:29.000-07:00
Summary: Pull Request resolved: #111257 As title Test Plan: CI Reviewed By: minsii, jianyuh Differential Revision: D50280766 fbshipit-source-id: ed3e18f86ab07d1a3ae47729ccf513d130775d4f
diff --git a/aten/src/ATen/native/cuda/UnarySpecialOpsKernel.cu b/aten/src/ATen/native/cuda/UnarySpecialOpsKernel.cu
@@ -281,18 +281,19 @@ void erfc_kernel_cuda(TensorIteratorBase& iter) {
 CONSTEXPR_EXCEPT_WIN_CUDA char erfinv_name[] = "erfinv_kernel";
 void erfinv_kernel_cuda(TensorIteratorBase& iter) {
   #if AT_USE_JITERATOR()
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.common_dtype(), "erfinv_cuda", [&]() {
+    AT_DISPATCH_FLOATING_TYPES_AND2(ScalarType::Half, ScalarType::BFloat16, iter.common_dtype(), "erfinv_cuda", [&]() {
       jitted_gpu_kernel</*name=*/erfinv_name,
                         /*return_dtype=*/ scalar_t,
                         /*common_dtype=*/ scalar_t,
                         /*arity=*/ 1>(iter, erfinv_string);
       });
   #else
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.common_dtype(), "erfinv_cuda", [&]() {
-      gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> scalar_t {
-        return ::erfinv(a);
-      });
-    });
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(ScalarType::Half, ScalarType::BFloat16,
+        iter.common_dtype(), "erfinv_cuda", [&]() {
+          gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> scalar_t {
+            return ::erfinv(a);
+          });
+        });
   #endif
 }
 
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -17351,7 +17351,7 @@ def reference_flatten(input, start_dim=0, end_dim=-1):
                                                   torch.bfloat16: 1e-2,
                                                   torch.float32: 1e-4}),),
                    dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16),
-                   dtypesIfCUDA=all_types_and(torch.bool, torch.half),
+                   dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16),
                    supports_sparse_csr=True,
                    supports_sparse_csc=True,
                    supports_sparse_bsr=True,
