Cleanup numerical stability, constants, compound functions in kernels

Author: andrjohns

stan/math/opencl/kernels/device_functions/Phi.hpp

@@ -24,11 +24,11 @@ static const char* phi_device_function
             if (x < -37.5) {
               return 0;
             } else if (x < -5.0) {
-              return 0.5 * erfc(-1.0 / sqrt(2.0) * x);
+              return 0.5 * erfc(-M_SQRT1_2 * x);
             } else if (x > 8.25) {
               return 1;
             } else {
-              return 0.5 * (1.0 + erf(1.0 / sqrt(2.0) * x));
+              return 0.5 * (1.0 + erf(M_SQRT1_2 * x));
             }
           }
           // \cond

stan/math/opencl/kernels/device_functions/lbeta.hpp

@@ -95,21 +95,22 @@ static const char* lbeta_device_function
               return lgamma(x) + lgamma(y) - lgamma(x + y);
             }
             double x_over_xy = x / (x + y);
+            double log_xpy = log(x + y);
             if (x < LGAMMA_STIRLING_DIFF_USEFUL) {
               // y large, x small
               double stirling_diff
                   = lgamma_stirling_diff(y) - lgamma_stirling_diff(x + y);
               double stirling
-                  = (y - 0.5) * log1p(-x_over_xy) + x * (1 - log(x + y));
+                  = (y - 0.5) * log1p(-x_over_xy) + x * (1 - log_xpy);
               return stirling + lgamma(x) + stirling_diff;
             }
 
             // both large
             double stirling_diff = lgamma_stirling_diff(x)
                                    + lgamma_stirling_diff(y)
                                    - lgamma_stirling_diff(x + y);
-            double stirling = (x - 0.5) * log(x_over_xy) + y * log1p(-x_over_xy)
-                              + 0.5 * log(2.0 * M_PI) - 0.5 * log(y);
+            double stirling = (x - 0.5) * (log(x) - log_xpy) + y * log1p(-x_over_xy)
+                              + 0.5 * (M_LN2 + log(M_PI)) - 0.5 * log(y);
             return stirling + stirling_diff;
           }
           // \cond

stan/math/opencl/kernels/device_functions/lgamma_stirling.hpp

@@ -28,7 +28,7 @@ static const char* lgamma_stirling_device_function
            * @return Stirling's approximation to lgamma(x).
            */
           double lgamma_stirling(double x) {
-            return 0.5 * log(2.0 * M_PI) + (x - 0.5) * log(x) - x;
+            return 0.5 * (M_LN2 + log(M_PI)) + (x - 0.5) * log(x) - x;
           }
           // \cond
           ) "\n#endif\n";  // NOLINT

stan/math/opencl/kernels/device_functions/logit.hpp

@@ -49,7 +49,7 @@ static const char* logit_device_function
           * @param x argument
           * @return log odds of argument
           */
-          double logit(double x) { return log(x / (1 - x)); }
+          double logit(double x) { return log(x) - log1m(x); }
           // \cond
           ) "\n#endif\n";  // NOLINT
 // \endcond

stan/math/opencl/kernels/neg_binomial_2_log_glm_lpmf.hpp

@@ -92,9 +92,9 @@ static const char* neg_binomial_2_log_glm_kernel_code = STRINGIFY(
         double log_phi = log(phi);
         double logsumexp_theta_logphi;
         if (theta > log_phi) {
-          logsumexp_theta_logphi = theta + log1p(exp(log_phi - theta));
+          logsumexp_theta_logphi = theta + log1p_exp(log_phi - theta);
         } else {
-          logsumexp_theta_logphi = log_phi + log1p(exp(theta - log_phi));
+          logsumexp_theta_logphi = log_phi + log1p_exp(theta - log_phi);
         }
         double y_plus_phi = y + phi;
         if (need_logp1) {

stan/math/opencl/kernels/ordered_logistic_glm_lpmf.hpp

@@ -87,20 +87,10 @@ static const char* ordered_logistic_glm_kernel_code = STRINGIFY(
 
           if (need_location_derivative || need_cuts_derivative) {
             double exp_cuts_diff = exp(cut_y2 - cut_y1);
-            if (cut2 > 0) {
-              double exp_m_cut2 = exp(-cut2);
-              d1 = exp_m_cut2 / (1 + exp_m_cut2);
-            } else {
-              d1 = 1 / (1 + exp(cut2));
-            }
+            d1 = inv_logit(-cut2);
             d1 -= exp_cuts_diff / (exp_cuts_diff - 1);
             d2 = 1 / (1 - exp_cuts_diff);
-            if (cut1 > 0) {
-              double exp_m_cut1 = exp(-cut1);
-              d2 -= exp_m_cut1 / (1 + exp_m_cut1);
-            } else {
-              d2 -= 1 / (1 + exp(cut1));
-            }
+            d2 -= inv_logit(-cut1);
 
             if (need_location_derivative) {
               location_derivative[gid] = d1 - d2;

stan/math/opencl/kernels/ordered_logistic_lpmf.hpp

@@ -83,20 +83,10 @@ static const char* ordered_logistic_kernel_code = STRINGIFY(
 
           if (need_lambda_derivative || need_cuts_derivative) {
             double exp_cuts_diff = exp(cut_y2 - cut_y1);
-            if (cut2 > 0) {
-              double exp_m_cut2 = exp(-cut2);
-              d1 = exp_m_cut2 / (1 + exp_m_cut2);
-            } else {
-              d1 = 1 / (1 + exp(cut2));
-            }
+            d1 = inv_logit(-cut2);
             d1 -= exp_cuts_diff / (exp_cuts_diff - 1);
             d2 = 1 / (1 - exp_cuts_diff);
-            if (cut1 > 0) {
-              double exp_m_cut1 = exp(-cut1);
-              d2 -= exp_m_cut1 / (1 + exp_m_cut1);
-            } else {
-              d2 -= 1 / (1 + exp(cut1));
-            }
+            d2 -= inv_logit(-cut1);
 
             if (need_lambda_derivative) {
               lambda_derivative[gid] = d1 - d2;

stan/math/opencl/kernels/tridiagonalization.hpp

@@ -84,15 +84,15 @@ static const char* tridiagonalization_householder_kernel_code = STRINGIFY(
       q = q_local[0];
       alpha = q_local[1];
       if (q != 0) {
-        double multi = sqrt(2.) / q;
+        double multi = M_SQRT2 / q;
         // normalize the Householder vector
         for (int i = lid + 1; i < P_span; i += lsize) {
           P[P_start + i] *= multi;
         }
       }
       if (gid == 0) {
         P[P_rows * (k + j + 1) + k + j]
-            = P[P_rows * (k + j) + k + j + 1] * q / sqrt(2.) + alpha;
+            = P[P_rows * (k + j) + k + j + 1] * q / M_SQRT2 + alpha;
       }
     }
     // \cond
@@ -291,7 +291,7 @@ static const char* tridiagonalization_v_step_3_kernel_code = STRINGIFY(
         v[i] -= acc * u[i];
       }
       if (gid == 0) {
-        P[P_rows * (k + j + 1) + k + j] -= *q / sqrt(2.) * u[0];
+        P[P_rows * (k + j + 1) + k + j] -= *q / M_SQRT2 * u[0];
       }
     }
     // \cond