Parser: Use malloc if Arena is not ready

This allows Parser/IParser to be used before Arena is ready. For example, amrex.the_arena_init_size=1e9 now works.
AMReX-Codes · Nov 21, 2024 · 17f3cf6 · 17f3cf6
1 parent b98eaaf
commit 17f3cf6
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 13 deletions.
diff --git a/Src/Base/Parser/AMReX_IParser.H b/Src/Base/Parser/AMReX_IParser.H
@@ -88,12 +88,13 @@ private:
         std::string m_expression;
         struct amrex_iparser* m_iparser = nullptr;
         int m_nvars = 0;
-        mutable char* m_host_executor = nullptr;
+        bool m_use_arena = true;
+        char* m_host_executor = nullptr;
 #ifdef AMREX_USE_GPU
-        mutable char* m_device_executor = nullptr;
+        char* m_device_executor = nullptr;
 #endif
-        mutable int m_max_stack_size = 0;
-        mutable int m_exe_size = 0;
+        int m_max_stack_size = 0;
+        int m_exe_size = 0;
         Data () = default;
         ~Data ();
         Data (Data const&) = delete;
@@ -129,6 +130,10 @@ IParser::compileHost () const
             }
 
             m_data->m_host_executor = (char*)The_Pinned_Arena()->alloc(m_data->m_exe_size);
+            if (m_data->m_host_executor == nullptr) { // Arena is not ready yet
+                m_data->m_host_executor = (char*) std::malloc(m_data->m_exe_size);
+                m_data->m_use_arena = false;
+            }
 
             try {
                 iparser_compile(m_data->m_iparser, m_data->m_host_executor);
@@ -155,7 +160,9 @@ IParser::compile () const
     auto exe = compileHost<N>();
 
 #ifdef AMREX_USE_GPU
-    if (m_data && m_data->m_iparser && !(m_data->m_device_executor)) {
+    if (m_data && m_data->m_iparser && !(m_data->m_device_executor)
+        && m_data->m_use_arena)
+    {
         m_data->m_device_executor = (char*)The_Arena()->alloc(m_data->m_exe_size);
         Gpu::htod_memcpy_async(m_data->m_device_executor, m_data->m_host_executor,
                                m_data->m_exe_size);

diff --git a/Src/Base/Parser/AMReX_IParser.cpp b/Src/Base/Parser/AMReX_IParser.cpp
@@ -42,7 +42,13 @@ IParser::Data::~Data ()
 {
     m_expression.clear();
     if (m_iparser) { amrex_iparser_delete(m_iparser); }
-    if (m_host_executor) { The_Pinned_Arena()->free(m_host_executor); }
+    if (m_host_executor) {
+        if (m_use_arena) {
+            The_Pinned_Arena()->free(m_host_executor);
+        } else {
+            std::free(m_host_executor);
+        }
+    }
 #ifdef AMREX_USE_GPU
     if (m_device_executor) { The_Arena()->free(m_device_executor); }
 #endif

diff --git a/Src/Base/Parser/AMReX_Parser.H b/Src/Base/Parser/AMReX_Parser.H
@@ -99,13 +99,14 @@ private:
         std::string m_expression;
         struct amrex_parser* m_parser = nullptr;
         int m_nvars = 0;
-        mutable char* m_host_executor = nullptr;
+        bool m_use_arena = true;
+        char* m_host_executor = nullptr;
 #ifdef AMREX_USE_GPU
-        mutable char* m_device_executor = nullptr;
+        char* m_device_executor = nullptr;
 #endif
-        mutable int m_max_stack_size = 0;
-        mutable int m_exe_size = 0;
-        mutable Vector<char const*> m_locals;
+        int m_max_stack_size = 0;
+        int m_exe_size = 0;
+        Vector<char const*> m_locals;
         Data () = default;
         ~Data ();
         Data (Data const&) = delete;
@@ -142,6 +143,10 @@ Parser::compileHost () const
             }
 
             m_data->m_host_executor = (char*)The_Pinned_Arena()->alloc(m_data->m_exe_size);
+            if (m_data->m_host_executor == nullptr) { // Arena is not ready yet
+                m_data->m_host_executor = (char*) std::malloc(m_data->m_exe_size);
+                m_data->m_use_arena = false;
+            }
 
             try {
                 m_data->m_locals = parser_compile(m_data->m_parser,
@@ -169,7 +174,9 @@ Parser::compile () const
     auto exe = compileHost<N>();
 
 #ifdef AMREX_USE_GPU
-    if (m_data && m_data->m_parser && !(m_data->m_device_executor)) {
+    if (m_data && m_data->m_parser && !(m_data->m_device_executor)
+        && m_data->m_use_arena)
+    {
         m_data->m_device_executor = (char*)The_Arena()->alloc(m_data->m_exe_size);
         Gpu::htod_memcpy_async(m_data->m_device_executor, m_data->m_host_executor,
                                m_data->m_exe_size);

diff --git a/Src/Base/Parser/AMReX_Parser.cpp b/Src/Base/Parser/AMReX_Parser.cpp
@@ -42,7 +42,13 @@ Parser::Data::~Data ()
 {
     m_expression.clear();
     if (m_parser) { amrex_parser_delete(m_parser); }
-    if (m_host_executor) { The_Pinned_Arena()->free(m_host_executor); }
+    if (m_host_executor) {
+        if (m_use_arena) {
+            The_Pinned_Arena()->free(m_host_executor);
+        } else {
+            std::free(m_host_executor);
+        }
+    }
 #ifdef AMREX_USE_GPU
     if (m_device_executor) { The_Arena()->free(m_device_executor); }
 #endif