Merge pull request #2470 from lammps/kk_finalize

Fix issue with Kokkos::finalize and library interface
This commit is contained in:
Axel Kohlmeyer
2021-07-01 21:11:49 -04:00
committed by GitHub
13 changed files with 130 additions and 35 deletions

View File

@ -9,6 +9,7 @@ This section documents the following functions:
- :cpp:func:`lammps_close`
- :cpp:func:`lammps_mpi_init`
- :cpp:func:`lammps_mpi_finalize`
- :cpp:func:`lammps_kokkos_finalize`
--------------------

View File

@ -134,7 +134,10 @@ compiled with.
The :py:func:`lmp.close() <lammps.lammps.close()>` call is
optional since the LAMMPS class instance will also be deleted
automatically during the :py:class:`lammps <lammps.lammps>` class
destructor.
destructor. Instead of :py:func:`lmp.close() <lammps.lammps.close()>`
it is also possible to call :py:func:`lmp.finalize() <lammps.lammps.finalize()>`;
this will destruct the LAMMPS instance, but also finalized and release
the MPI and/or Kokkos environment if enabled and active.
Note that you can create multiple LAMMPS objects in your Python
script, and coordinate and run multiple simulations, e.g.

View File

@ -76,17 +76,15 @@ MODULE LIBLAMMPS
TYPE(c_ptr), VALUE :: handle
END SUBROUTINE lammps_close
SUBROUTINE lammps_mpi_init(handle) BIND(C, name='lammps_mpi_init')
IMPORT :: c_ptr
TYPE(c_ptr), VALUE :: handle
SUBROUTINE lammps_mpi_init() BIND(C, name='lammps_mpi_init')
END SUBROUTINE lammps_mpi_init
SUBROUTINE lammps_mpi_finalize(handle) &
BIND(C, name='lammps_mpi_finalize')
IMPORT :: c_ptr
TYPE(c_ptr), VALUE :: handle
SUBROUTINE lammps_mpi_finalize() BIND(C, name='lammps_mpi_finalize')
END SUBROUTINE lammps_mpi_finalize
SUBROUTINE lammps_kokkos_finalize() BIND(C, name='lammps_kokkos_finalize')
END SUBROUTINE lammps_kokkos_finalize
SUBROUTINE lammps_file(handle,filename) BIND(C, name='lammps_file')
IMPORT :: c_ptr
TYPE(c_ptr), VALUE :: handle
@ -188,7 +186,8 @@ CONTAINS
IF (PRESENT(finalize)) THEN
IF (finalize) THEN
CALL lammps_mpi_finalize(self%handle)
CALL lammps_kokkos_finalize()
CALL lammps_mpi_finalize()
END IF
END IF
END SUBROUTINE lmp_close

View File

@ -460,10 +460,16 @@ class lammps(object):
# -------------------------------------------------------------------------
def finalize(self):
"""Shut down the MPI communication through the library interface by calling :cpp:func:`lammps_finalize`.
"""Shut down the MPI communication and Kokkos environment (if active) through the
library interface by calling :cpp:func:`lammps_mpi_finalize` and
:cpp:func:`lammps_kokkos_finalize`.
You cannot create or use any LAMMPS instances after this function is called
unless LAMMPS was compiled without MPI and without Kokkos support.
"""
self.close()
self.lib.lammps_finalize()
self.lib.lammps_kokkos_finalize()
self.lib.lammps_mpi_finalize()
# -------------------------------------------------------------------------

View File

@ -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;
}
/* ----------------------------------------------------------------------

View File

@ -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

View File

@ -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;

View File

@ -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);
}

View File

@ -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
// ----------------------------------------------------------------------

View File

@ -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

View File

@ -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();
}

View File

@ -63,6 +63,7 @@ extern void *lammps_open_fortran(int argc, char **argv, int f_comm);
extern void lammps_close(void *handle);
extern void lammps_mpi_init();
extern void lammps_mpi_finalize();
extern void lammps_kokkos_finalize();
extern void lammps_file(void *handle, const char *file);
extern char *lammps_command(void *handle, const char *cmd);
extern void lammps_commands_list(void *handle, int ncmd, const char **cmds);
@ -185,6 +186,7 @@ extern void *lammps_open_fortran(int argc, char **argv, int f_comm);
extern void lammps_close(void *handle);
extern void lammps_mpi_init();
extern void lammps_mpi_finalize();
extern void lammps_kokkos_finalize();
extern void lammps_file(void *handle, const char *file);
extern char *lammps_command(void *handle, const char *cmd);
extern void lammps_commands_list(void *handle, int ncmd, const char **cmds);

View File

@ -165,6 +165,13 @@ class PythonCapabilities(unittest.TestCase):
if self.cmake_cache['GPU_PREC'].lower() == 'single':
self.assertIn('single',settings['GPU']['precision'])
if self.cmake_cache['PKG_KOKKOS']:
if self.cmake_cache['Kokkos_ENABLE_OPENMP']:
self.assertIn('openmp',settings['KOKKOS']['api'])
if self.cmake_cache['Kokkos_ENABLE_SERIAL']:
self.assertIn('serial',settings['KOKKOS']['api'])
self.assertIn('double',settings['KOKKOS']['precision'])
def test_gpu_device(self):
info = self.lmp.get_gpu_device_info()