Merge pull request #2470 from lammps/kk_finalize

Fix issue with Kokkos::finalize and library interface

src/KOKKOS/kokkos.cpp

@@ -69,6 +69,10 @@ GPU_AWARE_UNKNOWN
 
 using namespace LAMMPS_NS;
 
+Kokkos::InitArguments KokkosLMP::args{-1, -1, -1, false};
+int KokkosLMP::is_finalized = 0;
+int KokkosLMP::init_ngpus = 0;
+
 /* ---------------------------------------------------------------------- */
 
 KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
@@ -155,6 +159,10 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
     } else if (strcmp(arg[iarg],"t") == 0 ||
                strcmp(arg[iarg],"threads") == 0) {
       nthreads = atoi(arg[iarg+1]);
+
+      if (nthreads <= 0)
+        error->all(FLERR,"Invalid number of threads requested for Kokkos: must be 1 or greater");
+
       iarg += 2;
 
     } else if (strcmp(arg[iarg],"n") == 0 ||
@@ -165,13 +173,27 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
     } else error->all(FLERR,"Invalid Kokkos command-line args");
   }
 
-  // initialize Kokkos
+  // Initialize Kokkos. However, we cannot change any
+  // Kokkos library parameters after the first initalization
 
-  if (me == 0) {
-    if (screen) fprintf(screen,"  will use up to %d GPU(s) per node\n",ngpus);
-    if (logfile) fprintf(logfile,"  will use up to %d GPU(s) per node\n",ngpus);
+  if (args.num_threads != -1) {
+    if (args.num_threads != nthreads || args.num_numa != numa || args.device_id != device)
+      if (me == 0)
+        error->warning(FLERR,"Kokkos package already initalized, cannot reinitialize with different parameters");
+    nthreads = args.num_threads;
+    numa = args.num_numa;
+    device = args.device_id;
+    ngpus = init_ngpus;
+  } else {
+    args.num_threads = nthreads;
+    args.num_numa = numa;
+    args.device_id = device;
+    init_ngpus = ngpus;
   }
 
+  if (me == 0)
+    utils::logmesg(lmp, "  will use up to {} GPU(s) per node\n",ngpus);
+
 #ifdef LMP_KOKKOS_GPU
   if (ngpus <= 0)
     error->all(FLERR,"Kokkos has been compiled for CUDA, HIP, or SYCL but no GPUs are requested");
@@ -184,12 +206,7 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
                          "than the OpenMP backend");
 #endif
 
-  Kokkos::InitArguments args;
-  args.num_threads = nthreads;
-  args.num_numa = numa;
-  args.device_id = device;
-
-  Kokkos::initialize(args);
+  KokkosLMP::initialize(args,error);
 
   // default settings for package kokkos command
 
