amrex::callNoinline: Call given function without inline (#4606)

This works on lambdas, functors, normal functions. But it does not work
on overloaded functions like std::sin. If needed, one could however wrap
functions like std::sin inside a lambda function. It also does not work
with normal
functions for SYCL and one would have to wrap it inside a lambda.     

Here is the motivation behind this PR. In this impactx PR
(BLAST-ImpactX/impactx#964), a GPU kernel uses 8
amrex::Parser's. The CUDA CI fails if more than one job is used in
build. Apparently the kernel is too big because all those parser
functions are inlined. This PR provides a way to reduce the size by
forcing noinline.

Author: WeiqunZhang

Src/Base/AMReX_TypeTraits.H

@@ -87,6 +87,9 @@ namespace amrex
     struct IsParticleContainer : public std::is_base_of<ParticleContainerBase, T>::type {};
 #endif
 
+    template <class T, class Enable = void>
+    struct DefinitelyNotHostRunnable : std::false_type {};
+
 #ifdef AMREX_USE_GPU
 
     template <class T, class Enable = void>
@@ -95,9 +98,6 @@ namespace amrex
     template <class T, class Enable = void>
     struct MaybeHostDeviceRunnable : std::true_type {};
 
-    template <class T, class Enable = void>
-    struct DefinitelyNotHostRunnable : std::false_type {};
-
 #if defined(AMREX_USE_CUDA) && defined(__NVCC__)
 
     template <class T>

Src/Base/AMReX_Utility.H

@@ -16,6 +16,7 @@
 #include <AMReX_GpuQualifiers.H>
 #include <AMReX_FileSystem.H>
 #include <AMReX_String.H>
+#include <AMReX_TypeTraits.H>
 
 #include <cfloat>
 #include <chrono>
@@ -242,6 +243,19 @@ namespace amrex
                          int limit = 100,
                          std::ostringstream ss = {});
 
+    /**
+     * \brief Call given function without inline.
+     *
+     * This works on lambdas, functors, normal functions. But it does not
+     * work with overloaded functions like std::sin. If needed, one could
+     * however wrap functions like std::sin inside a lambda function. It
+     * also does not work with normal functions for SYCL and one would have
+     * to wrap it inside a lambda.
+     */
+    template <typename F, typename... T>
+    AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
+    auto callNoinline (F const& f, T&&... arg)
+        -> decltype(std::declval<F>()(std::declval<T>()...)); // needed for nvcc
 }
 
 template <typename T>
@@ -506,5 +520,17 @@ std::string amrex::ToString(const T& t, const char* symbol_begin, const char* sy
     return ss.str();
 }
 
+template <typename F, typename... T>
+AMREX_GPU_HOST_DEVICE AMREX_NO_INLINE
+auto amrex::callNoinline (F const& f, T&&... arg)
+    -> decltype(std::declval<F>()(std::declval<T>()...)) // needed for nvcc
+{
+    AMREX_IF_ON_HOST((
+        if constexpr (!amrex::DefinitelyNotHostRunnable<F>::value) {
+            return f(std::forward<T>(arg)...);
+        }
+    ))
+    AMREX_IF_ON_DEVICE(( return f(std::forward<T>(arg)...); ))
+}
 
 #endif /*BL_UTILITY_H*/

Tests/CMakeLists.txt

@@ -125,7 +125,7 @@ else()
    #
    # List of subdirectories to search for CMakeLists.
    #
