Fix deadlock for CUDA

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 cudaMall seems work and using a single GPU stream also works. A workaround is implemented to avoid cudaFree when there are host callback functions inside an MFIter loop.
AMReX-Codes · Jul 22, 2024 · 189d7b7 · 189d7b7
1 parent cdacc4d
commit 189d7b7
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 single
+#endif
+    {
+        depth = 0;
+    }
 }
 
 void