Vectorize unary operator erfinv (pytorch#26629)

Summary:
Resolves pytorch#19088 for erfinv.

erfinv speedup (MKL, AMD Ryzen Threadripper 2970WX 24-Core Processor): 22x
Pull Request resolved: pytorch#26629

Differential Revision: D17527230

Pulled By: ezyang

fbshipit-source-id: 0a5a53a88f7eb219617120383a454a01ad78279a

Author: pearu

aten/src/ATen/core/TensorMethods.h

@@ -5120,7 +5120,13 @@ inline Tensor Tensor::polygamma(int64_t n) const {
 }
 inline Tensor Tensor::erfinv() const {
 #ifdef USE_STATIC_DISPATCH
-    return TypeDefault::erfinv(const_cast<Tensor&>(*this));
+    switch(tensorTypeIdToBackend(impl::dispatchTypeId(type_set()))) {
+        case Backend::CPU:
+            return CPUType::erfinv(const_cast<Tensor&>(*this));
+            break;
+        default:
+            AT_ERROR("erfinv not implemented for ", at::toString(type_set()));
+    }
 #else
     static c10::OperatorHandle op = c10::Dispatcher::singleton().findSchema({"aten::erfinv", ""}).value();
     return c10::Dispatcher::singleton().callUnboxed<Tensor, const Tensor &>(
@@ -5129,7 +5135,13 @@ inline Tensor Tensor::erfinv() const {
 }
 inline Tensor & Tensor::erfinv_() const {
 #ifdef USE_STATIC_DISPATCH
-    return TypeDefault::erfinv_(const_cast<Tensor&>(*this));
+    switch(tensorTypeIdToBackend(impl::dispatchTypeId(type_set()))) {
+        case Backend::CPU:
+            return CPUType::erfinv_(const_cast<Tensor&>(*this));
+            break;
+        default:
+            AT_ERROR("erfinv_ not implemented for ", at::toString(type_set()));
+    }
 #else
     static c10::OperatorHandle op = c10::Dispatcher::singleton().findSchema({"aten::erfinv_", ""}).value();
     return c10::Dispatcher::singleton().callUnboxedOnly<Tensor &, Tensor &>(

aten/src/ATen/cpu/vec256/vec256_base.h

@@ -8,6 +8,7 @@
 
 #include <ATen/Utils.h>
 #include <ATen/native/Copy.h>
+#include <ATen/native/Math.h>
 #include <ATen/NumericUtils.h>
 #include <c10/util/C++17.h>
 #include <c10/util/BFloat16.h>
@@ -197,6 +198,9 @@ struct Vec256 {
   Vec256<T> erfc() const {
     return map(std::erfc);
   }
+  Vec256<T> erfinv() const {
+    return map(calc_erfinv);
+  }
   Vec256<T> exp() const {
     return map(std::exp);
   }

aten/src/ATen/cpu/vec256/vec256_double.h

@@ -109,6 +109,9 @@ template <> class Vec256<double> {
   Vec256<double> erfc() const {
     return Vec256<double>(Sleef_erfcd4_u15(values));
   }
+  Vec256<double> erfinv() const {
+    return map(calc_erfinv);
+  }
   Vec256<double> exp() const {
     return Vec256<double>(Sleef_expd4_u10(values));
   }

aten/src/ATen/cpu/vec256/vec256_float.h

@@ -117,6 +117,9 @@ template <> class Vec256<float> {
   Vec256<float> erfc() const {
     return Vec256<float>(Sleef_erfcf8_u15(values));
   }
+  Vec256<float> erfinv() const {
+    return map(calc_erfinv);
+  }
   Vec256<float> exp() const {
     return Vec256<float>(Sleef_expf8_u10(values));
   }

aten/src/ATen/cpu/vml.h

@@ -106,6 +106,7 @@ IMPLEMENT_VML_BUG(cos)
 // IMPLEMENT_VML_BUG(cosh)
 IMPLEMENT_VML_BUG(erf)
 IMPLEMENT_VML_BUG(erfc)
+IMPLEMENT_VML(erfinv)
 IMPLEMENT_VML_BUG(exp)
 IMPLEMENT_VML_BUG(expm1)
 IMPLEMENT_VML_BUG(floor)
@@ -174,6 +175,7 @@ IMPLEMENT_VML_MKL(cos, Cos)
 // IMPLEMENT_VML_MKL(cosh, Cosh)
 IMPLEMENT_VML_MKL(erf, Erf)
 IMPLEMENT_VML_MKL(erfc, Erfc)
+IMPLEMENT_VML_MKL(erfinv, ErfInv)
 IMPLEMENT_VML_MKL(exp, Exp)
 IMPLEMENT_VML_MKL(expm1, Expm1)
 IMPLEMENT_VML_MKL(log, Ln)

aten/src/ATen/native/Math.h

@@ -1,8 +1,14 @@
+#pragma once
+
 #include <cstdlib>
 #include <cmath>
 #include <limits>
 #include <type_traits>
 
+#ifndef M_PIf
+#define M_PIf 3.1415926535f
+#endif  // M_PIf
+
 /* The next function is taken from  https://github.com/antelopeusersgroup/antelope_contrib/blob/master/lib/location/libgenloc/erfinv.c.
 Below is the copyright.
 Output was modified to be inf or -inf when input is 1 or -1. */

aten/src/ATen/native/UnaryOps.cpp

@@ -65,10 +65,6 @@ Tensor& ceil_out(Tensor& result, const Tensor& self) { return unary_op_impl_out(
 Tensor ceil(const Tensor& self) { return unary_op_impl(self, at::ceil_out); }
 Tensor& ceil_(Tensor& self) { return unary_op_impl_(self, at::ceil_out); }
 
-Tensor& erfinv_out(Tensor& result, const Tensor& self) { return unary_op_impl_out(result, self, erfinv_stub); }
-Tensor erfinv(const Tensor& self) { return unary_op_impl(self, at::erfinv_out); }
-Tensor& erfinv_(Tensor& self) { return unary_op_impl_(self, at::erfinv_out); }
-
 Tensor& floor_out(Tensor& result, const Tensor& self) { return unary_op_impl_out(result, self, floor_stub); }
 Tensor floor(const Tensor& self) { return unary_op_impl(self, at::floor_out); }
 Tensor& floor_(Tensor& self) { return unary_op_impl_(self, at::floor_out); }
@@ -281,6 +277,7 @@ IMPLEMENT_UNARY_OP_VEC(cos)
 IMPLEMENT_UNARY_OP_VEC(cosh)
 IMPLEMENT_UNARY_OP_VEC(erf)
 IMPLEMENT_UNARY_OP_VEC(erfc)
+IMPLEMENT_UNARY_OP_VEC_CUDA(erfinv)
 IMPLEMENT_UNARY_OP_VEC(exp)
 IMPLEMENT_UNARY_OP_VEC(expm1)
 IMPLEMENT_UNARY_OP_VEC(frac)

aten/src/ATen/native/cpu/UnaryOpsKernel.cpp

@@ -156,14 +156,6 @@ static void cosh_kernel(TensorIterator& iter) {
   });
 }
 
-static void erfinv_kernel(TensorIterator& iter) {
-  AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "erfinv_cpu", [&]() {
-    cpu_kernel(
-        iter,
-        [=](scalar_t a) -> scalar_t { return calc_erfinv(a); });
-  });
-}
-
 static void digamma_kernel(TensorIterator& iter) {
   AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "digamma", [&]() {
     cpu_kernel(
@@ -337,7 +329,6 @@ REGISTER_DISPATCH(neg_stub, &neg_kernel);
 REGISTER_DISPATCH(sign_stub, &sign_kernel);
 REGISTER_DISPATCH(sinh_stub, &sinh_kernel);
 REGISTER_DISPATCH(cosh_stub, &cosh_kernel);
-REGISTER_DISPATCH(erfinv_stub, &erfinv_kernel);
 REGISTER_DISPATCH(digamma_stub, &digamma_kernel);
 REGISTER_DISPATCH(trigamma_stub, &trigamma_kernel);
 REGISTER_DISPATCH(polygamma_stub, &polygamma_kernel);
@@ -355,6 +346,7 @@ IMPLEMENT_FLOAT_KERNEL(FLOATING, cos)
 // IMPLEMENT_FLOAT_KERNEL(FLOATING, cosh)
 IMPLEMENT_FLOAT_KERNEL(FLOATING, erf)
 IMPLEMENT_FLOAT_KERNEL(FLOATING, erfc)
+IMPLEMENT_FLOAT_KERNEL(FLOATING, erfinv)
 IMPLEMENT_FLOAT_KERNEL(FLOATING, exp)
 IMPLEMENT_FLOAT_KERNEL(FLOATING, expm1)
 IMPLEMENT_FLOAT_KERNEL(FLOATING, floor)

aten/src/ATen/native/native_functions.yaml

@@ -4842,17 +4842,23 @@
   use_c10_dispatcher: full
   supports_named_tensor: True
   variants: method, function
+  dispatch:
+    CPU: erfinv
+    CUDA: erfinv
 
 - func: erfinv_(Tensor(a!) self) -> Tensor(a!)
   use_c10_dispatcher: unboxed_only
   supports_named_tensor: True
   variants: method
+  dispatch:
+    CPU: _erfinv__cpu
+    CUDA: _erfinv__cuda
 
 - func: erfinv.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
   supports_named_tensor: True
   dispatch:
-    CPU: erfinv_out
-    CUDA: erfinv_out
+    CPU: _erfinv_out_cpu
+    CUDA: _erfinv_out_cuda
 
 - func: sign(Tensor self) -> Tensor
   use_c10_dispatcher: unboxed_only