-   set( AMREX_TESTS_SUBDIRS Amr AsyncOut CLZ CommType CTOParFor DeviceGlobal Enum
+   set( AMREX_TESTS_SUBDIRS Amr AsyncOut CallNoinline CLZ CommType CTOParFor DeviceGlobal Enum
                             MultiBlock MultiPeriod ParmParse Parser Parser2 ParserUserFn Reinit
                             RoundoffDomain SmallMatrix)
 

Tests/CallNoinline/CMakeLists.txt

@@ -0,0 +1,9 @@
+foreach(D IN LISTS AMReX_SPACEDIM)
+    set(_sources     main.cpp)
+    set(_input_files)
+
+    setup_test(${D} _sources _input_files)
+
+    unset(_sources)
+    unset(_input_files)
+endforeach()

Tests/CallNoinline/GNUmakefile

@@ -0,0 +1,24 @@
+AMREX_HOME ?= ../../amrex
+
+DEBUG	= FALSE
+
+DIM	= 3
+
+COMP    = gcc
+
+USE_MPI   = FALSE
+USE_OMP   = FALSE
+USE_CUDA  = FALSE
+USE_HIP   = FALSE
+USE_SYCL  = FALSE
+
+BL_NO_FORT = TRUE
+
+TINY_PROFILE = FALSE
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.defs
+
+include ./Make.package
+include $(AMREX_HOME)/Src/Base/Make.package
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.rules

Tests/CallNoinline/Make.package

@@ -0,0 +1 @@
+CEXE_sources += main.cpp

Tests/CallNoinline/main.cpp

@@ -0,0 +1,94 @@
+#include <AMReX.H>
+#include <AMReX_Gpu.H>
+#include <AMReX_Utility.H>
+
+#include <cmath>
+
+using namespace amrex;
+
+AMREX_GPU_DEVICE AMREX_FORCE_INLINE
+void halve (double& x)
+{
+    x *= 0.5;
+}
+
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+double f (double x, double y)
+{
+    return std::cos(x*y);
+}
+
+struct S {
+    AMREX_GPU_DEVICE AMREX_FORCE_INLINE
+    double operator() (double a, double b, double c) const {
+        return std::cos(a*b*c);
+    }
+};
+
+struct P {
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    double operator() (double a, double b, double c) const {
+        return std::cos(a*b*c);
+    }
+};
+
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+#endif
+
+int main(int argc, char* argv[])
+{
+    amrex::Initialize(argc,argv);
+    {
+        double pi = amrex::Math::pi<double>();
+        double one = 1.0;
+
+        auto g = [] AMREX_GPU_DEVICE (double a, double b, double c)
+        {
+            return std::sin(a*b*c);
+        };
+
+        S s{};
+        P p{};
+
+        amrex::Gpu::HostVector<double> ones(2);
+        amrex::Gpu::HostVector<double> zeroes(3);
+        auto* p1 = ones.data();
+        auto* p0 = zeroes.data();
+
+        amrex::ParallelFor(1, [=] AMREX_GPU_DEVICE (int)
+        {
+            p1[0] = callNoinline([] (double a) { return std::sin(a); }, pi*one*0.5);
+            p1[1] = callNoinline(g, pi, one, 0.5);
+
+            auto half = one;
+#ifdef AMREX_USE_SYCL
+            callNoinline([] (double& x) { halve(x); }, half);
+            p0[0] = callNoinline([] (double a, double b) { return f(a,b); }, pi, half);
+#else
+            callNoinline(halve, half);
+            p0[0] = callNoinline(f, pi, half);
+#endif
+            auto half2 = one;
+            callNoinline([] (double& a) { a *= 0.5; }, half2);
+            p0[1] = callNoinline(s, pi, one, half2);
+            p0[2] = callNoinline(p, pi, one, half2);
+        });
+        Gpu::streamSynchronize();
+
+        zeroes.push_back(callNoinline(f, pi, 0.5));
+        zeroes.push_back(callNoinline(P{}, pi, one, 0.5));
+
+        amrex::Print() << "ones: " << amrex::ToString(ones) << "\n"
+                       << "zeroes: " << amrex::ToString(zeroes) << "\n";
+
+        for (auto x : ones) {
+            AMREX_ALWAYS_ASSERT(almostEqual(x, 1.0, 10));
+        }
+
+        for (auto x : zeroes) {
+            AMREX_ALWAYS_ASSERT(std::abs(x) < 1.e-15);
+        }
+    }
+    amrex::Finalize();
+}