From 23a7f34fd7a4938f29761472f42d362023b00bd5 Mon Sep 17 00:00:00 2001 From: "S. Eric Clark" <25495882+clarkse@users.noreply.github.com> Date: Tue, 17 Sep 2024 17:21:30 -0700 Subject: [PATCH] Adding computation of complete elliptic integrals into amrex::Math (#4151) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit … so that it can be executed on accelerator devices. ## Summary Currently the parser recognizes comp_ellint_1 and comp_ellint_2, but only executes on CPU when compiled with gcc. An implementation of the calculations have been added to amrex::Math to compute these quantities using the method described here: https://dlmf.nist.gov/19.8 It is a quadratically converging iterative method that is compatible with device execution. Gauss's Arithmetic-Geometric Mean   --- Docs/sphinx_documentation/source/Basics.rst | 4 +- Src/Base/AMReX_Math.H | 67 +++++++++++++++++++++ Src/Base/Parser/AMReX_Parser_Y.H | 20 +++--- 3 files changed, 77 insertions(+), 14 deletions(-) diff --git a/Docs/sphinx_documentation/source/Basics.rst b/Docs/sphinx_documentation/source/Basics.rst index 97e68fd1e6..2de0afae59 100644 --- a/Docs/sphinx_documentation/source/Basics.rst +++ b/Docs/sphinx_documentation/source/Basics.rst @@ -595,8 +595,8 @@ numbers can be computed with ``min`` and ``max``, respectively. It supports the Heaviside step function, ``heaviside(x1,x2)`` that gives ``0``, ``x2``, ``1``, for ``x1 < 0``, ``x1 = 0`` and ``x1 > 0``, respectively. It supports the Bessel function of the first kind of order ``n`` -``jn(n,x)``. Complete elliptic integrals of the first and second kind, ``comp_ellint_1`` and ``comp_ellint_2``, -are supported only for gcc and CPUs. +``jn(n,x)``. Complete elliptic integrals of the first and second kind, ``comp_ellint_1(k)`` and ``comp_ellint_2(k)``, +are supported. There is ``if(a,b,c)`` that gives ``b`` or ``c`` depending on the value of ``a``. A number of comparison operators are supported, including ``<``, ``>``, ``==``, ``!=``, ``<=``, and ``>=``. The Boolean results from diff --git a/Src/Base/AMReX_Math.H b/Src/Base/AMReX_Math.H index c4d8d524af..7bc086bde1 100644 --- a/Src/Base/AMReX_Math.H +++ b/Src/Base/AMReX_Math.H @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -225,6 +226,72 @@ std::uint64_t umulhi (std::uint64_t a, std::uint64_t b) } #endif +template +AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE +T comp_ellint_1 (T k) +{ + // Computing K based on DLMF + // https://dlmf.nist.gov/19.8 + T tol = std::numeric_limits::epsilon(); + + T a0 = 1.0; + T g0 = std::sqrt(1.0 - k*k); + T a = a0; + T g = g0; + + // 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()/a; +} + +template +AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE +T comp_ellint_2 (T k) +{ + // Computing E based on DLMF + // https://dlmf.nist.gov/19.8 + T Kcomp = amrex::Math::comp_ellint_1(k); + T tol = std::numeric_limits::epsilon(); + + // Step Zero + T a0 = 1.0; + T g0 = std::sqrt(1.0 - k*k); + 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 diff --git a/Src/Base/Parser/AMReX_Parser_Y.H b/Src/Base/Parser/AMReX_Parser_Y.H index 38f561ad78..0f4f4470b0 100644 --- a/Src/Base/Parser/AMReX_Parser_Y.H +++ b/Src/Base/Parser/AMReX_Parser_Y.H @@ -349,28 +349,24 @@ AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE T parser_math_atanh (T a) { return std::atanh(a); } template -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); #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(k); #endif } template -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); #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(k); #endif }