Adding computation of complete elliptic integrals into amrex::Math so…

… that it can be executed on accelerator devices.

Src/Base/AMReX_Math.H

@@ -5,10 +5,12 @@
 #include <AMReX_GpuQualifiers.H>
 #include <AMReX_Extension.H>
 #include <AMReX_INT.H>
+#include <AMReX_REAL.H>
 #include <cmath>
 #include <cstdlib>
 #include <type_traits>
 #include <utility>
+#include <iostream>
 
 #ifdef AMREX_USE_SYCL
 # include <sycl/sycl.hpp>
@@ -225,6 +227,70 @@ std::uint64_t umulhi (std::uint64_t a, std::uint64_t b)
 }
 #endif
 
+template <typename T>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T comp_ellint_1 (T m)
+{
+ // Computing K based on DLMF
+ // https://dlmf.nist.gov/19.8
+ T tol = 1e-12;
+ T a0 = 1.0;
+ T g0 = std::sqrt(1.0 - m);
+ T a, g;
+
+ // Find Arithmetic Geometric mean
+ while(std::abs(a0 - g0) > tol) {
+ a = 0.5*(a0 + g0);
+ g = std::sqrt(a0 * g0);
+
+ a0 = a;
+ g0 = g;
+ }
+
+ return 0.5*pi<T>()/a;
+}
+
+template <typename T>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T comp_ellint_2 (T m)
+{
+ // Computing E based on DLMF
+ // https://dlmf.nist.gov/19.8
+ T Kcomp = amrex::Math::comp_ellint_1<T>(m);
+ T tol = 1e-12;
+
+ // Step Zero
+ T a0 = 1.0;
+ T g0 = std::sqrt(1.0 - m);
+ T cn = std::sqrt(a0*a0 - g0*g0);
+
+ // Step 1
+ int n = 1;
+ T a = 0.5 * (a0 + g0);
+ T g = std::sqrt(a0*g0);
+ cn = 0.25*cn*cn/a;
+
+ T sum_val = a*a;
+ a0 = a;
+ g0 = g;
+
+ while(std::abs(cn*cn) > tol) {
+ // Compute coefficients for this iteration
+ a = 0.5 * (a0 + g0);
+ g = std::sqrt(a0*g0);
+ cn = 0.25*cn*cn/a;
+
+ n++;
+ sum_val -= std::pow(2,n-1)*cn*cn;
+
+ // Save a and g for next iteration
+ a0 = a;
+ g0 = g;
+ }
+
+ return Kcomp*sum_val;
+}
+
 /***************************************************************************************************
  * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: BSD-3-Clause

Src/Base/Parser/AMReX_Parser_Y.H

@@ -349,28 +349,26 @@ AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
 T parser_math_atanh (T a) { return std::atanh(a); }
 
 template <typename T>
-AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
-T parser_math_comp_ellint_1 (T a)
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T parser_math_comp_ellint_1 (T k)
 {
 #if defined(__GNUC__) && !defined(__clang__) && !defined(__CUDA_ARCH__) && !defined(__NVCOMPILER)
- return std::comp_ellint_1(a);
+ // return std::comp_ellint_1(k);
+ return amrex::Math::comp_ellint_1<T>(k);
 #else
- amrex::ignore_unused(a);
- AMREX_ALWAYS_ASSERT_WITH_MESSAGE(false, "parser: comp_ellint_1 only supported with gcc and cpu");
- return 0.0;
+ return amrex::Math::comp_ellint_1<T>(k);
 #endif
 }
 
 template <typename T>
-AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
-T parser_math_comp_ellint_2 (T a)
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+T parser_math_comp_ellint_2 (T k)
 {
 #if defined(__GNUC__) && !defined(__clang__) && !defined(__CUDA_ARCH__) && !defined(__NVCOMPILER)
- return std::comp_ellint_2(a);
+ // return std::comp_ellint_2(k);
+ return amrex::Math::comp_ellint_2<T>(k);
 #else
- amrex::ignore_unused(a);
- AMREX_ALWAYS_ASSERT_WITH_MESSAGE(false, "parser: comp_ellint_2 only supported with gcc and cpu");
- return 0.0;
+ return amrex::Math::comp_ellint_2<T>(k);
 #endif
 }