From 15999f6518249d6b3104c502b1928dda0fca476f Mon Sep 17 00:00:00 2001 From: ljwoods2 Date: Tue, 12 Nov 2024 14:30:10 -0700 Subject: [PATCH] initial commit --- src/MISC/fix_imd.cpp | 779 +++++++++++++++++++++++++++++++++++++------ src/MISC/fix_imd.h | 26 +- 2 files changed, 704 insertions(+), 101 deletions(-) diff --git a/src/MISC/fix_imd.cpp b/src/MISC/fix_imd.cpp index 41ce97ee65..c295fdfaf8 100644 --- a/src/MISC/fix_imd.cpp +++ b/src/MISC/fix_imd.cpp @@ -76,6 +76,7 @@ negotiate an appropriate license for such distribution." #endif #include +#include using namespace LAMMPS_NS; using namespace FixConst; @@ -358,7 +359,6 @@ typedef struct { } IMDheader; #define IMDHEADERSIZE 8 -#define IMDVERSION 2 typedef enum IMDType_t { IMD_DISCONNECT, /**< close IMD connection, leaving sim running */ @@ -370,7 +370,15 @@ typedef enum IMDType_t { IMD_MDCOMM, /**< MDComm style force data */ IMD_PAUSE, /**< pause the running simulation */ IMD_TRATE, /**< set IMD update transmission rate */ - IMD_IOERROR /**< indicate an I/O error */ + IMD_IOERROR, /**< indicate an I/O error */ + /* IMDv3 only headers */ + IMD_SESSIONINFO, + IMD_RESUME, + IMD_TIME, + IMD_BOX, + IMD_VELOCITIES, + IMD_FORCES, + IMD_WAIT, } IMDType; /**< IMD command message type enumerations */ typedef struct { @@ -386,8 +394,20 @@ typedef struct { float Eimpr; /**< Improper energy, Kcal/mol */ } IMDEnergies; /**< IMD simulation energy report structure */ +/* IMDv3 only */ +typedef struct IMDSessionInfo { + bool time; + bool box; + bool coords; + bool unwrap; + bool velocities; + bool forces; + bool energies; +} IMDSessionInfo; + /** Send control messages - these consist of a header with no subsequent data */ -static int imd_handshake(void *); /**< check endianness, version compat */ +static int imd_handshake_v2(void *); /**< check endianness, version compat */ +static int imd_handshake_v3(void *, IMDSessionInfo *); /** Receive header and data */ static IMDType imd_recv_header(void *, int32 *); /** Receive MDComm-style forces, units are Kcal/mol/angstrom */ @@ -426,19 +446,22 @@ static void imdsock_destroy(void *); * The implementation follows at the end of the file. * ***************************************************************/ -/* struct for packed data communication of coordinates and forces. */ +/* struct for packed data communication of coordinates, velocities, and forces. */ struct commdata { tagint tag; float x,y,z; }; +MPI_Datatype MPI_CommData; + /*************************************************************** * create class and parse arguments in LAMMPS script. Syntax: - * fix ID group-ID imd [unwrap (on|off)] [fscale ] + * fix ID group-ID imd [version (2|3)] [unwrap (on|off)] [fscale ] [time (on|off)] [box (on|off)] [coordinates (on|off)] [velocities (on|off)] [forces (on|off)] ***************************************************************/ FixIMD::FixIMD(LAMMPS *lmp, int narg, char **arg) : Fix(lmp, narg, arg) { + if (narg < 4) error->all(FLERR,"Illegal fix imd command"); @@ -447,12 +470,21 @@ FixIMD::FixIMD(LAMMPS *lmp, int narg, char **arg) : error->all(FLERR,"Illegal fix imd parameter: port < 1024"); /* default values for optional flags */ + imd_version = 2; + unwrap_flag = 0; nowait_flag = 0; connect_msg = 1; imd_fscale = 1.0; imd_trate = 1; + /* IMDv3-only flags. Aren't stored as class attributes since they're converted into IMDSessionInfo */ + int time_flag = 1; + int box_flag = 1; + int coord_flag = 1; + int vel_flag = 1; + int force_flag = 1; + /* parse optional arguments */ int iarg = 4; while (iarg+1 < narg) { @@ -464,9 +496,19 @@ FixIMD::FixIMD(LAMMPS *lmp, int narg, char **arg) : imd_fscale = utils::numeric(FLERR,arg[iarg+1],false,lmp); } else if (0 == strcmp(arg[iarg], "trate")) { imd_trate = utils::inumeric(FLERR,arg[iarg+1],false,lmp); - } else { - error->all(FLERR,"Unknown fix imd parameter"); - } + } else if (0 == strcmp(arg[iarg], "version")) { + imd_version = utils::inumeric(FLERR,arg[iarg+1],false,lmp); + } else if (0 == strcmp(arg[iarg], "time")) { + time_flag = utils::logical(FLERR, arg[iarg+1], false, lmp); + } else if (0 == strcmp(arg[iarg], "box")) { + box_flag = utils::logical(FLERR, arg[iarg+1], false, lmp); + } else if (0 == strcmp(arg[iarg], "coordinates")) { + coord_flag = utils::logical(FLERR, arg[iarg+1], false, lmp); + } else if (0 == strcmp(arg[iarg], "velocities")) { + vel_flag = utils::logical(FLERR, arg[iarg+1], false, lmp); + } else if (0 == strcmp(arg[iarg], "forces")) { + force_flag = utils::logical(FLERR, arg[iarg+1], false, lmp); + } else error->all(FLERR,"Unknown fix imd parameter"); ++iarg; ++iarg; } @@ -474,12 +516,43 @@ FixIMD::FixIMD(LAMMPS *lmp, int narg, char **arg) : if (imd_trate < 1) error->all(FLERR,"Illegal fix imd parameter. trate < 1."); + if (imd_version != 2 && imd_version != 3) + error->all(FLERR,"Illegal fix imd parameter. version != 2 or 3."); + + imdsinfo = new IMDSessionInfo; + + /* In IMDv2 in LAMMPS, only coordinates are sent*/ + if (imd_version == 2) { + imdsinfo->time = false; + imdsinfo->box = false; + imdsinfo->coords = true; + imdsinfo->unwrap = unwrap_flag; + imdsinfo->velocities = false; + imdsinfo->forces = false; + imdsinfo->energies = false; + } + + if (imd_version == 3) { + imdsinfo->time = time_flag; + imdsinfo->box = box_flag; + imdsinfo->coords = coord_flag; + imdsinfo->unwrap = unwrap_flag; + imdsinfo->velocities = vel_flag; + imdsinfo->forces = force_flag; + imdsinfo->energies = false; + } + + bigint n = group->count(igroup); if (n > MAXSMALLINT) error->all(FLERR,"Too many atoms for fix imd"); num_coords = static_cast (n); - MPI_Comm_rank(world,&me); + // Define MPI dtype for passing x/v/f data + MPI_Type_contiguous(4, MPI_FLOAT, &MPI_CommData); + MPI_Type_commit(&MPI_CommData); + MPI_Comm_rank(world,&me); + /* initialize various imd state variables. */ clientsock = nullptr; localsock = nullptr; @@ -487,14 +560,33 @@ FixIMD::FixIMD(LAMMPS *lmp, int narg, char **arg) : imd_inactive = 0; imd_terminate = 0; imd_forces = 0; - force_buf = nullptr; + recv_force_buf = nullptr; maxbuf = 0; - msgdata = nullptr; - msglen = 0; - comm_buf = nullptr; + coord_data = nullptr; + vel_data = nullptr; + force_data = nullptr; idmap = nullptr; rev_idmap = nullptr; + + msglen = 0; + if (imdsinfo->time) { + msglen += 24+IMDHEADERSIZE; + } + if (imdsinfo->box) { + msglen += 9*4+IMDHEADERSIZE; + } + if (imdsinfo->coords) { + msglen += 3*4*num_coords+IMDHEADERSIZE; + } + if (imdsinfo->velocities) { + msglen += 3*4*num_coords+IMDHEADERSIZE; + } + if (imdsinfo->forces) { + msglen += 3*4*num_coords+IMDHEADERSIZE; + } + msgdata = new char[msglen]; + if (me == 0) { /* set up incoming socket on MPI rank 0. */ imdsock_init(); @@ -512,7 +604,7 @@ FixIMD::FixIMD(LAMMPS *lmp, int narg, char **arg) : if (imd_terminate) error->all(FLERR,"LAMMPS Terminated on error in IMD."); - /* storage required to communicate a single coordinate or force. */ + /* storage required to communicate a single coordinate, velocity, or force. */ size_one = sizeof(struct commdata); #if defined(LAMMPS_ASYNC_IMD) @@ -546,7 +638,6 @@ FixIMD::FixIMD(LAMMPS *lmp, int narg, char **arg) : *********************************/ FixIMD::~FixIMD() { - #if defined(LAMMPS_ASYNC_IMD) if (me == 0) { pthread_mutex_lock(&write_mutex); @@ -561,13 +652,17 @@ FixIMD::~FixIMD() pthread_cond_destroy(&write_cond); } #endif - auto hashtable = (taginthash_t *)idmap; - memory->destroy(comm_buf); - memory->destroy(force_buf); + memory->destroy(coord_data); + memory->destroy(vel_data); + memory->destroy(force_data); + + memory->destroy(msgdata); + memory->destroy(recv_force_buf); taginthash_destroy(hashtable); delete hashtable; free(rev_idmap); + free(imdsinfo); // close sockets imdsock_shutdown(clientsock); imdsock_destroy(clientsock); @@ -583,6 +678,7 @@ int FixIMD::setmask() int mask = 0; mask |= POST_FORCE; mask |= POST_FORCE_RESPA; + mask |= END_OF_STEP; return mask; } @@ -611,6 +707,7 @@ int FixIMD::reconnect() } else { fprintf(screen,"Waiting for IMD connection on port %d. Transfer rate %d.\n",imd_port, imd_trate); } + fflush(screen); } connect_msg = 0; clientsock = nullptr; @@ -637,7 +734,8 @@ int FixIMD::reconnect() return 0; } else { /* check endianness and IMD protocol version. */ - if (imd_handshake(clientsock)) { + if ((imd_version == 2 && imd_handshake_v2(clientsock)) || + (imd_version == 3 && imd_handshake_v3(clientsock, imdsinfo))) { if (screen) fprintf(screen, "IMD handshake error. Dropping connection.\n"); imdsock_destroy(clientsock); @@ -680,9 +778,21 @@ void FixIMD::setup(int) if (mask[i] & groupbit) ++nme; MPI_Allreduce(&nme,&nmax,1,MPI_INT,MPI_MAX,world); - memory->destroy(comm_buf); + maxbuf = nmax*size_one; - comm_buf = (void *) memory->smalloc(maxbuf,"imd:comm_buf"); + + if (imdsinfo->coords) { + memory->destroy(coord_data); + coord_data = (void *) memory->smalloc(maxbuf,"imd:coord_data"); + } + if (imdsinfo->velocities) { + memory->destroy(vel_data); + vel_data = (void *) memory->smalloc(maxbuf,"imd:vel_data"); + } + if (imdsinfo->forces) { + memory->destroy(force_data); + force_data = (void *) memory->smalloc(maxbuf,"imd:force_data"); + } connect_msg = 1; reconnect(); @@ -697,11 +807,11 @@ void FixIMD::setup(int) idmap = (void *)hashtable; int tmp, ndata; - auto buf = static_cast(comm_buf); + auto buf = static_cast(coord_data); if (me == 0) { - MPI_Status status; - MPI_Request request; + std::vector statuses; + std::vector requests; auto taglist = new tagint[num_coords]; int numtag=0; /* counter to map atom tags to a 0-based consecutive index list */ @@ -714,15 +824,39 @@ void FixIMD::setup(int) /* loop over procs to receive remote data */ for (i=1; i < comm->nprocs; ++i) { - MPI_Irecv(comm_buf, maxbuf, MPI_BYTE, i, 0, world, &request); - MPI_Send(&tmp, 0, MPI_INT, i, 0, world); - MPI_Wait(&request, &status); - MPI_Get_count(&status, MPI_BYTE, &ndata); - ndata /= size_one; + /* We're assuming tags are consistent across x,v,f */ + bool tag_recvd = false; + statuses.clear(); + requests.clear(); - for (j=0; j < ndata; ++j) { - taglist[numtag] = buf[j].tag; - ++numtag; + if (imdsinfo->coords) { + requests.push_back(MPI_Request()); + MPI_Irecv(coord_data, maxbuf, MPI_BYTE, i, 0, world, &requests.back()); + } + if (imdsinfo->velocities) { + requests.push_back(MPI_Request()); + MPI_Irecv(vel_data, maxbuf, MPI_BYTE, i, 0, world, &requests.back()); + } + if (imdsinfo->forces) { + requests.push_back(MPI_Request()); + MPI_Irecv(vel_data, maxbuf, MPI_BYTE, i, 0, world, &requests.back()); + } + statuses.resize(requests.size()); + MPI_Send(&tmp, 0, MPI_INT, i, 0, world); + MPI_Waitall(requests.size(), requests.data(), statuses.data()); + + for (size_t k=0; k < statuses.size(); ++k) { + if (!tag_recvd) { + MPI_Get_count(&statuses[k], MPI_BYTE, &ndata); + ndata /= size_one; + for (j=0; j < ndata; ++j) { + taglist[numtag] = buf[j].tag; + ++numtag; + } + tag_recvd = true; + } else { + break; + } } } @@ -747,7 +881,15 @@ void FixIMD::setup(int) } /* blocking receive to wait until it is our turn to send data. */ MPI_Recv(&tmp, 0, MPI_INT, 0, 0, world, MPI_STATUS_IGNORE); - MPI_Rsend(comm_buf, nme*size_one, MPI_BYTE, 0, 0, world); + if (imdsinfo->coords) { + MPI_Rsend(coord_data, nme*size_one, MPI_BYTE, 0, 0, world); + } + if (imdsinfo->velocities) { + MPI_Rsend(vel_data, nme*size_one, MPI_BYTE, 0, 0, world); + } + if (imdsinfo->forces) { + MPI_Rsend(force_data, nme*size_one, MPI_BYTE, 0, 0, world); + } } } @@ -795,6 +937,39 @@ void FixIMD::ioworker() * Send coodinates, energies, and add IMD forces to atoms. */ void FixIMD::post_force(int /*vflag*/) { + fflush(screen); + if (imd_version == 2) { + handle_step_v2(); + } + else if (imd_version == 3) { + handle_client_input_v3(); + } + + } + +/* ---------------------------------------------------------------------- */ +void FixIMD::post_force_respa(int vflag, int ilevel, int /*iloop*/) +{ + /* only process IMD on the outmost RESPA level. */ + if (ilevel == nlevels_respa-1) post_force(vflag); +} + +void FixIMD::end_of_step() +{ + if (imd_version == 3 && update->ntimestep % imd_trate == 0) { + handle_output_v3(); + } +} + +/* ---------------------------------------------------------------------- */ +/* local memory usage. approximately. */ +double FixIMD::memory_usage() +{ + return static_cast(num_coords+maxbuf+imd_forces)*size_one; +} + +void FixIMD::handle_step_v2() { + /* check for reconnect */ if (imd_inactive) { reconnect(); @@ -825,22 +1000,12 @@ void FixIMD::post_force(int /*vflag*/) switch(msg) { - case IMD_GO: - if (screen) - fprintf(screen, "Ignoring unexpected IMD_GO message.\n"); - break; - - case IMD_IOERROR: - if (screen) - fprintf(screen, "IMD connection lost.\n"); - /* fallthrough */ - case IMD_DISCONNECT: { /* disconnect from client. wait for new connection. */ imd_paused = 0; imd_forces = 0; - memory->destroy(force_buf); - force_buf = nullptr; + memory->destroy(recv_force_buf); + recv_force_buf = nullptr; imdsock_destroy(clientsock); clientsock = nullptr; if (screen) @@ -866,9 +1031,9 @@ void FixIMD::post_force(int /*vflag*/) case IMD_PAUSE: /* pause the running simulation. wait for second IMD_PAUSE to continue. */ if (imd_paused) { - if (screen) - fprintf(screen, "Continuing run on IMD client request.\n"); - imd_paused = 0; + if (screen) + fprintf(screen, "Continuing run on IMD client request.\n"); + imd_paused = 0; } else { if (screen) fprintf(screen, "Pausing run on IMD client request.\n"); @@ -884,31 +1049,18 @@ void FixIMD::post_force(int /*vflag*/) fprintf(screen, "IMD client requested change of transfer rate. Now it is %d.\n", imd_trate); break; - case IMD_ENERGIES: { - IMDEnergies dummy_energies; - imd_recv_energies(clientsock, &dummy_energies); - break; - } - - case IMD_FCOORDS: { - auto dummy_coords = new float[3*length]; - imd_recv_fcoords(clientsock, length, dummy_coords); - delete[] dummy_coords; - break; - } - case IMD_MDCOMM: { auto imd_tags = new int32[length]; auto imd_fdat = new float[3*length]; imd_recv_mdcomm(clientsock, length, imd_tags, imd_fdat); if (imd_forces < length) { /* grow holding space for forces, if needed. */ - memory->destroy(force_buf); - force_buf = (void *) memory->smalloc((bigint)length*size_one, - "imd:force_buf"); + memory->destroy(recv_force_buf); + recv_force_buf = (void *) memory->smalloc((bigint)length*size_one, + "imd:recv_force_buf"); } imd_forces = length; - buf = static_cast(force_buf); + buf = static_cast(recv_force_buf); /* compare data to hash table */ for (int ii=0; ii < length; ++ii) { @@ -943,12 +1095,12 @@ void FixIMD::post_force(int /*vflag*/) /* check if we need to readjust the forces comm buffer on the receiving nodes. */ if (me != 0) { if (old_imd_forces < imd_forces) { /* grow holding space for forces, if needed. */ - if (force_buf != nullptr) - memory->sfree(force_buf); - force_buf = memory->smalloc((bigint)imd_forces*size_one, "imd:force_buf"); + if (recv_force_buf != nullptr) + memory->sfree(recv_force_buf); + recv_force_buf = memory->smalloc((bigint)imd_forces*size_one, "imd:recv_force_buf"); } } - MPI_Bcast(force_buf, imd_forces*size_one, MPI_BYTE, 0, world); + MPI_Bcast(recv_force_buf, imd_forces*size_one, MPI_BYTE, 0, world); } /* Check if we need to communicate coordinates to the client. @@ -961,7 +1113,7 @@ void FixIMD::post_force(int /*vflag*/) if (update->ntimestep % imd_trate) { if (imd_forces > 0) { double **f = atom->f; - buf = static_cast(force_buf); + buf = static_cast(recv_force_buf); /* XXX. this is in principle O(N**2) == not good. * however we assume for now that the number of atoms @@ -989,27 +1141,25 @@ void FixIMD::post_force(int /*vflag*/) MPI_Allreduce(&nme,&nmax,1,MPI_INT,MPI_MAX,world); if (nmax*size_one > maxbuf) { - memory->destroy(comm_buf); + memory->destroy(coord_data); maxbuf = nmax*size_one; - comm_buf = memory->smalloc(maxbuf,"imd:comm_buf"); + coord_data = memory->smalloc(maxbuf,"imd:coord_data"); } int tmp, ndata; - buf = static_cast(comm_buf); + buf = static_cast(coord_data); if (me == 0) { MPI_Status status; MPI_Request request; /* collect data into new array. we bypass the IMD API to save * us one extra copy of the data. */ - msglen = 3*sizeof(float)*num_coords+IMDHEADERSIZE; - msgdata = new char[msglen]; imd_fill_header((IMDheader *)msgdata, IMD_FCOORDS, num_coords); /* array pointer, to the offset where we receive the coordinates. */ - auto recvcoord = (float *) (msgdata+IMDHEADERSIZE); + auto recvcoord = (float *) (msgdata+IMDHEADERSIZE); /* add local data */ - if (unwrap_flag) { + if (imdsinfo->unwrap) { double xprd = domain->xprd; double yprd = domain->yprd; double zprd = domain->zprd; @@ -1049,10 +1199,9 @@ void FixIMD::post_force(int /*vflag*/) } } } - /* loop over procs to receive remote data */ for (i=1; i < comm->nprocs; ++i) { - MPI_Irecv(comm_buf, maxbuf, MPI_BYTE, i, 0, world, &request); + MPI_Irecv(coord_data, maxbuf, MPI_BYTE, i, 0, world, &request); MPI_Send(&tmp, 0, MPI_INT, i, 0, world); MPI_Wait(&request, &status); MPI_Get_count(&status, MPI_BYTE, &ndata); @@ -1067,7 +1216,6 @@ void FixIMD::post_force(int /*vflag*/) } } } - /* done collecting frame data now communicate with IMD client. */ #if defined(LAMMPS_ASYNC_IMD) @@ -1081,13 +1229,12 @@ void FixIMD::post_force(int /*vflag*/) if (clientsock && imdsock_selwrite(clientsock,0)) { imd_writen(clientsock, msgdata, msglen); } - delete[] msgdata; #endif } else { /* copy coordinate data into communication buffer */ nme = 0; - if (unwrap_flag) { + if (imdsinfo->unwrap) { double xprd = domain->xprd; double yprd = domain->yprd; double zprd = domain->zprd; @@ -1128,23 +1275,439 @@ void FixIMD::post_force(int /*vflag*/) } /* blocking receive to wait until it is our turn to send data. */ MPI_Recv(&tmp, 0, MPI_INT, 0, 0, world, MPI_STATUS_IGNORE); - MPI_Rsend(comm_buf, nme*size_one, MPI_BYTE, 0, 0, world); + MPI_Rsend(coord_data, nme*size_one, MPI_BYTE, 0, 0, world); } - } - -/* ---------------------------------------------------------------------- */ -void FixIMD::post_force_respa(int vflag, int ilevel, int /*iloop*/) -{ - /* only process IMD on the outmost RESPA level. */ - if (ilevel == nlevels_respa-1) post_force(vflag); } -/* ---------------------------------------------------------------------- */ -/* local memory usage. approximately. */ -double FixIMD::memory_usage() -{ - return static_cast(num_coords+maxbuf+imd_forces)*size_one; +void FixIMD::handle_client_input_v3() { + // IMDV3 + + /* check for reconnect */ + if (imd_inactive) { + reconnect(); + MPI_Bcast(&imd_inactive, 1, MPI_INT, 0, world); + MPI_Bcast(&imd_terminate, 1, MPI_INT, 0, world); + if (imd_terminate) + error->all(FLERR,"LAMMPS terminated on error in setting up IMD connection."); + if (imd_inactive) + return; /* IMD client has detached and not yet come back. do nothing. */ + } + + struct commdata *buf; + int nlocal = atom->nlocal; + int *mask = atom->mask; + tagint *tag = atom->tag; + double **f = atom->f; + + if (me == 0) { + /* process all pending incoming data. */ + int imd_paused=0; + while ((imdsock_selread(clientsock, 0) > 0) || imd_paused) { + /* if something requested to turn off IMD while paused get out */ + if (imd_inactive) break; + + int32 length; + int msg = imd_recv_header(clientsock, &length); + + switch(msg) { + + case IMD_DISCONNECT: { + /* disconnect from client. wait for new connection. */ + imd_paused = 0; + imd_forces = 0; + memory->destroy(recv_force_buf); + recv_force_buf = nullptr; + imdsock_destroy(clientsock); + clientsock = nullptr; + if (screen) + fprintf(screen, "IMD client detached. LAMMPS run continues.\n"); + + connect_msg = 1; + reconnect(); + if (imd_terminate) imd_inactive = 1; + break; + } + + case IMD_KILL: + /* stop the simulation job and shutdown IMD */ + if (screen) + fprintf(screen, "IMD client requested termination of run.\n"); + imd_inactive = 1; + imd_terminate = 1; + imd_paused = 0; + imdsock_destroy(clientsock); + clientsock = nullptr; + break; + + case IMD_PAUSE: + /* pause the running simulation. wait for second IMD_PAUSE to continue. */ + if (!imd_paused) { + if (screen) + fprintf(screen, "Pausing run on IMD client request.\n"); + imd_paused = 1; + } else { + // Pause in IMDv3 is idempotent + continue; + } + break; + + case IMD_RESUME: + /* resume the running simulation. */ + if (imd_paused) { + if (screen) + fprintf(screen, "Continuing run on IMD client request.\n"); + imd_paused = 0; + } else { + // Resume in IMDv3 is idempotent + continue; + } + break; + + case IMD_TRATE: + /* change the IMD transmission data rate */ + if (length > 0) + imd_trate = length; + if (screen) + fprintf(screen, "IMD client requested change of transfer rate. Now it is %d.\n", imd_trate); + break; + + case IMD_MDCOMM: { + auto imd_tags = new int32[length]; + auto imd_fdat = new float[3*length]; + imd_recv_mdcomm(clientsock, length, imd_tags, imd_fdat); + + if (imd_forces < length) { /* grow holding space for forces, if needed. */ + memory->destroy(recv_force_buf); + recv_force_buf = (void *) memory->smalloc((bigint)length*size_one, + "imd:recv_force_buf"); + } + imd_forces = length; + buf = static_cast(recv_force_buf); + + /* compare data to hash table */ + for (int ii=0; ii < length; ++ii) { + buf[ii].tag = rev_idmap[imd_tags[ii]]; + buf[ii].x = imd_fdat[3*ii]; + buf[ii].y = imd_fdat[3*ii+1]; + buf[ii].z = imd_fdat[3*ii+2]; + } + delete[] imd_tags; + delete[] imd_fdat; + break; + } + case IMD_WAIT: { + /* Change IMD waiting behavior mid-session */ + if (length) { + nowait_flag = 0; + } + else { + nowait_flag = 1; + } + break; + } + + default: + if (screen) + fprintf(screen, "Unhandled incoming IMD message #%d. length=%d\n", msg, length); + break; + } + } + } + + /* update all tasks with current settings. */ + int old_imd_forces = imd_forces; + MPI_Bcast(&imd_trate, 1, MPI_INT, 0, world); + MPI_Bcast(&imd_inactive, 1, MPI_INT, 0, world); + MPI_Bcast(&imd_forces, 1, MPI_INT, 0, world); + MPI_Bcast(&imd_terminate, 1, MPI_INT, 0, world); + if (imd_terminate) + error->all(FLERR,"LAMMPS terminated on IMD request."); + + if (imd_forces > 0) { + /* check if we need to readjust the forces comm buffer on the receiving nodes. */ + if (me != 0) { + if (old_imd_forces < imd_forces) { /* grow holding space for forces, if needed. */ + if (recv_force_buf != nullptr) + memory->sfree(recv_force_buf); + recv_force_buf = memory->smalloc((bigint)imd_forces*size_one, "imd:recv_force_buf"); + } + } + MPI_Bcast(recv_force_buf, imd_forces*size_one, MPI_BYTE, 0, world); + } + + /* Check if we need to communicate coordinates to the client. + * Tuning imd_trate allows to keep the overhead for IMD low + * at the expense of a more jumpy display. Rather than using + * end_of_step() we do everything here in one go. + * + * If we don't communicate, only check if we have forces + * stored away and apply them. */ + if (imd_forces > 0) { + buf = static_cast(recv_force_buf); + + /* XXX. this is in principle O(N**2) == not good. + * however we assume for now that the number of atoms + * that we manipulate via IMD will be small compared + * to the total system size, so we don't hurt too much. */ + for (int j=0; j < imd_forces; ++j) { + for (int i=0; i < nlocal; ++i) { + if (mask[i] & groupbit) { + if (buf[j].tag == tag[i]) { + f[i][0] += imd_fscale*buf[j].x; + f[i][1] += imd_fscale*buf[j].y; + f[i][2] += imd_fscale*buf[j].z; + } + } + } + } + } +} + +void FixIMD::handle_output_v3() { + + tagint *tag = atom->tag; + double **x = atom->x; + double **v = atom->v; + double **f = atom->f; + imageint *image = atom->image; + int nlocal = atom->nlocal; + int *mask = atom->mask; + struct commdata *buf; + + // Only main process should use: + float *global_coords = nullptr; + float *global_vel = nullptr; + float *global_force = nullptr; + + // Prepare offsets in outgoing buffer + // Fill what we can wihtout collecting from other processes + if (me == 0) { + int offset = 0; + if (imdsinfo->time) { + imd_fill_header((IMDheader *)msgdata, IMD_TIME, 1); + double dt = update->dt; + + double currtime = update->atime + ((update->ntimestep - update->atimestep) * update->dt); + long long currstep = update->ntimestep; + char *time = (msgdata+IMDHEADERSIZE); + + memcpy(time, &dt, sizeof(dt)); + memcpy(time+sizeof(dt), &currtime, sizeof(currtime)); + memcpy(time+sizeof(dt)+sizeof(currtime), &currstep, sizeof(currstep)); + offset += IMDHEADERSIZE+sizeof(dt)+sizeof(currtime)+sizeof(currstep); + } + if (imdsinfo->box) { + imd_fill_header((IMDheader *)(msgdata + offset), IMD_BOX, 1); + // Get triclinic box vectors + float *box = (float *)(msgdata+offset+IMDHEADERSIZE); + box[0] = domain->h[0]; + box[1] = 0.0; + box[2] = 0.0; + box[3] = domain->h[5]; + box[4] = domain->h[1]; + box[5] = 0.0; + box[6] = domain->h[4]; + box[7] = domain->h[3]; + box[8] = domain->h[2]; + + offset += (9*4)+IMDHEADERSIZE; + + } + if (imdsinfo->coords) { + imd_fill_header((IMDheader *)(msgdata + offset), IMD_FCOORDS, num_coords); + global_coords = (float *) (msgdata + offset + IMDHEADERSIZE); + offset += 3*4*num_coords+IMDHEADERSIZE; + } + if (imdsinfo->velocities) { + imd_fill_header((IMDheader *)(msgdata + offset), IMD_VELOCITIES, num_coords); + global_vel = (float *) (msgdata + offset + IMDHEADERSIZE); + offset += 3*4*num_coords+IMDHEADERSIZE; + } + if (imdsinfo->forces) { + imd_fill_header((IMDheader *)(msgdata + offset), IMD_FORCES, num_coords); + global_force = (float *) (msgdata + offset + IMDHEADERSIZE); + offset += 3*4*num_coords+IMDHEADERSIZE; + } + } + + int ntotal, nmax, nme=0; + for (int i=0; i < nlocal; ++i) + if (mask[i] & groupbit) ++nme; + + // Atoms per proc + int *recvcounts = nullptr; + // Displacements in recv for each proc + int *displs = nullptr; + + + if (me == 0) { + recvcounts = new int[comm->nprocs]; + displs = new int[comm->nprocs]; + } + + MPI_Gather(&nme, 1, MPI_INT, recvcounts, 1, MPI_INT, 0, world); + + if (me == 0) { + displs[0] = 0; + ntotal = recvcounts[0]; + for (int i = 1; i < comm->nprocs; ++i) { + displs[i] = displs[i - 1] + recvcounts[i - 1]; + ntotal += recvcounts[i]; + } + } + + if (imdsinfo->coords) { + commdata *recvcoord = nullptr; + memory->destroy(coord_data); + coord_data = memory->smalloc(nme*size_one, "imd:coord_data"); + buf = static_cast(coord_data); + int idx = 0; + if (imdsinfo->unwrap) { + double xprd = domain->xprd; + double yprd = domain->yprd; + double zprd = domain->zprd; + double xy = domain->xy; + double xz = domain->xz; + double yz = domain->yz; + for (int i = 0; i < nlocal; ++i) { + if (mask[i] & groupbit) { + int ix = (image[i] & IMGMASK) - IMGMAX; + int iy = (image[i] >> IMGBITS & IMGMASK) - IMGMAX; + int iz = (image[i] >> IMG2BITS) - IMGMAX; + + if (domain->triclinic) { + buf[idx].tag = tag[i]; + buf[idx].x = x[i][0]; + ix * xprd + iy * xy + iz * xz; + buf[idx].y = x[i][1]; + iy * yprd + iz * yz; + buf[idx].z = x[i][2]; + iz * zprd; + } + else { + buf[idx].tag = tag[i]; + buf[idx].x = x[i][0]; + ix * xprd; + buf[idx].y = x[i][1]; + iy * yprd; + buf[idx].z = x[i][2]; + iz * zprd; + } + ++idx; + } + } + } + else { + for (int i = 0; i < nlocal; ++i) { + if (mask[i] & groupbit) { + buf[idx].tag = tag[i]; + buf[idx].x = x[i][0]; + buf[idx].y = x[i][1]; + buf[idx].z = x[i][2]; + ++idx; + } + } + } + if (me == 0) { + recvcoord = new commdata[ntotal]; + } + MPI_Gatherv(buf, nme, MPI_CommData, + recvcoord, recvcounts, displs, MPI_CommData, 0, world); + if (me == 0) { + // Sort the coordinates by tag, place in global_coords + for (int i = 0; i < comm->nprocs; ++i) { + for (int j = 0; j < recvcounts[i]; ++j) { + int idx = displs[i]+j; + const tagint t = 3*taginthash_lookup((taginthash_t *)idmap, recvcoord[idx].tag); + if (t != 3*HASH_FAIL) { + global_coords[t] = recvcoord[idx].x; + global_coords[t + 1] = recvcoord[idx].y; + global_coords[t + 2] = recvcoord[idx].z; + } + } + } + } + } + if (imdsinfo->velocities) { + commdata *recvvel = nullptr; + memory->destroy(vel_data); + vel_data = memory->smalloc(nme*size_one, "imd:vel_data"); + buf = static_cast(vel_data); + int idx = 0; + for (int i = 0; i < nlocal; ++i) { + if (mask[i] & groupbit) { + buf[idx].tag = tag[i]; + buf[idx].x = v[i][0]; + buf[idx].y = v[i][1]; + buf[idx].z = v[i][2]; + ++idx; + } + } + if (me == 0) { + recvvel = new commdata[ntotal]; + } + MPI_Gatherv(buf, nme, MPI_CommData, + recvvel, recvcounts, displs, MPI_CommData, 0, world); + if (me == 0) { + // Sort the coordinates by tag, place in global_vels + for (int i = 0; i < comm->nprocs; ++i) { + for (int j = 0; j < recvcounts[i]; ++j) { + int idx = displs[i]+j; + const tagint t = 3*taginthash_lookup((taginthash_t *)idmap, recvvel[idx].tag); + if (t != 3*HASH_FAIL) { + global_vel[t] = recvvel[idx].x; + global_vel[t + 1] = recvvel[idx].y; + global_vel[t + 2] = recvvel[idx].z; + } + } + } + } + } + if (imdsinfo->forces) { + commdata *recvforce = nullptr; + memory->destroy(force_data); + force_data = memory->smalloc(nme*size_one, "imd:force_data"); + buf = static_cast(force_data); + int idx = 0; + for (int i = 0; i < nlocal; ++i) { + if (mask[i] & groupbit) { + buf[idx].tag = tag[i]; + buf[idx].x = f[i][0]; + buf[idx].y = f[i][1]; + buf[idx].z = f[i][2]; + ++idx; + } + } + if (me == 0) { + recvforce = new commdata[ntotal]; + } + MPI_Gatherv(buf, nme, MPI_CommData, + recvforce, recvcounts, displs, MPI_CommData, 0, world); + if (me == 0) { + // Sort the coordinates by tag, place in global_coords + for (int i = 0; i < comm->nprocs; ++i) { + for (int j = 0; j < recvcounts[i]; ++j) { + int idx = displs[i]+j; + const tagint t = 3*taginthash_lookup((taginthash_t *)idmap, recvforce[idx].tag); + if (t != 3*HASH_FAIL) { + global_force[t] = recvforce[idx].x; + global_force[t + 1] = recvforce[idx].y; + global_force[t + 2] = recvforce[idx].z; + } + } + } + } + } + +/* done collecting frame data now communicate with IMD client. */ + +#if defined(LAMMPS_ASYNC_IMD) + /* wake up i/o worker thread and release lock on i/o buffer + * we can go back to our MD and let the i/o thread do the rest */ + buf_has_data=1; + pthread_cond_signal(&write_cond); + pthread_mutex_unlock(&write_mutex); +#else + /* send coordinate data, if client is able to accept */ + if (clientsock && imdsock_selwrite(clientsock,0)) { + imd_writen(clientsock, msgdata, msglen); + } +#endif } /* End of FixIMD class implementation. */ @@ -1401,13 +1964,35 @@ static int32 imd_writen(void *s, const char *ptr, int32 n) { return n; } -int imd_handshake(void *s) { +int imd_handshake_v2(void *s) { IMDheader header; imd_fill_header(&header, IMD_HANDSHAKE, 1); - header.length = IMDVERSION; /* Not byteswapped! */ + header.length = 2; /* Not byteswapped! */ return (imd_writen(s, (char *)&header, IMDHEADERSIZE) != IMDHEADERSIZE); } +int imd_handshake_v3(void *s, IMDSessionInfo *imdsinfo) { + IMDheader header; + imd_fill_header(&header, IMD_HANDSHAKE, 1); + header.length = 3; /* Not byteswapped so client can determine native endinaness */ + + if (imd_writen(s, (char *)&header, IMDHEADERSIZE) != IMDHEADERSIZE) return -1; + + imd_fill_header(&header, IMD_SESSIONINFO, 7); + unsigned char body[7] = {0}; + body[0] = imdsinfo->time; + body[1] = imdsinfo->energies; + body[2] = imdsinfo->box; + body[3] = imdsinfo->coords; + body[4] = !imdsinfo->unwrap; + body[5] = imdsinfo->velocities; + body[6] = imdsinfo->forces; + + if (imd_writen(s, (char *)&header, IMDHEADERSIZE) != IMDHEADERSIZE || + imd_writen(s, (char *)&body, 7) != 7) return -1; + return 0; +} + /* The IMD receive functions */ IMDType imd_recv_header(void *s, int32 *length) { diff --git a/src/MISC/fix_imd.h b/src/MISC/fix_imd.h index 2312540e66..6b8778dbf0 100644 --- a/src/MISC/fix_imd.h +++ b/src/MISC/fix_imd.h @@ -56,6 +56,9 @@ FixStyle(imd,FixIMD); #include #endif +/* IMDv3 session information */ +struct IMDSessionInfo; + /* prototype for c wrapper that calls the real worker */ extern "C" void *fix_imd_ioworker(void *); @@ -69,8 +72,11 @@ class FixIMD : public Fix { void init() override; void setup(int) override; void post_force(int) override; + void end_of_step() override; void post_force_respa(int, int, int) override; double memory_usage() override; + // Fix nevery at 1, use trate to skip in 'end_of_step` + int nevery = 1; protected: int imd_port; @@ -80,13 +86,17 @@ class FixIMD : public Fix { int num_coords; // total number of atoms controlled by this fix int size_one; // bytes per atom in communication buffer. int maxbuf; // size of atom communication buffer. - void *comm_buf; // communication buffer + void *coord_data; // communication buffer for coordinates + void *vel_data; // communication buffer for velocities + void *force_data; // communication buffer for forces void *idmap; // hash for mapping atom indices to consistent order. tagint *rev_idmap; // list of the hash keys for reverse mapping. - int imd_forces; // number of forces communicated via IMD. - void *force_buf; // force data buffer - double imd_fscale; // scale factor for forces. in case VMD's units are off. + int imd_version; // version of the IMD protocol to be used. + + int imd_forces; // number of forces communicated via IMD. + void *recv_force_buf; // force data buffer + double imd_fscale; // scale factor for forces. in case VMD's units are off. int imd_inactive; // true if IMD connection stopped. int imd_terminate; // true if IMD requests termination of run. @@ -96,12 +106,20 @@ class FixIMD : public Fix { int nowait_flag; // true if LAMMPS should not wait with the execution for VMD. int connect_msg; // flag to indicate whether a "listen for connection message" is needed. + /* IMDv3-only */ + IMDSessionInfo *imdsinfo; // session information for IMDv3 + int me; // my MPI rank in this "world". int nlevels_respa; // flag to determine respa levels. int msglen; char *msgdata; + private: + void handle_step_v2(); + void handle_client_input_v3(); + void handle_output_v3(); + #if defined(LAMMPS_ASYNC_IMD) int buf_has_data; // flag to indicate to the i/o thread what to do. pthread_mutex_t write_mutex; // mutex for sending coordinates to i/o thread