From 07d26630824132ed4afa33f3cc7b7fdb5ae12588 Mon Sep 17 00:00:00 2001 From: Daniel Arndt 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 +#include #include #include @@ -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 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(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