From 07d26630824132ed4afa33f3cc7b7fdb5ae12588 Mon Sep 17 00:00:00 2001
From: Daniel Arndt <arndtd@ornl.gov>
Date: Thu, 9 Mar 2023 13:50:19 -0500
Subject: [PATCH] Initialize Kokkos with MultithreadInfo::n_threads()

---
 source/base/kokkos.cc |  6 +++++-
 source/base/mpi.cc    | 12 +++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/source/base/kokkos.cc b/source/base/kokkos.cc
index 5ea3b374c1..18f41861fa 100644
--- a/source/base/kokkos.cc
+++ b/source/base/kokkos.cc
@@ -14,6 +14,7 @@
 // ---------------------------------------------------------------------
 
 #include <deal.II/base/kokkos.h>
+#include <deal.II/base/multithread_info.h>
 
 #include <deal.II/lac/la_parallel_vector.h>
 #include <deal.II/lac/vector_memory.h>
@@ -34,7 +35,10 @@ namespace internal
         // only execute once
         static bool dummy = [] {
           dealii_initialized_kokkos = true;
-          Kokkos::initialize();
+          const auto settings =
+            Kokkos::InitializationSettings().set_num_threads(
+              MultithreadInfo::n_threads());
+          Kokkos::initialize(settings);
           std::atexit(Kokkos::finalize);
           return true;
         }();
diff --git a/source/base/mpi.cc b/source/base/mpi.cc
index cd1d1ea8d1..cd06ceaadd 100644
--- a/source/base/mpi.cc
+++ b/source/base/mpi.cc
@@ -776,7 +776,17 @@ namespace Utilities
 #endif
 
       // Initialize Kokkos
-      Kokkos::initialize(argc, argv);
+      {
+        std::vector<char *> argv_new;
+        for (int i = 0; i < argc; ++i)
+          argv_new.push_back(argv[i]);
+        std::stringstream threads_flag;
+        threads_flag << "--kokkos-threads=" << MultithreadInfo::n_threads();
+        argv_new.push_back(const_cast<char *>(threads_flag.str().c_str()));
+        argv_new.push_back(nullptr);
+        int argc_new = argc + 1;
+        Kokkos::initialize(argc_new, argv_new.data());
+      }
 
       // we are allowed to call MPI_Init ourselves and PETScInitialize will
       // detect this. This allows us to use MPI_Init_thread instead.
-- 
2.39.5