Improve Kokkos initialization

diff --git a/include/deal.II/base/kokkos.h b/include/deal.II/base/kokkos.h
index caccb9446621d9d30d0a463cc2ee273d0c8e8b01..6c020fd8251f481eb09ca91f26571c3796fdfbb8 100644 (file)
--- a/include/deal.II/base/kokkos.h
+++ b/include/deal.II/base/kokkos.h
@@ -22,9 +22,18 @@ DEAL_II_NAMESPACE_OPEN
 
 namespace Impl
 {
+  /**
+   * Records if Kokkos has been initialized by deal.II. The value stored is only
+   * meaningful after ensure_kokkos_initialized() has been called.
+   */
+  extern bool dealii_initialized_kokkos;
+
+  /**
+   * Makes sure that Kokkos is initialized. Sets dealii_initialized_kokkos.
+   */
   void
   ensure_kokkos_initialized();
-}
+} // namespace Impl
 
 DEAL_II_NAMESPACE_CLOSE
 

diff --git a/include/deal.II/lac/vector_memory.templates.h b/include/deal.II/lac/vector_memory.templates.h
index e6c13b649c52100a376457ec47423d1ccb1edf58..48db94c9b1a1e55d83ac42bec0aca324dd2e84de 100644 (file)
--- a/include/deal.II/lac/vector_memory.templates.h
+++ b/include/deal.II/lac/vector_memory.templates.h
@@ -19,6 +19,7 @@
 
 #include <deal.II/base/config.h>
 
+#include <deal.II/base/kokkos.h>
 #include <deal.II/base/memory_consumption.h>
 
 #include <deal.II/lac/vector_memory.h>
@@ -32,7 +33,21 @@ template <typename VectorType>
 typename GrowingVectorMemory<VectorType>::Pool &
 GrowingVectorMemory<VectorType>::get_pool()
 {
-  static GrowingVectorMemory<VectorType>::Pool pool;
+  // Kokkos needs to be initialized before constructing the static Pool for
+  // vector types that use Kokkos.
+  // If Kokkos is initialized by deal.II, this make sure that it is finalized
+  // after the Pool has been destroyed.
+  // If Kokkos is not initialized by deal.II, we assume that Kokkos is not
+  // finalized past program end together with static variables and we need to
+  // make sure to empty the Pool when finalizing Kokkos so that the destruction
+  // of the Pool doesn't call Kokkos functions.
+  Impl::ensure_kokkos_initialized();
+  static auto pool = []() {
+    if (!Impl::dealii_initialized_kokkos)
+      Kokkos::push_finalize_hook(
+        GrowingVectorMemory<VectorType>::release_unused_memory);
+    return GrowingVectorMemory<VectorType>::Pool{};
+  }();
   return pool;
 }
 

diff --git a/source/base/kokkos.cc b/source/base/kokkos.cc
index 93ccad321bd39b3e8e9b55738600acfa3853f7f5..2b5678c7e3a96ce7c466724c640df6e43b95276f 100644 (file)
--- a/source/base/kokkos.cc
+++ b/source/base/kokkos.cc
@@ -24,39 +24,14 @@ DEAL_II_NAMESPACE_OPEN
 
 namespace Impl
 {
+  bool dealii_initialized_kokkos = false;
+
   void
   ensure_kokkos_initialized()
   {
     if (!Kokkos::is_initialized())
       {
-        GrowingVectorMemory<
-          LinearAlgebra::distributed::Vector<double, MemorySpace::Host>>{};
-        GrowingVectorMemory<
-          LinearAlgebra::distributed::Vector<float, MemorySpace::Host>>{};
-#ifdef DEAL_II_WITH_CUDA
-        GrowingVectorMemory<
-          LinearAlgebra::distributed::Vector<double, MemorySpace::CUDA>>{};
-        GrowingVectorMemory<
-          LinearAlgebra::distributed::Vector<float, MemorySpace::CUDA>>{};
-#endif
-        Kokkos::push_finalize_hook(
-          GrowingVectorMemory<
-            LinearAlgebra::distributed::Vector<double, MemorySpace::Host>>::
-            release_unused_memory);
-        Kokkos::push_finalize_hook(
-          GrowingVectorMemory<
-            LinearAlgebra::distributed::Vector<float, MemorySpace::Host>>::
-            release_unused_memory);
-#ifdef DEAL_II_WITH_CUDA
-        Kokkos::push_finalize_hook(
-          GrowingVectorMemory<
-            LinearAlgebra::distributed::Vector<double, MemorySpace::CUDA>>::
-            release_unused_memory);
-        Kokkos::push_finalize_hook(
-          GrowingVectorMemory<
-            LinearAlgebra::distributed::Vector<float, MemorySpace::CUDA>>::
-            release_unused_memory);
-#endif
+        dealii_initialized_kokkos = true;
         Kokkos::initialize();
         std::atexit(Kokkos::finalize);
       }