Environment variable: AMREX_THE_ARENA_INIT_SIZE (#4786)

We can now set the default value of amrex.the_arena_init_size with an
environment variable. This is convenient for CI jobs.

---------

Co-authored-by: Alexander Sinn <[email protected]>

Author: WeiqunZhang

Docs/sphinx_documentation/source/GPU.rst

@@ -560,8 +560,16 @@ to :cpp:`The_Arena()` to reduce memory fragmentation.
 In :cpp:`amrex::Initialize`, a large amount of GPU device memory is
 allocated and is kept in :cpp:`The_Arena()`.  The default is 3/4 of the
 total device memory, and it can be changed with a :cpp:`ParmParse`
-parameter, ``amrex.the_arena_init_size``, in the unit of bytes.  The default
-initial size for other arenas is 8388608 (i.e., 8 MB).  For
+parameter, ``amrex.the_arena_init_size``, in the unit of bytes. The default
+can also be changed with an environment variable
+``AMREX_THE_ARENA_INIT_SIZE=X``, where ``X`` is the number of bytes. When
+both the :cpp:`ParmParse` parameter and the environment variable are
+present, the former will override the latter. In both cases, the number
+string could have optional single quotes ``'`` as separators (e.g.,
+``10'000'000'000``). It may also use floating-point notation (``2.5e10``),
+as long as converting it does not introduce any loss of precision.
+
+The default initial size for other arenas is 8388608 (i.e., 8 MB).  For
 :cpp:`The_Managed_Arena()` and :cpp:`The_Device_Arena()`, it can be changed
 with ``amrex.the_managed_arena_init_size`` and
 ``amrex.the_device_arena_init_size``, respectively, if they are not an alias

Src/Base/AMReX_Arena.cpp

@@ -6,6 +6,7 @@
 
 #include <AMReX.H>
 #include <AMReX_BLProfiler.H>
+#include <AMReX_IParser.H>
 #include <AMReX_Print.H>
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_ParmParse.H>
@@ -362,6 +363,13 @@ Arena::Initialize (bool minimal)
     the_pinned_arena_release_threshold = Gpu::Device::totalGlobalMem() / Gpu::Device::numDevicePartners() / 2L;
 #endif
 
+    // Overwrite the initial size with environment variables
+    if (char const* init_size_p = std::getenv("AMREX_THE_ARENA_INIT_SIZE")) {
+        IParser iparser(init_size_p);
+        auto exe = iparser.compileHost<0>();
+        the_arena_init_size = exe();
+    }
+
     ParmParse pp("amrex");
     pp.queryAdd(        "the_arena_init_size",         the_arena_init_size);
     pp.queryAdd( "the_device_arena_init_size",  the_device_arena_init_size);