From 16810b84eba91c85e23dd9d0f62392d0d7f4ec87 Mon Sep 17 00:00:00 2001 From: Stan Gerald Moore Date: Mon, 24 Jan 2022 14:35:27 -0700 Subject: [PATCH] Add option to balance dump file output --- src/dump.cpp | 165 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/dump.h | 2 + 2 files changed, 167 insertions(+) diff --git a/src/dump.cpp b/src/dump.cpp index e5652c210c..cdf2f57d9e 100644 --- a/src/dump.cpp +++ b/src/dump.cpp @@ -75,6 +75,7 @@ Dump::Dump(LAMMPS *lmp, int /*narg*/, char **arg) : Pointers(lmp) clearstep = 0; sort_flag = 0; + balance_flag = 0; append_flag = 0; buffer_allow = 0; buffer_flag = 0; @@ -222,6 +223,9 @@ void Dump::init() ids = idsort = nullptr; index = proclist = nullptr; irregular = nullptr; + + if (balance_flag) + error->all(FLERR,"Cannot balance dump output without sorting enabled"); } if (sort_flag) { @@ -417,6 +421,7 @@ void Dump::write() if (sort_flag && sortcol == 0) pack(ids); else pack(nullptr); if (sort_flag) sort(); + if (balance_flag) balance(); // write timestep header // for multiproc, @@ -888,6 +893,161 @@ int Dump::bufcompare_reverse(const int i, const int j, void *ptr) #endif +/* ---------------------------------------------------------------------- + parallel load balance of buf across all procs + must come after sort +------------------------------------------------------------------------- */ + +void Dump::balance() +{ + bigint *proc_offsets,*proc_new_offsets; + memory->create(proc_offsets,nprocs+1,"dump:proc_offsets"); + memory->create(proc_new_offsets,nprocs+1,"dump:proc_new_offsets"); + + // compute atom offset for this proc + + bigint offset; + bigint bnme = nme; + MPI_Scan(&bnme,&offset,1,MPI_LMP_BIGINT,MPI_SUM,world); + + // gather atom offsets for all procs + + MPI_Allgather(&offset,1,MPI_LMP_BIGINT,&proc_offsets[1],1,MPI_LMP_BIGINT,world); + + proc_offsets[0] = 0; + + // how many atoms should I own after balance + + int nme_balance = static_cast(ntotal/nprocs); + + // include remainder atoms on first procs + + int remainder = ntotal % nprocs; + if (me < remainder) nme_balance += 1; + + // compute new atom offset for this proc + + bigint offset_balance; + bigint bnme_balance = nme_balance; + MPI_Scan(&bnme_balance,&offset_balance,1,MPI_LMP_BIGINT,MPI_SUM,world); + + // gather new atom offsets for all procs + + MPI_Allgather(&offset_balance,1,MPI_LMP_BIGINT,&proc_new_offsets[1],1,MPI_LMP_BIGINT,world); + + proc_new_offsets[0] = 0; + + bigint start = proc_new_offsets[me]; + bigint end = proc_new_offsets[me+1]; + + // reset buf size to largest of any post-balance nme values + // this insures proc 0 can receive everyone's info + + int nmax; + MPI_Allreduce(&nme_balance,&nmax,1,MPI_INT,MPI_MAX,world); + + // allocate a second buffer for balanced data + + double* buf_balance; + memory->create(buf_balance,nmax*size_one,"dump:buf_balance"); + + // compute from which procs I am receiving atoms + // post recvs first + + int nswap = 0; + MPI_Request *request = new MPI_Request[nprocs]; + int procstart = start; + int iproc = me; + int iproc_prev; + + for (bigint i = start; i < end; i++) { + + // find which proc this atom belongs to + + while (i < proc_offsets[iproc]) iproc--; + while (i > proc_offsets[iproc+1]-1) iproc++; + + if (i != start && (iproc != iproc_prev || i == end-1)) { + + // finished with proc + + int procrecv = iproc; + if (iproc != iproc_prev) procrecv = iproc_prev; + + int procnrecv = i - procstart + 1; + if (iproc != iproc_prev) procnrecv--; + + // post receive for this proc + + if (iproc_prev != me) + MPI_Irecv(&buf_balance[(procstart - start)*size_one],procnrecv*size_one,MPI_DOUBLE, + procrecv,0,world,&request[nswap++]); + + procstart = i; + } + + iproc_prev = iproc; + } + + // compute which atoms I am sending and to which procs + + procstart = 0; + iproc = me; + for (int i = 0; i < nme; i++) { + + // find which proc this atom should belong to + + while (proc_offsets[me] + i < proc_new_offsets[iproc]) iproc--; + while (proc_offsets[me] + i > proc_new_offsets[iproc+1] - 1) iproc++; + + if (i != 0 && (iproc != iproc_prev || i == nme - 1)) { + + // finished with proc + + int procsend = iproc; + if (iproc != iproc_prev) procsend = iproc_prev; + + int procnsend = i - procstart + 1; + if (iproc != iproc_prev) procnsend--; + + // send for this proc + + if (iproc_prev != me) { + MPI_Send(&buf[procstart*size_one],procnsend*size_one,MPI_DOUBLE,procsend,0,world); + } else { + + // sending to self, copy buffers + + int offset_me = proc_offsets[me] - start; + memcpy(&buf_balance[(offset_me + procstart)*size_one],&buf[procstart*size_one],procnsend*size_one*sizeof(double)); + } + + procstart = i; + } + + iproc_prev = iproc; + } + + // wait for all recvs + + for (int n = 0; n < nswap; n++) + MPI_Wait(&request[n],MPI_STATUS_IGNORE); + + nme = nme_balance; + + // swap buffers + + double *tmp = buf; + buf = buf_balance; + + // cleanup + + memory->destroy(tmp); + memory->destroy(proc_offsets); + memory->destroy(proc_new_offsets); + delete [] request; +} + /* ---------------------------------------------------------------------- process params common to all dumps here if unknown param, call modify_param specific to the dump @@ -1108,6 +1268,11 @@ void Dump::modify_params(int narg, char **arg) } iarg += 2; + } else if (strcmp(arg[iarg],"balance") == 0) { + if (iarg+2 > narg) error->all(FLERR,"Illegal dump_modify command"); + balance_flag = utils::logical(FLERR,arg[iarg+1],false,lmp); + iarg += 2; + } else if (strcmp(arg[iarg],"time") == 0) { if (iarg+2 > narg) error->all(FLERR,"Illegal dump_modify command"); time_flag = utils::logical(FLERR,arg[iarg+1],false,lmp); diff --git a/src/dump.h b/src/dump.h index 482e87a207..83edda3810 100644 --- a/src/dump.h +++ b/src/dump.h @@ -67,6 +67,7 @@ class Dump : protected Pointers { int header_flag; // 0 = item, 2 = xyz int flush_flag; // 0 if no flush, 1 if flush every dump int sort_flag; // 1 if sorted output + int balance_flag; // 1 if load balanced output int append_flag; // 1 if open file in append mode, 0 if not int buffer_allow; // 1 if style allows for buffer_flag, 0 if not int buffer_flag; // 1 if buffer output as one big string, 0 if not @@ -161,6 +162,7 @@ class Dump : protected Pointers { static int bufcompare(const int, const int, void *); static int bufcompare_reverse(const int, const int, void *); #endif + void balance(); }; } // namespace LAMMPS_NS