@@ -299,9 +316,27 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
 
 KokkosLMP::~KokkosLMP()
 {
-  // finalize Kokkos
 
-  Kokkos::finalize();
+}
+
+/* ---------------------------------------------------------------------- */
+
+void KokkosLMP::initialize(Kokkos::InitArguments args, Error *error)
+{
+  if (!Kokkos::is_initialized()) {
+    if (is_finalized)
+      error->all(FLERR,"Kokkos package already finalized, cannot re-initialize");
+    Kokkos::initialize(args);
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
+void KokkosLMP::finalize()
+{
+  if (Kokkos::is_initialized() && !is_finalized)
+    Kokkos::finalize();
+  is_finalized = 1;
 }
 
 /* ----------------------------------------------------------------------

src/KOKKOS/kokkos.h

@@ -49,8 +49,14 @@ class KokkosLMP : protected Pointers {
   int newtonflag;
   double binsize;
 
+  static int is_finalized;
+  static Kokkos::InitArguments args;
+  static int init_ngpus;
+
   KokkosLMP(class LAMMPS *, int, char **);
   ~KokkosLMP();
+  static void initialize(Kokkos::InitArguments, Error *);
+  static void finalize();
   void accelerator(int, char **);
   int neigh_count(int);
 
@@ -84,13 +90,21 @@ because MPI library not recognized
 
 The local MPI rank was not found in one of four supported environment variables.
 
+E: Invalid number of threads requested for Kokkos: must be 1 or greater
+
+Self-explanatory.
+
 E: GPUs are requested but Kokkos has not been compiled for CUDA
 
 Recompile Kokkos with CUDA support to use GPUs.
 
-E: Kokkos has been compiled for CUDA but no GPUs are requested
+E: Kokkos has been compiled for CUDA, HIP, or SYCL but no GPUs are requested
 
-One or more GPUs must be used when Kokkos is compiled for CUDA.
+One or more GPUs must be used when Kokkos is compiled for CUDA/HIP/SYCL.
+
+W: Kokkos package already initalized, cannot reinitialize with different parameters
+
+Self-explanatory.
 
 E: Illegal ... command
 

src/accelerator_kokkos.h

@@ -56,16 +56,12 @@ class KokkosLMP {
 
   KokkosLMP(class LAMMPS *, int, char **) { kokkos_exists = 0; }
   ~KokkosLMP() {}
+  static void finalize() {}
   void accelerator(int, char **) {}
   int neigh_list_kokkos(int) { return 0; }
   int neigh_count(int) { return 0; }
 };
 
-class Kokkos {
- public:
-  static void finalize() {}
-};
-
 class AtomKokkos : public Atom {
  public:
   tagint **k_special;

src/error.cpp

@@ -81,7 +81,7 @@ void Error::universe_all(const std::string &file, int line, const std::string &s
 
   throw LAMMPSException(mesg);
 #else
-  if (lmp->kokkos) Kokkos::finalize();
+  KokkosLMP::finalize();
   MPI_Finalize();
   exit(1);
 #endif
@@ -107,6 +107,7 @@ void Error::universe_one(const std::string &file, int line, const std::string &s
 
   throw LAMMPSAbortException(mesg, universe->uworld);
 #else
+  KokkosLMP::finalize();
   MPI_Abort(universe->uworld,1);
   exit(1); // to trick "smart" compilers into believing this does not return
 #endif
@@ -173,8 +174,8 @@ void Error::all(const std::string &file, int line, const std::string &str)
   if (screen && screen != stdout) fclose(screen);
   if (logfile) fclose(logfile);
 
+  KokkosLMP::finalize();
   if (universe->nworlds > 1) MPI_Abort(universe->uworld,1);
-  if (lmp->kokkos) Kokkos::finalize();
   MPI_Finalize();
   exit(1);
 #endif
@@ -213,6 +214,7 @@ void Error::one(const std::string &file, int line, const std::string &str)
 #else
   if (screen) fflush(screen);
   if (logfile) fflush(logfile);
+  KokkosLMP::finalize();
   MPI_Abort(world,1);
   exit(1); // to trick "smart" compilers into believing this does not return
 #endif
@@ -315,7 +317,7 @@ void Error::done(int status)
   if (screen && screen != stdout) fclose(screen);
   if (logfile) fclose(logfile);
 
-  if (lmp->kokkos) Kokkos::finalize();
+  KokkosLMP::finalize();
   MPI_Finalize();
   exit(status);
 }

src/library.cpp

@@ -19,6 +19,7 @@
 #include "library.h"
 #include <mpi.h>
 
+#include "accelerator_kokkos.h"
 #include "atom.h"
 #include "atom_vec.h"
 #include "comm.h"
@@ -333,8 +334,8 @@ The MPI standard requires that any MPI application calls
 do any MPI calls, MPI is still initialized internally to avoid errors
 accessing any MPI functions.  This function should then be called right
 before exiting the program to wait until all (parallel) tasks are
-completed and then MPI is cleanly shut down.  After this function no
-more MPI calls may be made.
+completed and then MPI is cleanly shut down.  After calling this
+function no more MPI calls may be made.
 
 .. versionadded:: 18Sep2020
 
@@ -353,6 +354,28 @@ void lammps_mpi_finalize()
   }
 }
 
+/* ---------------------------------------------------------------------- */
+
+/** Shut down the Kokkos library environment.
+ *
+\verbatim embed:rst
+
+The Kokkos library may only be initialized once during the execution of
+a process.  This is done automatically the first time Kokkos
+functionality is used.  This requires that the Kokkos environment
+must be explicitly shut down after any LAMMPS instance using it is
+closed (to release associated resources).
+After calling this function no Kokkos functionality may be used.
+
+.. versionadded:: TBD
+
+\endverbatim */
+
+void lammps_kokkos_finalize()
+{
+  KokkosLMP::finalize();
+}
+
 // ----------------------------------------------------------------------
 // Library functions to process commands
 // ----------------------------------------------------------------------

src/library.h

@@ -94,6 +94,7 @@ void lammps_close(void *handle);
 
 void lammps_mpi_init();
 void lammps_mpi_finalize();
+void lammps_kokkos_finalize();
 
 /* ----------------------------------------------------------------------
  * Library functions to process commands

src/main.cpp

@@ -14,6 +14,7 @@
 #include "lammps.h"
 
 #include "input.h"
+#include "accelerator_kokkos.h"
 #if defined(LAMMPS_EXCEPTIONS)
 #include "exceptions.h"
 #endif
@@ -77,13 +78,16 @@ int main(int argc, char **argv)
     lammps->input->file();
     delete lammps;
   } catch (LAMMPSAbortException &ae) {
+    KokkosLMP::finalize();
     MPI_Abort(ae.universe, 1);
   } catch (LAMMPSException &e) {
+    KokkosLMP::finalize();
     MPI_Barrier(lammps_comm);
     MPI_Finalize();
     exit(1);
   } catch (fmt::format_error &fe) {
     fprintf(stderr, "fmt::format_error: %s\n", fe.what());
+    KokkosLMP::finalize();
     MPI_Abort(MPI_COMM_WORLD, 1);
     exit(1);
   }
@@ -94,10 +98,12 @@ int main(int argc, char **argv)
     delete lammps;
   } catch (fmt::format_error &fe) {
     fprintf(stderr, "fmt::format_error: %s\n", fe.what());
+    KokkosLMP::finalize();
     MPI_Abort(MPI_COMM_WORLD, 1);
     exit(1);
   }
 #endif
+  KokkosLMP::finalize();
   MPI_Barrier(lammps_comm);
   MPI_Finalize();
 }