New parameter amrex::init_snan (AMReX-Codes#4030)

Uninitialized FArrayBox, BaseFab<double|float>, PODVectors<double|float>, Gpu::DeviceVector<double|float>, etc. will be initialized to signaling NaNs, if this parameter is true. The default value for the new parameter is false unless AMREX_DEBUG or AMREX_TESTING is true. Note that previously this was only performed for FArrayBox, and it was controlled by fab.init_snan.
ajnonaka · Jul 11, 2024 · dcb9cc0 · dcb9cc0
1 parent cca5aad
commit dcb9cc0
Show file tree

Hide file tree

Showing 8 changed files with 137 additions and 26 deletions.
diff --git a/.github/workflows/gcc.yml b/.github/workflows/gcc.yml
@@ -641,7 +641,7 @@ jobs:
     - name: Run tests
       run: |
         cd build
-        ctest --output-on-failure -R
+        ctest --output-on-failure
 
   test_hdf5:
     name: [email protected] HDF5 I/O Test [tests]

diff --git a/Src/Base/AMReX.H b/Src/Base/AMReX.H
@@ -66,6 +66,8 @@ namespace amrex
 
         extern AMREX_EXPORT ErrorHandler error_handler;
         extern AMREX_EXPORT int abort_on_unused_inputs;
+
+        extern AMREX_EXPORT bool init_snan;
     }
 
     /** the AMReX "git describe" version */
@@ -199,6 +201,9 @@ namespace amrex
     [[nodiscard]] int Verbose () noexcept;
     void SetVerbose (int v) noexcept;
 
+    [[nodiscard]] bool InitSNaN () noexcept;
+    void SetInitSNaN (bool v) noexcept;
+
     // ! Get the entire command line including the executable
     [[nodiscard]] std::string get_command ();
 

diff --git a/Src/Base/AMReX.cpp b/Src/Base/AMReX.cpp
@@ -113,6 +113,11 @@ namespace system
     std::ostream* osout = &std::cout;
     std::ostream* oserr = &std::cerr;
     ErrorHandler error_handler = nullptr;
+#if defined(AMREX_DEBUG) || defined(AMREX_TESTING)
+    bool init_snan = true;
+#else
+    bool init_snan = false;
+#endif
 }
 }
 
@@ -156,6 +161,10 @@ int amrex::Verbose () noexcept { return amrex::system::verbose; }
 
 void amrex::SetVerbose (int v) noexcept { amrex::system::verbose = v; }
 
+bool amrex::InitSNaN () noexcept { return amrex::system::init_snan; }
+
+void amrex::SetInitSNaN (bool v) noexcept  { amrex::system::init_snan = v; }
+
 void amrex::SetErrorHandler (amrex::ErrorHandler f) {
     amrex::system::error_handler = f;
 }
