Fix deadlock for CUDA (#4044)

It has been noticed that Tests/GPU/CNS/Exec/RT hangs with `amrex.the_arena_init_size=0 amrex.the_arena_release_threshold=0`. The issue appears to be CUDA host callback functions do not work well with cudaFree in the main host thread. Note that we don't have any CUDA API calls in the host callback function. Also note that cudaMalloc seems to work and using a single GPU stream also works. A workaround is implemented to avoid cudaFree inside an MFIter loop.
AMReX-Codes · Jul 22, 2024 · 4392b19 · 4392b19
1 parent 1232089
commit 4392b19
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 9 deletions.
diff --git a/Src/Base/AMReX_CArena.cpp b/Src/Base/AMReX_CArena.cpp
@@ -2,6 +2,7 @@
 #include <AMReX_CArena.H>
 #include <AMReX_BLassert.H>
 #include <AMReX_Gpu.H>
+#include <AMReX_MFIter.H>
 #include <AMReX_ParallelReduce.H>
 
 #ifdef AMREX_TINY_PROFILING
@@ -57,7 +58,11 @@ CArena::alloc_protected (std::size_t nbytes)
     }
 #endif
 
-    if (static_cast<Long>(m_used+nbytes) >= arena_info.release_threshold) {
+    if (static_cast<Long>(m_used+nbytes) >= arena_info.release_threshold
+#ifdef AMREX_USE_GPU
+        && (MFIter::currentDepth() == 0)
+#endif
+        ) {
         freeUnused_protected();
     }
 
@@ -393,7 +398,11 @@ CArena::hasFreeDeviceMemory (std::size_t sz)
 
         std::size_t nbytes = Arena::align(sz == 0 ? 1 : sz);
 
-        if (static_cast<Long>(m_used+nbytes) >= arena_info.release_threshold) {
+        if (static_cast<Long>(m_used+nbytes) >= arena_info.release_threshold
+#ifdef AMREX_USE_GPU
+            && (MFIter::currentDepth() == 0)
+#endif
+            ) {
             freeUnused_protected();
         }
 

diff --git a/Src/Base/AMReX_MFIter.H b/Src/Base/AMReX_MFIter.H
@@ -167,6 +167,8 @@ public:
 
     static int allowMultipleMFIters (int allow);
 
+    static int currentDepth ();
+
     void Finalize ();
 
 protected:

diff --git a/Src/Base/AMReX_MFIter.cpp b/Src/Base/AMReX_MFIter.cpp
@@ -17,6 +17,17 @@ MFIter::allowMultipleMFIters (int allow)
     return allow;
 }
 
+int
+MFIter::currentDepth ()
+{
+    int r;
+#ifdef AMREX_USE_OMP
+#pragma omp atomic read
+#endif
+    r = MFIter::depth;
+    return r;
+}
+
 MFIter::MFIter (const FabArrayBase& fabarray_,
                 unsigned char       flags_)
     :
@@ -222,13 +233,6 @@ MFIter::Finalize ()
     // mark as invalid
     currentIndex = endIndex;
 
-#ifdef AMREX_USE_OMP
-#pragma omp master
-#endif
-    {
-        depth = 0;
-    }
-
 #ifdef BL_USE_TEAM
     if ( ! (flags & NoTeamBarrier) )
         ParallelDescriptor::MyTeam().MemoryBarrier();
@@ -257,6 +261,13 @@ MFIter::Finalize ()
     if (m_fa) {
         m_fa.reset(nullptr);
     }
+
+#ifdef AMREX_USE_OMP
+#pragma omp master
+#endif
+    {
+        depth = 0;
+    }
 }
 
 void