New parameter amrex::init_snan

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.
AMReX-Codes · Jul 10, 2024 · 82c9a97 · 82c9a97
1 parent cca5aad
commit 82c9a97
Show file tree

Hide file tree

Showing 7 changed files with 135 additions and 15 deletions.
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,28 @@ 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
+{
+    if (this->dptr && this->truesize > 0) {
+#ifdef AMREX_USE_GPU
+        if ((run_on == RunOn::Device) && Gpu::inLaunchRegion()) {
+            auto* p = this->dptr;
+            amrex::ParallelFor(this->truesize, [=] AMREX_GPU_DEVICE (Long i) noexcept
+            {
+                p[i] = std::numeric_limits<T>::signaling_NaN();
+            });
+        } else
+#endif
+        {
+            amrex::fill_snan(this->dptr, this->truesize);
+        }
+    }
+}
+
 template <class T>
 template <RunOn run_on>
 void
@@ -1924,6 +1955,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
@@ -14,4 +14,9 @@ extern "C" {
     void  amrex_array_init_snan (amrex_real* p, size_t nelems);
 }
 
+namespace amrex {
+    void fill_snan (float* p, size_t nelems);
+    void fill_snan (double* p, size_t nelems);
+}
+
 #endif
diff --git a/Src/Base/AMReX_MemPool.cpp b/Src/Base/AMReX_MemPool.cpp
@@ -114,26 +114,33 @@ 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
+    amrex::fill_snan(p, nelems);
+}
 
-#else
+}
 
+void amrex::fill_snan (float* p, size_t nelems)
+{
 #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));
     }
+#else
+    amrex::ignore_unused(p, nelems);
 #endif
+}
 
+void amrex::fill_snan (double* p, size_t nelems)
+{
+#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));
+    }
+#else
+    amrex::ignore_unused(p, nelems);
 #endif
 }
-}
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,38 @@ 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::ParallelFor(count, [=] AMREX_GPU_DEVICE (Size i) noexcept {
+                           data[i] = std::numeric_limits<T>::signaling_NaN();
+                       });
+                       Gpu::streamSynchronize();
+                       return;
+                   } else if constexpr (IsPolymorphicArenaAllocator<Allocator<T>>::value) {
+                       if (allocator.arena()->isManaged() ||
+                           allocator.arena()->isDevice())
+                       {
+                           amrex::ParallelFor(count, [=] AMREX_GPU_DEVICE (Size i) noexcept
+                           {
+                               data[i] = std::numeric_limits<T>::signaling_NaN();
+                           });
+                           Gpu::streamSynchronize();
+                           return;
+                       }
+                   }
+#endif
+                   fill_snan(data, count);
+                }
+            }
+        }
     }
 
     namespace VectorGrowthStrategy
@@ -254,6 +288,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 +629,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 +775,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;
+        }
     };
 }