@@ -444,6 +453,7 @@ amrex::Initialize (int& argc, char**& argv, bool build_parm_parse,
         ParmParse pp("amrex");
         pp.queryAdd("v", system::verbose);
         pp.queryAdd("verbose", system::verbose);
+        pp.queryAdd("init_snan", system::init_snan);
     }
 
     if (system::verbose > 0) {

diff --git a/Src/Base/AMReX_BaseFab.H b/Src/Base/AMReX_BaseFab.H
@@ -457,6 +457,15 @@ public:
     void getVal (T* data, const IntVect& pos, int N, int numcomp) const noexcept;
     //! Same as above, except that starts at component 0 and copies all comps.
     void getVal (T* data, const IntVect& pos) const noexcept;
+
+#if defined(AMREX_USE_GPU)
+    template <RunOn run_on,
+#else
+    template <RunOn run_on=RunOn::Host,
+#endif
+              class U=T, std::enable_if_t<std::is_same_v<U,float> || std::is_same_v<U,double>,int> FOO = 0>
+    void fill_snan () noexcept;
+
     /**
     * \brief The setVal functions set sub-regions in the BaseFab to a
     * constant value.  This most general form specifies the sub-box,
@@ -1829,6 +1838,15 @@ BaseFab<T>::shiftHalf (int idir, int n_cell) noexcept
     return *this;
 }
 
+template <class T>
+template <RunOn run_on, class U,
+          std::enable_if_t<std::is_same_v<U,float> || std::is_same_v<U,double>, int> FOO>
+void
+BaseFab<T>::fill_snan () noexcept
+{
+    amrex::fill_snan<run_on>(this->dptr, this->truesize);
+}
+
 template <class T>
 template <RunOn run_on>
 void
@@ -1924,6 +1942,20 @@ BaseFab<T>::define ()
     placementNew(this->dptr, this->truesize);
 
     amrex::update_fab_stats(this->domain.numPts(), this->truesize, sizeof(T));
+
+    if constexpr (std::is_same_v<T,float> || std::is_same_v<T,double>) {
+        if (amrex::InitSNaN() && this->truesize > 0) {
+#ifdef AMREX_USE_GPU
+            if (Gpu::inLaunchRegion() && arena()->isDeviceAccessible()) {
+                this->template fill_snan<RunOn::Device>();
+                Gpu::streamSynchronize();
+            } else
+#endif
+            {
+                this->template fill_snan<RunOn::Host>();
+            }
+        }
+    }
 }
 
 template <class T>

diff --git a/Src/Base/AMReX_FArrayBox.cpp b/Src/Base/AMReX_FArrayBox.cpp
@@ -144,6 +144,9 @@ FArrayBox::FArrayBox (const Box& b, int ncomp, Real const* p) noexcept
 void
 FArrayBox::initVal () noexcept
 {
+    // If amrex::InitSNaN is true, snans have been filled by BaseFab.
+    if (amrex::InitSNaN()) { return; }
+
     Real * p = dataPtr();
     Long s = size();
     if (p && s > 0) {

diff --git a/Src/Base/AMReX_MemPool.H b/Src/Base/AMReX_MemPool.H
@@ -2,8 +2,12 @@
 #define BL_MEMPOOL_H
 #include <AMReX_Config.H>
 
+#include <AMReX_GpuLaunch.H>
 #include <AMReX_REAL.H>
 
+#include <limits>
+#include <type_traits>
+
 extern "C" {
     void  amrex_mempool_init ();
     void  amrex_mempool_finalize ();
@@ -14,4 +18,40 @@ extern "C" {
     void  amrex_array_init_snan (amrex_real* p, size_t nelems);
 }
 
+namespace amrex {
+    template <RunOn run_on, typename T,
+              std::enable_if_t<std::is_same_v<T,double> || std::is_same_v<T,float>, int> FOO = 0>
+    void fill_snan (T* p, std::size_t nelems)
+    {
+        if (p == nullptr || nelems == 0) { return; }
+#ifdef AMREX_USE_GPU
+        if (Gpu::inLaunchRegion() && run_on == RunOn::Device) {
+            amrex::ParallelFor(nelems, [=] AMREX_GPU_DEVICE (Long i) noexcept
+            {
+                p[i] = std::numeric_limits<T>::signaling_NaN();
+            });
+        } else
+#endif
+        {
+            if constexpr (std::is_same_v<T,float>) {
+#ifdef UINT32_MAX
+                const uint32_t snan = UINT32_C(0x7fa00000);
+                static_assert(sizeof(float) == sizeof(uint32_t), "MemPool: sizeof float != sizeof uint32_t");
+                for (size_t i = 0; i < nelems; ++i) {
+                    std::memcpy(p++, &snan, sizeof(float));
+                }
+#endif
+            } else if constexpr (std::is_same_v<T,double>) {
+#ifdef UINT64_MAX
+                const uint64_t snan = UINT64_C(0x7ff0000080000001);
+                static_assert(sizeof(double) == sizeof(uint64_t), "MemPool: sizeof double != sizeof uint64_t");
+                for (size_t i = 0; i < nelems; ++i) {
+                    std::memcpy(p++, &snan, sizeof(double));
+                }
+#endif
+            }
+        }
+    }
+}
+
 #endif
diff --git a/Src/Base/AMReX_MemPool.cpp b/Src/Base/AMReX_MemPool.cpp
@@ -114,26 +114,7 @@ void amrex_real_array_init (Real* p, size_t nelems)
 
 void amrex_array_init_snan (Real* p, size_t nelems)
 {
-#ifdef BL_USE_DOUBLE
-
-#ifdef UINT64_MAX
-    const uint64_t snan = UINT64_C(0x7ff0000080000001);
-    static_assert(sizeof(double) == sizeof(uint64_t), "MemPool: sizeof double != sizeof uint64_t");
-    for (size_t i = 0; i < nelems; ++i) {
-        std::memcpy(p++, &snan, sizeof(double));
-    }
-#endif
-
-#else
-
-#ifdef UINT32_MAX
-    const uint32_t snan = UINT32_C(0x7fa00000);
-    static_assert(sizeof(float) == sizeof(uint32_t), "MemPool: sizeof float != sizeof uint32_t");
-    for (size_t i = 0; i < nelems; ++i) {
-        std::memcpy(p++, &snan, sizeof(float));
-    }
-#endif
-
-#endif
+    amrex::fill_snan<RunOn::Host>(p, nelems);
 }
+
 }
diff --git a/Src/Base/AMReX_PODVector.H b/Src/Base/AMReX_PODVector.H
@@ -2,10 +2,12 @@
 #define AMREX_PODVECTOR_H_
 #include <AMReX_Config.H>
 
+#include <AMReX.H>
 #include <AMReX_Arena.H>
 #include <AMReX_GpuLaunch.H>
 #include <AMReX_GpuAllocators.H>
 #include <AMReX_GpuDevice.H>
+#include <AMReX_MemPool.H>
 #include <AMReX_TypeTraits.H>
 
 #include <iterator>
@@ -196,6 +198,33 @@ namespace amrex
 #endif
             std::memmove(dst, src, count);
         }
+
+        template <typename T, typename Size, template<class> class Allocator>
+        void maybe_init_snan (T* data, Size count, Allocator<T> const& allocator)
+        {
+            amrex::ignore_unused(data, count, allocator);
+            if constexpr (std::is_same_v<float,  std::remove_cv_t<T>> ||
+                          std::is_same_v<double, std::remove_cv_t<T>>) {
+                if (amrex::InitSNaN()) {
+#ifdef AMREX_USE_GPU
+                   if constexpr (RunOnGpu<Allocator<T>>::value) {
+                       amrex::fill_snan<RunOn::Device>(data, count);
+                       Gpu::streamSynchronize();
+                       return;
+                   } else if constexpr (IsPolymorphicArenaAllocator<Allocator<T>>::value) {
+                       if (allocator.arena()->isManaged() ||
+                           allocator.arena()->isDevice())
+                       {
+                           amrex::fill_snan<RunOn::Device>(data, count);
+                           Gpu::streamSynchronize();
+                           return;
+                       }
+                   }
+#endif
+                   amrex::fill_snan<RunOn::Host>(data, count);
+                }
+            }
+        }
     }
 
     namespace VectorGrowthStrategy
@@ -254,6 +283,7 @@ namespace amrex
         {
             if (a_size != 0) {
                 m_data = allocate(m_size);
+                detail::maybe_init_snan(m_data, m_size, (Allocator const&)(*this));
             }
         }
 
@@ -594,16 +624,18 @@ namespace amrex
 
         void resize (size_type a_new_size)
         {
-            if (m_capacity < a_new_size) {
-                reserve(a_new_size);
+            auto old_size = m_size;
+            resize_without_init_snan(a_new_size);
+            if (old_size < a_new_size) {
+                detail::maybe_init_snan(m_data + old_size,
+                                        m_size - old_size, (Allocator const&)(*this));
             }
-            m_size = a_new_size;
         }
 
         void resize (size_type a_new_size, const T& a_val)
         {
             size_type old_size = m_size;
-            resize(a_new_size);
+            resize_without_init_snan(a_new_size);
             if (old_size < a_new_size)
             {
                 detail::uninitializedFillNImpl(m_data + old_size,
@@ -738,6 +770,14 @@ namespace amrex
             m_size = new_size;
             m_capacity = new_capacity;
         }
+
+        void resize_without_init_snan (size_type a_new_size)
+        {
+            if (m_capacity < a_new_size) {
+                reserve(a_new_size);
+            }
+            m_size = a_new_size;
+        }
     };
 }