From 4c0fa31cee3181b4f7a6243439ba44346f38cce1 Mon Sep 17 00:00:00 2001
From: Daniel Arndt <arndtd@ornl.gov>
Date: Sun, 4 Dec 2022 14:30:17 -0500
Subject: [PATCH] Improve Kokkos initialization

---
 include/deal.II/base/kokkos.h                 | 11 ++++++-
 include/deal.II/lac/vector_memory.templates.h | 17 +++++++++-
 source/base/kokkos.cc                         | 31 ++-----------------
 3 files changed, 29 insertions(+), 30 deletions(-)

diff --git a/include/deal.II/base/kokkos.h b/include/deal.II/base/kokkos.h
index caccb94466..6c020fd825 100644
--- 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 e6c13b649c..48db94c9b1 100644
--- 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 93ccad321b..2b5678c7e3 100644
--- 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);
       }
-- 
2.39.5